File size: 579 Bytes
c1be4e0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22

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]