linear regr

script.py

@@ -0,0 +1,34 @@
+from sklearn.linear_model import LinearRegression
+import pandas as pd
+import numpy as np
+
+
+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 = X.drop(["brand"], axis=1)
+    X = pd.get_dummies(X,  columns= ["engine_type"], drop_first=True)
+    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 = df_dev.drop(["brand"], axis=1)
+        df_dev = pd.get_dummies(df_dev, columns= ["engine_type"], drop_first=True)
+        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()