competition-test-model / model_wrapper.py
malimustafaa's picture
Upload model_wrapper.py with huggingface_hub
5cf607d
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)