Parinith-reddy/hinglish-chatbot

1"""2Predictor — FP16 PyTorch version3Loads the compressed FP16 model for deployment.4"""5 6import json, os7import torch8import torch.nn.functional as F9from transformers import AutoTokenizer, AutoModelForSequenceClassification10from normalizer_v2 import normalize11from online_learner import OnlineLearner12 13MODEL_DIR = "hinglish_deploy"14DEVICE    = torch.device("cpu")15 16ONLINE_BASE_WEIGHT = 0.1517ONLINE_MAX_WEIGHT  = 0.4018 19 20class Predictor:21    def __init__(self):22        print(f"[Predictor] Loading FP16 model...")23 24        self.tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR)25        self.model = AutoModelForSequenceClassification.from_pretrained(26            MODEL_DIR,27            torch_dtype=torch.float16,28        )29        self.model.to(DEVICE)30        self.model.eval()31 32        with open(os.path.join(MODEL_DIR, "label_map.json"), encoding="utf-8") as f:33            maps = json.load(f)34        self.label2id = maps["label2id"]35        self.id2label = {int(k): v for k, v in maps["id2label"].items()}36        self.labels   = list(self.label2id.keys())37 38        self.online = OnlineLearner(intent_labels=self.labels)39        print(f"[Predictor] Ready — {len(self.labels)} intents")40 41    def _predict_proba(self, text: str) -> dict:42        enc = self.tokenizer(43            text,44            return_tensors="pt",45            max_length=128,46            padding="max_length",47            truncation=True48        ).to(DEVICE)49 50        with torch.no_grad():51            logits = self.model(**enc).logits.float()52 53        probs = F.softmax(logits, dim=-1).squeeze().cpu().numpy()54        return {self.id2label[i]: float(p) for i, p in enumerate(probs)}55 56    def predict(self, raw_text: str, top_k: int = 3) -> dict:57        normalized  = normalize(raw_text)58        model_probs = self._predict_proba(normalized)59        model_top   = max(model_probs, key=model_probs.get)60        model_conf  = model_probs[model_top]61 62        online_probs = self.online.predict_proba(normalized)63        if online_probs is not None:64            n      = self.online.feedback_count65            weight = min(ONLINE_BASE_WEIGHT + (n / 5000) * ONLINE_MAX_WEIGHT,66                         ONLINE_MAX_WEIGHT)67            combined = {68                intent: (1 - weight) * model_probs.get(intent, 0.0)69                       + weight       * online_probs.get(intent, 0.0)70                for intent in self.labels71            }72        else:73            combined = model_probs74 75        sorted_intents = sorted(combined.items(), key=lambda x: x[1], reverse=True)76        top_intent, top_conf = sorted_intents[0]77 78        return {79            "normalized":  normalized,80            "top_intent":  top_intent,81            "confidence":  round(top_conf * 100, 1),82            "top_k":       [(i, round(c * 100, 1)) for i, c in sorted_intents[:top_k]],83            "muril_pred":  model_top,84            "muril_conf":  round(model_conf * 100, 1),85        }86 87    def feedback(self, raw_text: str, muril_pred: str,88                 muril_conf: float, chosen_intent: str):89        normalized = normalize(raw_text)90        self.online.log_feedback(normalized, muril_pred, muril_conf, chosen_intent)91        self.online.learn(normalized, chosen_intent)92