llama_backend.py

"""Primary backend: Gemma 4 E2B GGUF via llama.cpp.

- Text extraction runs in-process with `llama-cpp-python` using a JSON-schema
  grammar (GBNF), so the model can only ever emit valid JSON.
- Audio transcription runs through the `llama-mtmd-cli` binary (llama.cpp PR #21421),
  which decodes Gemma's audio encoder. The mmproj must be BF16 per that PR.

Anything missing (the GGUF weights, the BF16 mmproj, the binary, or the python
binding) raises `BackendUnavailable`, and the callers fall back to the Gemma
transformers path in gemma.py.
"""

from __future__ import annotations

import shutil
import subprocess
import threading

from config import (
    LLAMA_CTX,
    LLAMA_GPU_LAYERS,
    LLAMA_MAX_TOKENS,
    LLAMA_MMPROJ_PATH,
    LLAMA_MODEL_PATH,
    LLAMA_MTMD_CLI,
    LLAMA_TEMPERATURE,
    LLAMA_THREADS,
)


class BackendUnavailable(Exception):
    """Raised when the llama.cpp path can't run, so callers fall back to transformers."""


_llm = None
_lock = threading.Lock()

_ASR_PROMPT = "Transcribe this audio exactly. Output only the transcription text."


def _get_llm():
    """Lazily load the GGUF model for in-process text generation."""
    global _llm
    if _llm is None:
        with _lock:
            if _llm is None:
                if not LLAMA_MODEL_PATH.exists():
                    raise BackendUnavailable(f"GGUF not found: {LLAMA_MODEL_PATH}")
                try:
                    from llama_cpp import Llama
                except ImportError as e:
                    raise BackendUnavailable(f"llama-cpp-python not installed: {e}")
                _llm = Llama(
                    model_path=str(LLAMA_MODEL_PATH),
                    n_ctx=LLAMA_CTX,
                    n_gpu_layers=LLAMA_GPU_LAYERS,
                    n_threads=LLAMA_THREADS,
                    verbose=False,
                )
    return _llm


def extract_json(messages: list[dict], schema: dict) -> str:
    """Grammar-constrained chat completion. Returns the raw JSON string."""
    llm = _get_llm()
    resp = llm.create_chat_completion(
        messages=messages,
        response_format={"type": "json_object", "schema": schema},
        temperature=LLAMA_TEMPERATURE,
        max_tokens=LLAMA_MAX_TOKENS,
    )
    return resp["choices"][0]["message"]["content"]


def _resolve_cli() -> str:
    cli = shutil.which(LLAMA_MTMD_CLI)
    if cli:
        return cli
    from pathlib import Path

    if Path(LLAMA_MTMD_CLI).exists():
        return LLAMA_MTMD_CLI
    raise BackendUnavailable(f"llama-mtmd-cli not found: {LLAMA_MTMD_CLI}")


def _clean_cli_output(stdout: str) -> str:
    """Best-effort extraction of the generated text from llama-mtmd-cli stdout.

    The CLI interleaves a little logging with the generation; we drop obvious log
    lines and the echoed prompt. May need tuning once run against the real binary.
    """
    text = stdout
    if _ASR_PROMPT in text:                      # take everything after the prompt echo
        text = text.split(_ASR_PROMPT, 1)[1]
    keep = []
    for line in text.splitlines():
        s = line.strip()
        if not s:
            continue
        low = s.lower()
        if any(tok in low for tok in ("llama_", "ggml", "load time", "sample time",
                                      "prompt eval", "eval time", "total time",
                                      "main:", "encode", "mtmd_")):
            continue
        keep.append(s)
    return " ".join(keep).strip()


def transcribe_audio(audio_path: str) -> str:
    """Transcribe via the llama-mtmd-cli binary using Gemma's audio encoder."""
    if not LLAMA_MODEL_PATH.exists():
        raise BackendUnavailable(f"GGUF not found: {LLAMA_MODEL_PATH}")
    if not LLAMA_MMPROJ_PATH.exists():
        raise BackendUnavailable(f"BF16 mmproj not found: {LLAMA_MMPROJ_PATH}")
    cli = _resolve_cli()

    cmd = [
        cli,
        "-m", str(LLAMA_MODEL_PATH),
        "--mmproj", str(LLAMA_MMPROJ_PATH),
        "--audio", audio_path,
        "-p", _ASR_PROMPT,
        "--temp", "0",
        "-ngl", str(LLAMA_GPU_LAYERS),
    ]
    try:
        proc = subprocess.run(cmd, capture_output=True, text=True, timeout=300)
    except (OSError, subprocess.SubprocessError) as e:
        raise BackendUnavailable(f"llama-mtmd-cli failed to run: {e}")
    if proc.returncode != 0:
        raise BackendUnavailable(f"llama-mtmd-cli exit {proc.returncode}: {proc.stderr[-300:]}")
    return _clean_cli_output(proc.stdout)