Spaces:
Paused
Paused
File size: 2,317 Bytes
29fd274 15fe7f7 29fd274 4c32d46 29fd274 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 | 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)
|