Create services/intent_classifier_client.py

services/intent_classifier_client.py

@@ -0,0 +1,132 @@
+import os
+import re
+import threading
+from typing import Optional
+
+import torch
+from transformers import AutoTokenizer, Gemma3ForCausalLM
+
+from knowledge.classifier_prompt import (
+    build_system_prompt,
+    get_allowed_intents_for_state,
+)
+
+MODEL_ID = os.getenv("MODEL_ID", "google/gemma-3-1b-it")
+MAX_NEW_TOKENS = int(os.getenv("MAX_NEW_TOKENS", "12"))
+HF_TOKEN = os.getenv("HF_TOKEN")
+ENABLE_MODEL_CLASSIFIER = os.getenv("ENABLE_MODEL_CLASSIFIER", "true").lower() == "true"
+
+_model = None
+_tokenizer = None
+_model_lock = threading.Lock()
+
+
+def _normalize_label(text: str, allowed_intents: list[str]) -> str:
+    cleaned = (text or "").strip().lower()
+    cleaned = cleaned.replace("```", "").replace("`", "").strip()
+
+    for intent in allowed_intents:
+        if re.search(rf"\b{re.escape(intent.lower())}\b", cleaned):
+            return intent
+
+    return "unclear"
+
+
+def _load_model_once():
+    global _model, _tokenizer
+
+    if _model is not None and _tokenizer is not None:
+        return _model, _tokenizer
+
+    with _model_lock:
+        if _model is not None and _tokenizer is not None:
+            return _model, _tokenizer
+
+        if not HF_TOKEN:
+            raise RuntimeError("HF_TOKEN is missing. Add it in Hugging Face Space Secrets.")
+
+        print(f"[intent-classifier] loading model: {MODEL_ID}")
+
+        _tokenizer = AutoTokenizer.from_pretrained(
+            MODEL_ID,
+            token=HF_TOKEN
+        )
+
+        _model = Gemma3ForCausalLM.from_pretrained(
+            MODEL_ID,
+            token=HF_TOKEN
+        ).eval()
+
+        print("[intent-classifier] model loaded successfully")
+
+    return _model, _tokenizer
+
+
+def _run_generation(user_message: str, state: str, flow_data: Optional[dict] = None) -> dict:
+    model, tokenizer = _load_model_once()
+
+    allowed_intents = get_allowed_intents_for_state(state)
+    system_prompt = build_system_prompt(
+        state=state,
+        flow_data=flow_data or {},
+        allowed_intents=allowed_intents,
+    )
+
+    messages = [
+        {"role": "system", "content": system_prompt},
+        {"role": "user", "content": user_message},
+    ]
+
+    prompt = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True
+    )
+
+    inputs = tokenizer(prompt, return_tensors="pt")
+
+    with torch.inference_mode():
+        generation = model.generate(
+            **inputs,
+            max_new_tokens=MAX_NEW_TOKENS,
+            do_sample=False,
+            temperature=None,
+            top_p=None,
+        )
+
+    input_len = inputs["input_ids"].shape[-1]
+    generated_tokens = generation[0][input_len:]
+    raw_output = tokenizer.decode(generated_tokens, skip_special_tokens=True).strip()
+
+    final_intent = _normalize_label(raw_output, allowed_intents)
+
+    return {
+        "intent": final_intent,
+        "raw_output": raw_output,
+        "model": MODEL_ID,
+        "allowed_intents": allowed_intents,
+    }
+
+
+def classify_message_with_model(user_message: str, state: str, flow_data: Optional[dict] = None) -> Optional[dict]:
+    """
+    Returns:
+    {
+      "intent": "...",
+      "raw_output": "...",
+      "model": "...",
+      "allowed_intents": [...]
+    }
+    or None if classifier is disabled
+    """
+    if not ENABLE_MODEL_CLASSIFIER:
+        return None
+
+    if not user_message or not user_message.strip():
+        return None
+
+    return _run_generation(
+        user_message=user_message.strip(),
+        state=state,
+        flow_data=flow_data or {},
+    )