sport-text-classification-b.../main.py
2021-05-12 13:47:40 +02:00

52 lines
1.5 KiB
Python

import pandas as pd
from sklearn.naive_bayes import MultinomialNB
from sklearn.feature_extraction.text import TfidfVectorizer
def main():
clf = MultinomialNB()
vectorizer = TfidfVectorizer()
x_train, y_train = get_training_data(vectorizer)
clf.fit(x_train, y_train)
x_dev = get_dev_data(vectorizer)
Y_dev_predicted = clf.predict(x_dev)
save_to_tsv(Y_dev_predicted, 'dev-0/out.tsv')
x_test = get_test_data(vectorizer)
Y_test_predicted = clf.predict(x_test)
save_to_tsv(Y_test_predicted, 'test-A/out.tsv')
def save_to_tsv(data, path):
pd.DataFrame(data).to_csv(path, sep='\t', index=False, header=False)
def get_training_data(vectorizer):
train_dataset = pd.read_csv('train/train.tsv/train.tsv', sep='\t', header=None, error_bad_lines=False)
y_train = train_dataset[0]
X_train = train_dataset[1]
x_train = [str(item) for item in X_train.to_numpy()]
x_train = vectorizer.fit_transform(x_train)
return x_train, y_train
def get_dev_data(vectorizer):
dev_dataset = pd.read_csv('dev-0/in.tsv', sep='\t', header=None, error_bad_lines=False)
X_dev = [str(item) for item in dev_dataset.to_numpy()]
return vectorizer.transform(X_dev)
def get_test_data(vectorizer):
test_dataset = pd.read_csv('test-A/in.tsv', sep='\t', header=None, error_bad_lines=False)
X_test = [str(item) for item in test_dataset.to_numpy()]
return vectorizer.transform(X_test)
if __name__ == '__main__':
main()