test_team / test_team.py
Fizza03's picture
Add test_team.py
b764b4e 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)