444452

run.py

@@ -0,0 +1,66 @@
+#!/usr/bin/env python
+
+import os
+import pandas as pd
+from sklearn.metrics import accuracy_score
+from gzip import open as open_gz
+from sklearn.feature_extraction.text import TfidfVectorizer
+from sklearn.naive_bayes import MultinomialNB
+from sklearn.pipeline import make_pipeline
+
+
+def read_gz_file(path, filename, cutoff=None):
+    i = 0
+    data = []
+    with open_gz(os.path.join(path, filename), 'rt', 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=['y', 'x'])
+    return df
+
+
+def read_tsv(path, filename, 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())
+            i += 1
+    df = pd.DataFrame(data, columns=['col_name'])
+    return df
+
+
+def evaluate_and_save(path, file, model):
+    # print(path, file)
+    df = read_tsv(path, file)
+    # print(df)
+    predicted = model.predict(df['col_name'])
+    # expected = read_tsv(path, 'expected.tsv')['col_name']
+    # print('score: ', accuracy_score(expected, predicted))
+    # print(type(predicted))
+    with open(os.path.join(path, 'out.tsv'), 'w') as f:
+        for value in predicted:
+            f.write(f'{value}\n')
+
+
+def main():
+    train_data = read_gz_file('train', 'train.tsv.gz', 20000)
+
+    tfidf_vectorizer = TfidfVectorizer()
+    model = MultinomialNB()
+    pipe = make_pipeline(tfidf_vectorizer, model)
+    x_train = train_data['x']
+    y_train = train_data['y']
+
+    pipe.fit(x_train.values, y_train.values)
+    for path, file in [('dev-0', 'in.tsv'), ('test-A', 'in.tsv')]:
+        evaluate_and_save(path, file, pipe)
+
+
+if __name__ == '__main__':
+    main()