inference.py

"""
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