import os from pathlib import Path from dataclasses import dataclass import re import torch import numpy as np import sys sys.path.insert(0, str(Path(__file__).parent.parent)) try: from .dataset import ( Vocabulary, tokenize, DISTORTION_LABELS, CRISIS_LABELS, PAD_TOKEN, SOS_TOKEN, EOS_TOKEN, ) from .models import AnupamaModel except ImportError: from dataset import ( Vocabulary, tokenize, DISTORTION_LABELS, CRISIS_LABELS, PAD_TOKEN, SOS_TOKEN, EOS_TOKEN, ) from model.models import AnupamaModel @dataclass class ClassifierOutputs: crisis_label: str # "safe" | "at_risk" | "crisis" crisis_probs: list[float] mood_score: int # 1–5 valence: float # continuous 1–5 distortion: str # distortion label or "none" distortion_probs: list[float] @dataclass class EngineResponse: text: str classifiers: ClassifierOutputs is_crisis: bool conditioning_tokens: list[str] def tokens_to_sentence(tokens: list[str]) -> str: """Detokenize: join tokens and clean up spacing around punctuation.""" text = " ".join(tokens) text = re.sub(r" ([?.!,])", r"\1", text) text = re.sub(r"\s+", " ", text).strip() if text and text[0].islower(): text = text[0].upper() + text[1:] # Ensure sentence ends with punctuation if text and text[-1] not in ".?!": text += "." return text class Anupama: CRISIS_PROTOCOL = ( "I hear you, and I'm genuinely concerned about your safety right now. " "Please reach out to a crisis line — they're available 24/7:\n\n" "988 Suicide & Crisis Lifeline: Call or text 988 (US)\n" "Crisis Text Line: Text HOME to 741741\n" "Emergency: Call 911 if you're in immediate danger\n\n" "You don't have to go through this alone." ) MODE_COND_TOKENS = { "support": "", "cbt": "", "intake": "", } def __init__( self, model: AnupamaModel, vocab: Vocabulary, device: torch.device, max_gen_len: int = 80, temperature: float = 0.85, top_p: float = 0.92, ): self.model = model.to(device) self.model.eval() self.vocab = vocab self.device = device self.max_gen_len = max_gen_len self.temperature = temperature self.top_p = top_p @classmethod def load(cls, checkpoint_dir: str, device_str: str = "auto") -> "Anupama": """Load model from a training checkpoint directory.""" if device_str == "auto": device = torch.device("cuda" if torch.cuda.is_available() else "cpu") else: device = torch.device(device_str) ckpt_dir = Path(checkpoint_dir) vocab = Vocabulary.load(ckpt_dir / "vocab.pkl") embed_matrix = np.load(ckpt_dir / "embed_matrix.npy") model = AnupamaModel(embed_matrix, len(vocab), pad_idx=vocab.pad_idx) # Load best model or fall back to final ckpt_path = ckpt_dir / "best_model.pt" if not ckpt_path.exists(): ckpt_path = ckpt_dir / "final_model.pt" ckpt = torch.load(ckpt_path, map_location=device) model.load_state_dict(ckpt["model_state"]) print(f"[Engine] Loaded model from {ckpt_path}") return cls(model, vocab, device) # ── Tokenize & encode ────────────────────────────────────────────────── def _encode(self, text: str): tokens = tokenize(text) ids = self.vocab.encode(tokens) if not ids: ids = [self.vocab.unk_idx] id_tensor = torch.tensor([ids], dtype=torch.long, device=self.device) lengths = torch.tensor([len(ids)]) return id_tensor, lengths @torch.no_grad() def classify(self, text: str) -> ClassifierOutputs: id_tensor, lengths = self._encode(text) # Crisis crisis_labels, crisis_probs = self.model.crisis.predict(id_tensor, lengths) c_label = CRISIS_LABELS[crisis_labels[0].item()] c_probs = crisis_probs[0].tolist() # Sentiment mood_labels, valence, sent_probs = self.model.sentiment.predict(id_tensor, lengths) mood = mood_labels[0].item() v = valence[0].item() # Distortion dist_labels, dist_probs = self.model.distortion.predict(id_tensor, lengths) d_label = DISTORTION_LABELS[dist_labels[0].item()] d_probs = dist_probs[0].tolist() return ClassifierOutputs( crisis_label=c_label, crisis_probs=c_probs, mood_score=mood, valence=round(v, 2), distortion=d_label, distortion_probs=d_probs, ) # ── Build conditioning tokens from classifier outputs ────────────────── def _build_cond_tokens(self, cls_out: ClassifierOutputs, mode: str) -> list[str]: cond = [] # Crisis token cond.append(f"<{cls_out.crisis_label.upper()}>") # Mood token cond.append(f"") # Mode token cond.append(self.MODE_COND_TOKENS.get(mode, "")) # Distortion signal if cls_out.distortion != "none": cond.append("") else: cond.append("") return cond # ── Generate response ────────────────────────────────────────────────── @torch.no_grad() def _generate(self, text: str, cond_tokens: list[str]) -> str: id_tensor, lengths = self._encode(text) cond_ids = [ self.vocab.word2idx.get(tok, self.vocab.unk_idx) for tok in cond_tokens ] token_ids = self.model.generator.generate( src_ids=id_tensor, src_lengths=lengths, cond_ids=cond_ids, sos_idx=self.vocab.sos_idx, eos_idx=self.vocab.eos_idx, max_len=self.max_gen_len, temperature=self.temperature, top_p=self.top_p, ) tokens = self.vocab.decode(token_ids) # Strip special tokens that leaked through tokens = [t for t in tokens if not (t.startswith("<") and t.endswith(">"))] return tokens_to_sentence(tokens) def respond(self, text: str, mode: str = "support") -> EngineResponse: """ Full pipeline: classify input → build cond tokens → generate response. Crisis inputs skip generation and return the crisis protocol. """ cls_out = self.classify(text) if cls_out.crisis_label == "crisis": return EngineResponse( text=self.CRISIS_PROTOCOL, classifiers=cls_out, is_crisis=True, conditioning_tokens=[], ) cond_tokens = self._build_cond_tokens(cls_out, mode) response_text = self._generate(text, cond_tokens) # Safety fallback: if generated response is too short or empty if len(response_text.strip()) < 10: response_text = self._fallback_response(cls_out, mode) return EngineResponse( text=response_text, classifiers=cls_out, is_crisis=False, conditioning_tokens=cond_tokens, ) def _fallback_response(self, cls_out: ClassifierOutputs, mode: str) -> str: """Rule-based fallback when generation produces low-quality output.""" if cls_out.mood_score <= 2: return ("It sounds like you're going through something really difficult. " "I'm here to listen. Would you like to share more about what's been happening?") if cls_out.distortion != "none": return ("I notice there might be some challenging thought patterns at play. " "Let's slow down and look at this together — what's the situation you're facing?") return ("Thank you for sharing that with me. " "How long have you been feeling this way?") if __name__ == "__main__": import sys if len(sys.argv) < 2: print("Usage: python engine.py ") sys.exit(1) engine = Anupama.load(sys.argv[1]) test_inputs = [ ("I've been feeling really anxious about my job and can't sleep", "support"), ("I always fail everything, I'm just a useless person", "cbt"), ("I've been struggling for about three weeks now", "intake"), ("I don't see any point in going on", "support"), # should trigger crisis ] for text, mode in test_inputs: print(f"\n{'─'*60}") print(f"User [{mode}]: {text}") result = engine.respond(text, mode) print(f"Crisis: {result.is_crisis} | " f"Mood: {result.classifiers.mood_score} | " f"Distortion: {result.classifiers.distortion}") print(f"Cond: {result.conditioning_tokens}") print(f"Bot: {result.text}")