57 lines
1.3 KiB
Python
57 lines
1.3 KiB
Python
|
from naivebayes import NaiveBayesTextClassifier
|
||
|
import lzma
|
||
|
from spacy.lang.en.stop_words import STOP_WORDS as en_stop
|
||
|
|
||
|
categories_list = [0, 1]
|
||
|
|
||
|
classifier = NaiveBayesTextClassifier(
|
||
|
categories=categories_list,
|
||
|
stop_words=en_stop
|
||
|
)
|
||
|
|
||
|
X = []
|
||
|
Y = []
|
||
|
with lzma.open('train/in.tsv.xz', 'r') as file:
|
||
|
for line in file:
|
||
|
line = line.strip()
|
||
|
X.append(line.decode("utf-8"))
|
||
|
|
||
|
|
||
|
with open('train/expected.tsv', 'r') as file:
|
||
|
for line in file:
|
||
|
line = line.strip()
|
||
|
Y.append(int(line))
|
||
|
print(len(X), len(Y))
|
||
|
classifier.train(X[:15000], Y[:15000])
|
||
|
classifier.train(X[15000:30000], Y[15000:30000])
|
||
|
# classifier.train(X[30000:60000], Y[30000:60000])l
|
||
|
# classifier.train(X[60000:90000], Y[60000:90000])
|
||
|
|
||
|
|
||
|
test_x = []
|
||
|
with lzma.open('dev-0/in.tsv.xz', 'r') as file:
|
||
|
for line in file:
|
||
|
line = line.strip()
|
||
|
test_x.append(line.decode("utf-8"))
|
||
|
|
||
|
predicted_classes = classifier.classify(test_x)
|
||
|
|
||
|
|
||
|
f = open("dev-0/out.tsv", "a")
|
||
|
for p in predicted_classes:
|
||
|
f.write(str(p) + '\n')
|
||
|
f.close()
|
||
|
|
||
|
test_x = []
|
||
|
with lzma.open('test-A/in.tsv.xz', 'r') as file:
|
||
|
for line in file:
|
||
|
line = line.strip()
|
||
|
test_x.append(line.decode("utf-8"))
|
||
|
|
||
|
predicted_classes = classifier.classify(test_x)
|
||
|
|
||
|
|
||
|
f = open("test-A/out.tsv", "a")
|
||
|
for p in predicted_classes:
|
||
|
f.write(str(p) + '\n')
|
||
|
f.close()
|