""" Inference module for ASR (Whisper-small) and story Q&A (Qwen2.5-3B-Instruct). Models are loaded on-demand and cached globally for reuse. """ import logging import torch logger = logging.getLogger(__name__) # --------------------------------------------------------------------------- # ASR — Whisper-small (loaded on demand) # --------------------------------------------------------------------------- _asr_pipe = None def get_asr_pipeline(): """Load Whisper-small pipeline on first call, cache thereafter.""" global _asr_pipe if _asr_pipe is None: from transformers import pipeline logger.info("Loading Whisper-small for ASR...") _asr_pipe = pipeline( "automatic-speech-recognition", model="openai/whisper-small", device="cuda" if torch.cuda.is_available() else "cpu", torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32, ) logger.info("Whisper-small loaded.") return _asr_pipe def transcribe_audio(audio_path: str) -> str: """Transcribe an audio file to text using Whisper-small.""" if not audio_path: return "" import soundfile as sf import numpy as np audio_data, sample_rate = sf.read(audio_path, dtype="float32") # Convert stereo to mono if needed if len(audio_data.shape) > 1: audio_data = audio_data.mean(axis=1) pipe = get_asr_pipeline() result = pipe({"raw": audio_data, "sampling_rate": sample_rate}, generate_kwargs={"language": "en"}) return result.get("text", "").strip() # --------------------------------------------------------------------------- # Q&A — Qwen2.5-3B-Instruct (always loaded after first call) # --------------------------------------------------------------------------- _qa_tokenizer = None _qa_model = None def get_qa_model(): """Load Qwen2.5-3B-Instruct on first call, cache thereafter.""" global _qa_tokenizer, _qa_model if _qa_model is None: from transformers import AutoTokenizer, AutoModelForCausalLM model_id = "Qwen/Qwen2.5-3B-Instruct" logger.info("Loading %s...", model_id) _qa_tokenizer = AutoTokenizer.from_pretrained(model_id) # Check available VRAM — if less than 3GB free, use CPU use_gpu = False if torch.cuda.is_available(): free_vram = torch.cuda.get_device_properties(0).total_memory - torch.cuda.memory_allocated(0) free_vram_gb = free_vram / (1024**3) logger.info("Free VRAM: %.1f GB", free_vram_gb) if free_vram_gb >= 3.0: use_gpu = True if use_gpu: load_kwargs = {"device_map": "auto", "torch_dtype": torch.float16} # Enable FlashAttention-2 if available, else SDPA try: import flash_attn # noqa: F401 load_kwargs["attn_implementation"] = "flash_attention_2" logger.info("Using FlashAttention-2 for Q&A model.") except ImportError: load_kwargs["attn_implementation"] = "sdpa" try: from transformers import BitsAndBytesConfig load_kwargs["quantization_config"] = BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_compute_dtype=torch.float16, bnb_4bit_quant_type="nf4", ) logger.info("Using 4-bit quantization on GPU.") except Exception: logger.info("bitsandbytes unavailable, using float16 on GPU.") else: logger.info("Insufficient VRAM — loading Q&A model on CPU (float32).") load_kwargs = {"device_map": "cpu", "torch_dtype": torch.float32} _qa_model = AutoModelForCausalLM.from_pretrained(model_id, **load_kwargs) logger.info("Qwen2.5-3B-Instruct loaded on %s.", "GPU" if use_gpu else "CPU") return _qa_tokenizer, _qa_model def _get_relevant_context(paragraphs: list[str], current_idx: int, question: str) -> str: """Return full story with emphasis on current section for context.""" if not paragraphs: return "" # Build context: full story (truncated if too long) with current paragraph highlighted total_text = "\n\n".join(paragraphs) # If story is short enough (< 2000 chars), use it all if len(total_text) <= 2000: current_marker = f"\n\n[Currently reading]: {paragraphs[current_idx]}" if current_idx < len(paragraphs) else "" return total_text + current_marker # For longer stories: use top relevant paragraphs + surrounding context question_words = set(question.lower().split()) scored = [] for i, para in enumerate(paragraphs): para_words = set(para.lower().split()) overlap = len(question_words & para_words) # Boost paragraphs near current position proximity_bonus = max(0, 5 - abs(i - current_idx)) scored.append((overlap + proximity_bonus, i, para)) scored.sort(key=lambda x: x[0], reverse=True) # Take top 5 most relevant paragraphs top_paras = sorted(scored[:5], key=lambda x: x[1]) # sort by position context = "\n\n".join(s[2] for s in top_paras) # Add current paragraph marker if current_idx < len(paragraphs): context += f"\n\n[Currently reading]: {paragraphs[current_idx]}" return context def answer_story_question( question: str, paragraphs: list[str], current_idx: int = 0, ) -> str: """ Generate a short, grounded answer to a child's question about the story. Returns the answer text (1-2 sentences). """ if not question.strip(): return "" tokenizer, model = get_qa_model() context = _get_relevant_context(paragraphs, current_idx, question) if not context: context = "\n\n".join(paragraphs[:5]) messages = [ { "role": "system", "content": ( "You are a friendly storyteller answering a child's question about a bedtime story. " "Answer in 1-2 short, simple sentences using ONLY information from the story context below. " "If the story doesn't contain the answer, say so gently. " "Use warm, age-appropriate language." ), }, { "role": "user", "content": f"Story context:\n{context}\n\nChild's question: {question}", }, ] text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) inputs = tokenizer(text, return_tensors="pt").to(model.device) with torch.no_grad(): outputs = model.generate( **inputs, max_new_tokens=80, temperature=0.7, do_sample=True, pad_token_id=tokenizer.eos_token_id, ) answer_tokens = outputs[0][inputs["input_ids"].shape[1]:] answer = tokenizer.decode(answer_tokens, skip_special_tokens=True).strip() return answer