team1 / model_wrapper.py
k224742's picture
Upload 2 files
132ff2f verified
raw
history blame contribute delete
445 Bytes
# model_wrapper.py
import joblib
import numpy as np
LABELS = ['admiration','anger','disgust','fear','hope',
'joy','love','pride','sadness']
class MyModel:
def __init__(self):
bundle = joblib.load("model.pkl")
self.vectorizer = bundle["vectorizer"]
self.classifier = bundle["classifier"]
def predict(self, texts):
X = self.vectorizer.transform(texts)
return self.classifier.predict(X)