naive bayes

run.py

@@ -0,0 +1,19 @@
+from sklearn.feature_extraction.text import TfidfVectorizer
+from sklearn.naive_bayes import MultinomialNB
+from sklearn.pipeline import make_pipeline
+import sys
+import pandas as pd
+
+train_file = sys.argv[1]
+pred_x_file = sys.argv[2]
+
+train = pd.read_csv(train_file, sep='\t', header=None)
+pred_x = pd.read_csv(pred_x_file, sep='\t', header=None)
+train_x, train_y = train[1], train[0]
+
+model = make_pipeline(TfidfVectorizer(), MultinomialNB())
+
+model.fit(train_x, train_y)
+pred_y = model.predict(pred_x.stack())
+
+pd.DataFrame(pred_y).to_csv('out.tsv', header=False, index=None)