gemma.py

"""Unified Gemma 4 E2B backend.

One small multimodal model does BOTH jobs: speech-to-text (audio modality) and
structured extraction (text generation). The model + processor load once and stay
resident. See config.py for the knobs.

Note: Gemma 4 is a gated Google model. Before first run:
    huggingface-cli login
and accept the license on the model page. The `AutoModelForMultimodalLM` /
`processor.parse_response` API follows the model card; adjust if your installed
transformers version exposes different names.
"""

from __future__ import annotations

import threading

from config import (
    GEMMA_DEVICE_MAP,
    GEMMA_DTYPE,
    GEMMA_MAX_NEW_TOKENS,
    GEMMA_MODEL_ID,
    GEMMA_SAMPLE,
    GEMMA_TEMPERATURE,
    GEMMA_TOP_K,
    GEMMA_TOP_P,
)

_model = None
_processor = None
_lock = threading.Lock()


def _load():
    """Load model + processor once, behind a lock."""
    global _model, _processor
    if _model is None:
        with _lock:
            if _model is None:
                from transformers import AutoModelForMultimodalLM, AutoProcessor

                _processor = AutoProcessor.from_pretrained(GEMMA_MODEL_ID)
                _model = AutoModelForMultimodalLM.from_pretrained(
                    GEMMA_MODEL_ID,
                    dtype=GEMMA_DTYPE,
                    device_map=GEMMA_DEVICE_MAP,
                )
    return _model, _processor


def _generate(messages: list[dict], max_new_tokens: int | None = None) -> str:
    """Run a chat-template generation and return the decoded reply text."""
    model, processor = _load()
    inputs = processor.apply_chat_template(
        messages,
        tokenize=True,
        return_dict=True,
        return_tensors="pt",
        add_generation_prompt=True,
        enable_thinking=False,
    ).to(model.device)
    input_len = inputs["input_ids"].shape[-1]

    gen_kwargs = {"max_new_tokens": max_new_tokens or GEMMA_MAX_NEW_TOKENS}
    if GEMMA_SAMPLE:
        gen_kwargs.update(
            do_sample=True,
            temperature=GEMMA_TEMPERATURE,
            top_p=GEMMA_TOP_P,
            top_k=GEMMA_TOP_K,
        )
    else:
        gen_kwargs["do_sample"] = False  # greedy: deterministic

    outputs = model.generate(**inputs, **gen_kwargs)
    # enable_thinking=False means no reasoning block to strip, so a plain decode with
    # special tokens removed gives clean text without relying on parse_response().
    return processor.decode(outputs[0][input_len:], skip_special_tokens=True).strip()


def transcribe_audio(audio_path: str) -> str:
    """Transcribe a short (<=30s) audio file using Gemma's audio modality."""
    messages = [{
        "role": "user",
        "content": [
            {"type": "text", "text": (
                "Transcribe this voice note verbatim in its original language. "
                "Return only the transcription text, with no commentary or labels."
            )},
            {"type": "audio", "audio": audio_path},
        ],
    }]
    return _generate(messages, max_new_tokens=512)


def generate_chat(messages: list[dict], max_new_tokens: int | None = None) -> str:
    """Text-only chat generation (used by the extractor)."""
    return _generate(messages, max_new_tokens=max_new_tokens)