nei-demo-backup / svm_predict.py
SmitaGautam's picture
Update svm_predict.py
587b779 verified
raw
history blame
974 Bytes
import nltk
from nltk import word_tokenize
from nltk import pos_tag
import joblib
import numpy as np
from train import feature_vector, pos_tags
model = joblib.load('model.pkl')
scaler = joblib.load('scaler.pkl')
nltk.download('averaged_perceptron_tagger_eng')
nltk.download('punkt_tab')
def predict(sentence):
tokens = word_tokenize(sentence)
sent_pos_tags = pos_tag(tokens)
sent_features = []
l = len(tokens)
for idx, word in enumerate(tokens):
current_tag = sent_pos_tags[idx][1]
current_idx = pos_tags.index(current_tag) if current_tag in pos_tags else -1
prev_word = tokens[idx-1] if idx!=0 else ""
next_word = tokens[idx+1] if idx<l-1 else ""
word_features = feature_vector(word, (idx+1)/l, current_idx, prev_word, next_word)
sent_features.append(word_features)
scaled_features = scaler.transform(sent_features)
predictions = model.predict(scaled_features)
return tokens, predictions