forked from kubapok/auta-public
46 lines
1.5 KiB
Python
46 lines
1.5 KiB
Python
from sklearn.linear_model import LinearRegression
|
|
from sklearn.preprocessing import PolynomialFeatures
|
|
import pandas as pd
|
|
import numpy as np
|
|
|
|
brands = None
|
|
|
|
def process_data(df):
|
|
df["age"] = df.year.apply(lambda x: np.sqrt(2017-x))
|
|
df["sqrt_mileage"] = df.mileage.apply(lambda x: np.sqrt(x))
|
|
df["sqrt_engine_capacity"] = df.engine_capacity.apply(lambda x: np.sqrt(x))
|
|
global brands
|
|
if not brands:
|
|
brands = df.brand.value_counts()[:35].index.tolist()
|
|
df.brand = df.brand.apply(lambda x: x if x in brands else "0")
|
|
df = pd.get_dummies(df)
|
|
poly = PolynomialFeatures(2, interaction_only=True)
|
|
df = poly.fit_transform(df)
|
|
return df
|
|
|
|
def get_model():
|
|
df = pd.read_csv('./train/train.tsv', sep='\t',
|
|
names=["price", "mileage", "year", "brand", "engine_type", "engine_capacity"])
|
|
X = df.loc[:, df.columns != 'price']
|
|
y = df['price']
|
|
X = process_data(X)
|
|
regr = LinearRegression()
|
|
return regr.fit(X, y)
|
|
|
|
|
|
def predict_and_write(path, model):
|
|
with open(f'{path}out.tsv', 'w') as out:
|
|
df_dev = pd.read_csv(f'{path}in.tsv', sep='\t',
|
|
names=["mileage", "year", "brand", "engine_type", "engine_capacity"])
|
|
df_dev = process_data(df_dev)
|
|
predictions = model.predict(df_dev).astype(int)
|
|
for prediction in predictions:
|
|
out.write(f"{prediction}\n")
|
|
|
|
def main():
|
|
model = get_model()
|
|
predict_and_write('./dev-0/', model)
|
|
predict_and_write('./test-A/', model)
|
|
|
|
if __name__ == '__main__':
|
|
main() |