35 lines
1.0 KiB
Python
35 lines
1.0 KiB
Python
import numpy
|
|
from sklearn.preprocessing import LabelEncoder
|
|
from sklearn.naive_bayes import GaussianNB, MultinomialNB
|
|
from sklearn.pipeline import Pipeline
|
|
from sklearn.feature_extraction.text import TfidfVectorizer
|
|
|
|
|
|
def trainModel(trainFileIn, trainFileExpected):
|
|
with open(trainFileExpected, 'r') as f:
|
|
expectedData = f.readlines()
|
|
|
|
with open(trainFileIn, 'r') as f:
|
|
inData = f.readlines()
|
|
|
|
expectedDataEncoded = LabelEncoder().fit_transform(expectedData)
|
|
|
|
pipeline = Pipeline(steps=[
|
|
('tfidf', TfidfVectorizer()),
|
|
('naive-bayes', MultinomialNB())
|
|
])
|
|
|
|
return pipeline.fit(inData, expectedDataEncoded)
|
|
|
|
def evaluateModel(model, inFile, outFile):
|
|
with open(inFile, 'r') as f:
|
|
inData = f.readlines()
|
|
prediction = model.predict(inData)
|
|
numpy.savetxt(outFile, prediction, fmt='%d', delimiter='\n')
|
|
|
|
model = trainModel("train/in.tsv", "train/expected.tsv")
|
|
evaluateModel(model, "dev-0/in.tsv", "dev-0/out.tsv")
|
|
evaluateModel(model, "test-A/in.tsv", "test-A/out.tsv")
|
|
|
|
|