40 lines
1.2 KiB
Python
40 lines
1.2 KiB
Python
import lzma
|
|
|
|
from sklearn.naive_bayes import MultinomialNB
|
|
from sklearn.feature_extraction.text import TfidfVectorizer
|
|
|
|
|
|
def get_data(file_name, data_type):
|
|
lines = []
|
|
if data_type == "tsv":
|
|
with open(file_name, encoding="utf-8") as file:
|
|
for line in file.readlines():
|
|
lines.append(int(line.replace("\n", "")))
|
|
else:
|
|
with lzma.open(f"{file_name}.{data_type}") as file:
|
|
for line in file.readlines():
|
|
lines.append(line.rstrip().decode("utf-8"))
|
|
return lines
|
|
|
|
|
|
def bayes(train):
|
|
x_data = get_data(f"{train}/in.tsv", "xz")
|
|
Y_data = get_data(f"{train}/expected.tsv", "tsv")
|
|
|
|
vectorizer = TfidfVectorizer(stop_words="english")
|
|
X_data = vectorizer.fit_transform(x_data)
|
|
|
|
clf = MultinomialNB()
|
|
y_pred = clf.fit(X_data, Y_data)
|
|
|
|
for predct in ["test-A", "dev-0"]:
|
|
Y_test = get_data(f"{predct}/in.tsv", "xz")
|
|
y_prediction = y_pred.predict(vectorizer.transform(Y_test))
|
|
|
|
with open(f"{predct}\out.tsv", "w", encoding="UTF-8") as file_out:
|
|
for single_pred in y_prediction:
|
|
file_out.writelines(f"{str(single_pred)}\n")
|
|
|
|
|
|
bayes("train")
|