43 lines
1.5 KiB
Python
43 lines
1.5 KiB
Python
import pandas as pd
|
|
from sklearn.naive_bayes import MultinomialNB
|
|
from sklearn.feature_extraction.text import TfidfVectorizer
|
|
|
|
|
|
# * Training data loading
|
|
with open('train/in.tsv', 'r', encoding='utf-8') as f:
|
|
x_train = pd.DataFrame(f.readlines(), columns=['text'])
|
|
y_train = pd.read_csv('train/expected.tsv', sep='\t',
|
|
names=['paranormal'], encoding='utf-8')
|
|
# *Validation data loading
|
|
with open('dev-0/in.tsv', 'r', encoding='utf-8') as f:
|
|
x_dev = pd.DataFrame(f.readlines(), columns=['text'])
|
|
# * Test data loading
|
|
with open('test-A/in.tsv', 'r', encoding='utf-8') as f:
|
|
x_test = pd.DataFrame(f.readlines(), columns=['text'])
|
|
|
|
# * Training data preparation
|
|
tfidf_vectorizer = TfidfVectorizer(max_df=0.95, max_features=500)
|
|
x_train_vectorized = tfidf_vectorizer.fit_transform(
|
|
x_train['text'].values)
|
|
|
|
# * Model training
|
|
mnb_model = MultinomialNB().fit(x_train_vectorized, y_train.values.ravel())
|
|
|
|
# * Validation data preparation
|
|
x_dev_prepared = tfidf_vectorizer.transform(x_dev['text'].values)
|
|
# * Validation data predictions
|
|
predictions = mnb_model.predict(x_dev_prepared)
|
|
# * Validation predicitons saving
|
|
with open('dev-0/out.tsv', 'w') as f:
|
|
for pred in predictions:
|
|
f.write(f'{pred}\n')
|
|
|
|
# * Test data preparation
|
|
x_test_vectorized = tfidf_vectorizer.transform(x_test['text'].values)
|
|
# * Test data predictions
|
|
predictions = mnb_model.predict(x_test_vectorized)
|
|
# * Test predictions saving
|
|
with open('test-A/out.tsv', 'w') as f:
|
|
for pred in predictions:
|
|
f.write(f'{pred}\n')
|