20 lines
595 B
Python
20 lines
595 B
Python
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)
|