forked from kubapok/auta-public
59 lines
1.8 KiB
Python
59 lines
1.8 KiB
Python
import pandas as pd
|
|
from sklearn.linear_model import LinearRegression
|
|
from sklearn.metrics import mean_squared_error
|
|
from sklearn.preprocessing import PolynomialFeatures
|
|
|
|
col_names = ["Price","Mileage","Year","Brand","EngineType","EngineCapacity"]
|
|
|
|
|
|
|
|
def prepareData(df):
|
|
df["Age"] = 2018 - df["Year"]
|
|
df["SqrtAge"] = df.Age**0.5
|
|
df = pd.concat([df, df['EngineType'].str.get_dummies()], axis = 1 )
|
|
df = df.drop(['EngineType','Brand'], axis = 1)
|
|
df["SqrtMileage"] = df.Mileage ** 0.5
|
|
df["SqrtEngineCapacity"] = df.EngineCapacity ** 0.5
|
|
poly = PolynomialFeatures(2, interaction_only=True)
|
|
df = poly.fit_transform(df)
|
|
return df
|
|
|
|
def main():
|
|
|
|
df =pd.read_csv('train/train.tsv', sep="\t", names=col_names)
|
|
y_dev =pd.read_csv('dev-0/expected.tsv', sep="\t", names=["Price"])
|
|
x_dev =pd.read_csv('dev-0/in.tsv', sep="\t", names=["Mileage","Year","Brand","EngineType","EngineCapacity"])
|
|
x_test =pd.read_csv('test-A/in.tsv', sep="\t", names=["Mileage","Year","Brand","EngineType","EngineCapacity"])
|
|
|
|
y_train = df.Price
|
|
x_train = df.drop('Price', axis=1)
|
|
|
|
x_train = prepareData(x_train)
|
|
|
|
linReg = LinearRegression()
|
|
linReg.fit(x_train, y_train)
|
|
|
|
x_dev = prepareData(x_dev)
|
|
x_test = prepareData(x_test)
|
|
|
|
#Score modelu dla zbioru dev
|
|
score = linReg.score(x_dev, y_dev)
|
|
print(score)
|
|
|
|
#Wartość RMSE dla zbioru dev
|
|
y_pred = linReg.predict(x_dev)
|
|
data = {'Price':y_pred}
|
|
y_pred = pd.DataFrame(data)
|
|
rmse = mean_squared_error(y_dev, y_pred, squared=False)
|
|
print(rmse)
|
|
|
|
#predict dla test-A
|
|
y_pred_test = linReg.predict(x_test)
|
|
data = {'Price':y_pred_test}
|
|
y_pred_test = pd.DataFrame(data)
|
|
|
|
y_pred_test.to_csv(r'test-A/out.tsv', sep='\t')
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |