Spaces:
Paused
Paused
| from pathlib import Path | |
| from enum import Enum | |
| from app.config import Settings | |
| class StrEnum(str, Enum): | |
| pass | |
| class Intent(StrEnum): | |
| MATCH_STATISTICS = "MATCH_STATISTICS" | |
| MATCH_RECAP = "MATCH_RECAP" | |
| MOMENTUM_ANALYSIS = "MOMENTUM_ANALYSIS" | |
| PLAYER_COMPARISON = "PLAYER_COMPARISON" | |
| CRICKET_RULES = "CRICKET_RULES" | |
| GENERAL_MATCH_QUERY = "GENERAL_MATCH_QUERY" | |
| def classify(question: str) -> Intent: | |
| text = question.lower() | |
| if any(word in text for word in ("rule", "no-ball", "noball", "wide", "drs", "powerplay", "lbw")): | |
| return Intent.CRICKET_RULES | |
| if any(word in text for word in ("compare", "versus", "vs", "partnership")): | |
| return Intent.PLAYER_COMPARISON | |
| if any(word in text for word in ("why", "momentum", "dropped", "surge", "pressure", "run rate")): | |
| return Intent.MOMENTUM_ANALYSIS | |
| if any(word in text for word in ("what happened", "last", "recap", "catch me up")): | |
| return Intent.MATCH_RECAP | |
| if any(word in text for word in ("how many", "score", "runs", "wickets", "overs", "required")): | |
| return Intent.MATCH_STATISTICS | |
| return Intent.GENERAL_MATCH_QUERY | |
| class IntentClassifier: | |
| def __init__(self, settings: Settings): | |
| self._model = None | |
| self._tokenizer = None | |
| self._labels = list(Intent) | |
| self.source = "heuristic_fallback" | |
| if settings.intent_model_path and Path(settings.intent_model_path).exists(): | |
| try: | |
| from transformers import AutoModelForSequenceClassification, AutoTokenizer | |
| self._tokenizer = AutoTokenizer.from_pretrained(settings.intent_model_path) | |
| self._model = AutoModelForSequenceClassification.from_pretrained(settings.intent_model_path) | |
| self.source = "peft_model" | |
| except (ImportError, OSError, ValueError): | |
| self._model = None | |
| def predict(self, question: str) -> Intent: | |
| if self._model is not None and self._tokenizer is not None: | |
| import torch | |
| encoded = self._tokenizer(question, return_tensors="pt", truncation=True) | |
| with torch.no_grad(): | |
| label_id = int(self._model(**encoded).logits.argmax(dim=-1).item()) | |
| return self._labels[label_id] | |
| return classify(question) | |