s444452

run.py

@@ -0,0 +1,56 @@
+#!/usr/bin/env python
+
+import os
+import math
+import pandas as pd
+from sklearn.metrics import mean_squared_error
+from sklearn.feature_extraction.text import TfidfVectorizer
+from sklearn.linear_model import LinearRegression
+from sklearn.pipeline import make_pipeline
+from stop_words import get_stop_words
+
+
+def rmse(y_expected, y_predicted):
+    return math.sqrt(mean_squared_error(y_expected, y_predicted))
+
+
+def read_tsv(path, filename, col_names, cutoff=None):
+    i = 0
+    data = []
+    with open(os.path.join(path, filename), 'r', encoding='utf-8') as file:
+        for line in file:
+            if cutoff and i >= cutoff:
+                break
+            data.append(line.strip().split('\t'))
+            i += 1
+
+    df = pd.DataFrame(data, columns=col_names)
+    return df
+
+
+def evaluate_and_save(path, file, model):
+    df = read_tsv(path, file, ['y'])['y'].values
+    y_predicted = model.predict(df)
+    # y_expected = read_tsv(path, 'expected.tsv', ['y'])['y'].values
+    # print('score: ', rmse(y_expected, y_predicted))
+    with open(os.path.join(path, 'out.tsv'), 'w') as f:
+        for value in y_predicted:
+            f.write(f'{value}\n')
+
+
+def main():
+    train_data = read_tsv('train', 'train.tsv', ['date_start', 'date_end', 'x1', 'x2', 'text'], 50000)
+    train_data['date'] = (train_data['date_start'].astype(float) + train_data['date_end'].astype(float)) / 2
+    train_x = train_data['text'].values
+    train_y = train_data['date'].values
+    vectorizer = TfidfVectorizer(stop_words=get_stop_words('polish'))
+    model = LinearRegression()
+    pipe = make_pipeline(vectorizer, model)
+    pipe.fit(train_x, train_y)
+    # print('training done')
+    for path, file in [('dev-0', 'in.tsv'), ('dev-1', 'in.tsv'), ('test-A', 'in.tsv')]:
+        evaluate_and_save(path, file, pipe)
+
+
+if __name__ == '__main__':
+    main()