Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
import os
import pickle
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.naive_bayes import MultinomialNB
from sklearn.pipeline import Pipeline
def train(X, y,outpath=None, verbose=True):
def build(X, y=None):
"""
Inner build function that builds a single model.
"""
model = Pipeline([
('vectorizer', TfidfVectorizer(min_df=1, ngram_range=(1, 2))),
('clf', MultinomialNB(alpha=.01))
])
model.fit(X, y)
return model
model= build(X, y)
model.labels_ = y
if outpath:
with open(outpath, 'wb') as f:
pickle.dump(model, f)
if verbose:print("Model written out to {}".format(outpath))
return model
def predict(text,PATH):
try:
with open(PATH, 'rb') as f:
model = pickle.load(f)
except IOError:
return False
print(model.predict([
text
]))
return {
"class":model.predict([
text
]),
"accuracy":1
}
if __name__ == "__main__":
PATH = "model.pickle"
if not os.path.exists(PATH):
X = ["hello",
"fuck you",
"hey",
"hii",
"how are you ?"]
y = ["greeting",
"bad",
"greeting",
"greeting",
"cool"]
model = train(X, y, outpath=PATH)
else:
with open(PATH, 'rb') as f:
model = pickle.load(f)
with open(PATH, 'rb') as f:
model = pickle.load(f)
yhat = model.predict([
"sdlfkjdf lkjdlsj dl ksgjdlk"
])
if yhat.any():
all_labels = model.labels_.inverse_transform(yhat)
print all_labels[0][0]
os.remove(PATH)
else:
print False
os.remove(PATH)
# print(model.labels_.inverse_transform(yhat))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment