34 lines
941 B
Python
34 lines
941 B
Python
#!/usr/bin/env python3
|
|
|
|
"""
|
|
Predict using argument types model.
|
|
"""
|
|
|
|
import pandas as pd
|
|
import pickle
|
|
import string
|
|
|
|
|
|
class ArgPredict():
|
|
"""
|
|
Class for argument types model.
|
|
"""
|
|
unlabel_dict = {
|
|
0:'hipoteza',1:'rzeczowe', 2:'logiczne', 3:'emocjonalne', 4:'inne' }
|
|
|
|
|
|
def __init__(self, modelpath="arg_model.pkl", vectorpath="arg_vector.pkl"):
|
|
self.model = pickle.load(open(modelpath, 'rb'))
|
|
self.cv = pickle.load(open(vectorpath, 'rb'))
|
|
|
|
def __normalize(self, text):
|
|
punctuation = string.punctuation + "„“”«»‚’-–…"
|
|
return ''.join([c for c in text.lower() if c not in punctuation])
|
|
|
|
def predict(self, messages):
|
|
df = pd.DataFrame(messages, columns=['body_text'])
|
|
features = self.cv.transform(df['body_text'].apply(self.__normalize)).toarray()
|
|
df['predictions'] = self.model.predict(features)
|
|
df['predictions'] = df['predictions'].apply(lambda t: self.unlabel_dict[t])
|
|
return list(df['predictions'])
|