32 lines
1.3 KiB
Python
32 lines
1.3 KiB
Python
|
import lzma
|
||
|
|
import nltk
|
||
|
|
from sklearn.naive_bayes import MultinomialNB
|
||
|
|
from sklearn.feature_extraction.text import TfidfVectorizer
|
||
|
|
|
||
|
|
with lzma.open("train/in.tsv.xz", "rt", encoding="utf-8") as train_file:
|
||
|
|
in_train = [x.strip().lower() for x in train_file.readlines()]
|
||
|
|
|
||
|
|
with open("train/expected.tsv", "r", encoding="utf-8") as train_file:
|
||
|
|
out_train = [int(x.strip()) for x in train_file.readlines()]
|
||
|
|
|
||
|
|
with lzma.open("dev-0/in.tsv.xz", "rt", encoding="utf-8") as dev_file:
|
||
|
|
in_dev = [x.strip().lower() for x in dev_file.readlines()]
|
||
|
|
|
||
|
|
with lzma.open("test-A/in.tsv.xz", "rt", encoding="utf-8") as test_file:
|
||
|
|
in_test = [x.strip().lower() for x in test_file.readlines()]
|
||
|
|
|
||
|
|
tfidf_vectorizer=TfidfVectorizer()
|
||
|
|
IN_train = tfidf_vectorizer.fit_transform(in_train)
|
||
|
|
classifier = MultinomialNB()
|
||
|
|
y_pred = classifier.fit(IN_train, out_train)
|
||
|
|
|
||
|
|
y_prediction = y_pred.predict(tfidf_vectorizer.transform(in_test))
|
||
|
|
with open("test-A/out.tsv", "w", encoding="utf-8") as test_out_file:
|
||
|
|
for single_pred in y_prediction:
|
||
|
|
test_out_file.writelines(f"{str(single_pred)}\n")
|
||
|
|
|
||
|
|
pred_dev = y_pred.predict(tfidf_vectorizer.transform(in_test))
|
||
|
|
with open("dev-0/out.tsv", "w", encoding="utf-8") as dev_out_file:
|
||
|
|
for single_pred in pred_dev:
|
||
|
|
dev_out_file.writelines(f"{str(single_pred)}\n")
|