src/apps/p5_memory_quilt/llm_backend.py

from __future__ import annotations

"""Strict local LLM prompt-rewriter backend.

This module is import-safe: it only imports transformers/torch when invoked.
It runs offline-only (local_files_only=True) and raises if the local checkpoint
is not available or if generation does not produce the expected JSON payload.
"""

from dataclasses import dataclass
from pathlib import Path
import json
import os
import time


@dataclass(frozen=True)
class LLMConfig:
    model_id: str
    checkpoint_path: str
    max_new_tokens: int = 180
    temperature: float = 0.4
    top_p: float = 0.9


def available(checkpoint_path: str) -> bool:
    path = checkpoint_path.strip()
    return bool(path and Path(path).exists())


def _build_instruction(memory_text: str, location_tag: str, style: str) -> str:
    location = location_tag.strip() or "(none)"
    return (
        "You are a prompt rewriter for an offline image generation demo. "
        "Given a short neighborhood memory, produce JSON with keys: "
        "caption (<=4 words), story (1 sentence), flux_prompt (1 sentence, vivid), style_hint (short).\n\n"
        f"STYLE: {style}\n"
        f"LOCATION_TAG: {location}\n"
        f"MEMORY: {memory_text.strip()}\n\n"
        "Return ONLY valid JSON."
    )


def try_rewrite(memory_text: str, location_tag: str, style: str, cfg: LLMConfig) -> tuple[dict[str, str], dict[str, object]]:
    """Run a local checkpoint and return strict rewrite payload + logging metadata."""

    started_at = time.perf_counter()

    try:
        import torch  # type: ignore
    except Exception as exc:
        raise RuntimeError("torch is not installed; cannot run LLM backend") from exc

    try:
        from transformers import AutoModelForCausalLM, AutoTokenizer  # type: ignore
    except Exception as exc:
        raise RuntimeError("transformers is not installed; cannot run LLM backend") from exc

    os.environ.setdefault("HF_HUB_OFFLINE", "1")
    os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")

    tokenizer = AutoTokenizer.from_pretrained(cfg.checkpoint_path, local_files_only=True, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(
        cfg.checkpoint_path,
        local_files_only=True,
        trust_remote_code=True,
        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
        device_map="auto" if torch.cuda.is_available() else None,
    )

    prompt = _build_instruction(memory_text, location_tag, style)
    inputs = tokenizer(prompt, return_tensors="pt")
    if torch.cuda.is_available():
        inputs = {k: v.to("cuda") for k, v in inputs.items()}

    prompt_tokens = int(inputs["input_ids"].shape[-1])
    outputs = model.generate(
        **inputs,
        max_new_tokens=cfg.max_new_tokens,
        do_sample=True,
        temperature=cfg.temperature,
        top_p=cfg.top_p,
        eos_token_id=getattr(tokenizer, "eos_token_id", None),
    )

    decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
    start = decoded.rfind("{")
    end = decoded.rfind("}")
    if start == -1 or end == -1 or end <= start:
        raise RuntimeError("LLM output did not contain JSON object")

    raw_json = decoded[start : end + 1]
    try:
        payload = json.loads(raw_json)
    except Exception as exc:
        raise RuntimeError(f"Failed to parse JSON from LLM output: {exc}") from exc

    def _require(key: str) -> str:
        value = payload.get(key)
        if not isinstance(value, str):
            raise RuntimeError(f"LLM output missing required '{key}' field")
        value = value.strip()
        if not value:
            raise RuntimeError(f"LLM output field '{key}' was empty")
        return value

    normalized = {
        "caption": _require("caption"),
        "story": _require("story"),
        "flux_prompt": _require("flux_prompt"),
        "style_hint": _require("style_hint"),
    }

    generated_tokens = max(0, int(outputs.shape[-1]) - prompt_tokens)
    elapsed_ms = round((time.perf_counter() - started_at) * 1000.0, 2)
    generation_stats = {
        "prompt_tokens": prompt_tokens,
        "generated_tokens": generated_tokens,
        "max_new_tokens": cfg.max_new_tokens,
        "temperature": cfg.temperature,
        "top_p": cfg.top_p,
        "elapsed_ms": elapsed_ms,
        "device": "cuda" if torch.cuda.is_available() else "cpu",
    }
    meta = {
        "adapter_name": "local-transformers",
        "backend": "local-transformers",
        "model_id": cfg.model_id,
        "checkpoint_path": cfg.checkpoint_path,
        "generation_stats": generation_stats,
    }
    return normalized, meta