Modify script

main.py

@@ -0,0 +1,71 @@
+import pandas
+import os
+import sys
+from sklearn.linear_model import LinearRegression
+
+IN_FILE_NAME = "in.tsv"
+OUT_FILE_NAME = "out.tsv"
+MAIN_DIR = "."
+NAMES_FILE_NAME = "names"
+TRAIN_DIR = "train"
+TRAIN_FILE_NAME = "train.tsv"
+VALUE_SEP = "\t"
+LINE_SEP = "\n"
+CATEGORY_TYPE = "category"
+
+
+def main(dirname: str):
+    names = get_names()
+
+    X, Y = get_train_data(names)
+
+    clf = LinearRegression().fit(X, Y)
+
+    clf.predict(get_input_data(dirname, names)).tofile(os.path.join(
+        dirname, OUT_FILE_NAME), sep=LINE_SEP)
+
+
+def get_train_data(names: list):
+    train_path = os.path.join(MAIN_DIR, TRAIN_DIR, TRAIN_FILE_NAME)
+    check_file(train_path)
+    train_data = process_input(pandas.read_csv(
+        train_path, header=None, sep=VALUE_SEP, names=names), names)
+
+    train_data = train_data.loc[(train_data[names[0]] > 1000)]
+
+    X = train_data.loc[:, train_data.columns != names[0]]
+    Y = train_data[names[0]]
+
+    return X, Y
+
+
+def get_input_data(dirname, names):
+    in_path = os.path.join(dirname, IN_FILE_NAME)
+    check_file(in_path)
+    return process_input(pandas.read_csv(in_path, header=None, sep=VALUE_SEP, names=names[1:]), names)
+
+
+def process_input(df, names):
+    df = df.drop([names[4]], axis=1)
+    for c in df.select_dtypes(include=object).columns.values:
+        df[c] = df[c].astype(CATEGORY_TYPE).cat.codes
+    return df
+
+
+def get_names() -> list:
+    names_path = os.path.join(MAIN_DIR, NAMES_FILE_NAME)
+    check_file(names_path)
+
+    with open(names_path) as f_names:
+        return f_names.read().strip().split(VALUE_SEP)
+
+
+def check_file(filename: str):
+    if not os.path.exists(filename):
+        raise FileNotFoundError(filename)
+
+
+if __name__ == "__main__":
+    if len(sys.argv) < 2:
+        raise Exception("Name of working dir not specified!")
+    main(sys.argv[1])