sport-text-classification-b.../main.py

24 lines
795 B
Python

import pandas as pd
from sklearn.naive_bayes import MultinomialNB
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.pipeline import make_pipeline
train_set = pd.read_table('train/train.tsv.gz', error_bad_lines=False, header=None)
dev_set = pd.read_table('dev-0/in.tsv', error_bad_lines=False, header=None)
test_set = pd.read_table('test-A/in.tsv', error_bad_lines=False, header=None)
X_train = train_set[1].values
y_train = train_set[0].values
X_dev = dev_set[0].values
X_test = test_set[0].values
model = make_pipeline(TfidfVectorizer(), MultinomialNB())
model.fit(X_train, y_train)
dev_prediction = model.predict(X_dev)
test_prediction = model.predict(X_test)
dev_prediction.tofile('./dev-0/out.tsv', sep='\n')
test_prediction.tofile('./test-A/out.tsv', sep='\n')