intent-detector-v2 / inference.py
An-CNT's picture
Upload Vietnamese Intent Detector v2
c1be4e0 verified
import joblib
import json
from sentence_transformers import SentenceTransformer
import pandas as pd
def load_model():
clf = joblib.load("intent_classifier_lgb.pkl")
embedder = SentenceTransformer("sbert_embed")
with open("labels.json", "r") as f:
label2id = json.load(f)
id2label = {v:k for k,v in label2id.items()}
return clf, embedder, id2label
clf, embedder, id2label = load_model()
def predict_intent(text: str):
emb = embedder.encode([text])[0]
df_emb = pd.DataFrame([emb])
pred = clf.predict(df_emb)[0]
return id2label[pred]