61 lines
1.4 KiB
Python
61 lines
1.4 KiB
Python
from sklearn.naive_bayes import GaussianNB
|
|
import pandas as pd
|
|
from sklearn.naive_bayes import MultinomialNB
|
|
from sklearn.feature_extraction.text import TfidfVectorizer
|
|
|
|
PATHS = ['./train/train.tsv', './dev-0/in.tsv', './test-A/in.tsv']
|
|
PATHS_OUTPUT = ['./dev-0/out.tsv', './test-A/out.tsv']
|
|
|
|
def get_data(path):
|
|
return pd.read_table(path, error_bad_lines=False, sep='\t', header=None)
|
|
|
|
def get_X_y_train(data):
|
|
X_train = data[1].values
|
|
y_train = data[0].values
|
|
|
|
return X_train, y_train
|
|
|
|
def training(x, y):
|
|
vectorizer = TfidfVectorizer()
|
|
result = vectorizer.fit_transform(x)
|
|
classifier = MultinomialNB()
|
|
classifier.fit(result, y)
|
|
|
|
return classifier, vectorizer
|
|
|
|
|
|
def predict(vectorizer, classifier, x):
|
|
result = vectorizer.transform(x)
|
|
pred = classifier.predict(result)
|
|
|
|
return pred
|
|
|
|
def generate_output(pred, path):
|
|
pred.tofile(path, sep = '\n')
|
|
|
|
def main():
|
|
#prepare train
|
|
train = get_data(PATHS[0])
|
|
X_train, y_train = get_X_y_train(train)
|
|
|
|
#train
|
|
classifier, vectorizer = training(X_train, y_train)
|
|
|
|
#dev
|
|
X_dev = get_data(PATHS[1])
|
|
X_dev = X_dev[0].values
|
|
pred_dev = predict(vectorizer, classifier, X_dev)
|
|
|
|
#test
|
|
X_test = get_data(PATHS[2])
|
|
X_test = X_test[0].values
|
|
pred_test = predict(vectorizer, classifier, X_test)
|
|
|
|
#generate output
|
|
generate_output(pred_dev, PATHS_OUTPUT[0])
|
|
generate_output(pred_test, PATHS_OUTPUT[1])
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main() |