inference.py

"""
inference.py
------------
Model loading and NPC response generation for the NPCAlign demo.

ZeroGPU two-phase loading:
  Phase 1 — module import (no GPU required):
    · hf_login for gated Llama base model
    · Tokenizer load (no tensor ops intercepted by ZeroGPU)
    · snapshot_download of both LoRA adapter repos (HTTP only)

  Phase 2 — first @spaces.GPU call (GPU context active):
    · AutoModelForCausalLM.from_pretrained  ← requires GPU context
    · PeftModel.from_pretrained             ← requires GPU context
    The spaces package patches torch.__torch_dispatch__ globally and raises
    "No CUDA GPUs are available" if safetensors weights are loaded outside a
    GPU context, so all weight loading must happen here (lazy, once only).

  Subsequent @spaces.GPU calls:
    · Model already in CPU RAM → .to(device) → inference → .to("cpu")
"""

import os
import re
import traceback

import torch
from huggingface_hub import login as hf_login, snapshot_download
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer

try:
    import spaces
except ImportError:
    class spaces:  # noqa: N801
        @staticmethod
        def GPU(fn):
            return fn


# ── Model identifiers ──────────────────────────────────────────────────────────
BASE_MODEL_ID  = "meta-llama/Meta-Llama-3.1-8B-Instruct"
SFT_ADAPTER_ID = "HermitQ/NPCAlign-SFT"
DPO_ADAPTER_ID = "HermitQ/NPCAlign-DPO"

# ── Turn limits ────────────────────────────────────────────────────────────────
MAX_TURNS      = 16
SOFT_HINT_TURN = 10

# ── Token budget (mirrors training max_length = 3072) ─────────────────────────
_MAX_TOKENS     = 3072
_MAX_NEW_TOKENS = 256
_SYS_BUDGET     = 700
_HISTORY_BUDGET = _MAX_TOKENS - _SYS_BUDGET - _MAX_NEW_TOKENS  # ≈ 2116

# ── System prompt ──────────────────────────────────────────────────────────────
SYSTEM_TEMPLATE = (
    "You are {name}.\n"
    "Background: {background}\n"
    "Current Location: {location}\n"
    "Quest: {quest}\n\n"
    "Roleplaying Instructions:\n"
    "- Speak using appropriate tone and vocabulary\n"
    "- Reference your background and current surroundings naturally\n"
    "- Keep responses conversational and authentic\n"
    "- React to the player's words and intentions\n"
    "- Your first response should be a greeting to the player\n"
    "- Once the quest has been fully explained and the player shows readiness, "
    "or if the conversation becomes repetitive, proactively bring it to a "
    "natural close with a farewell.\n"
)

_SOFT_HINT = (
    "[Turn {turn}/{max_turns}] If the quest information has been conveyed, "
    "begin steering toward a natural, in-character farewell. "
    "When you close the conversation, append <End> at the end of your message."
)

# ── Global state ───────────────────────────────────────────────────────────────
_tokenizer: AutoTokenizer | None = None
_model: PeftModel | None = None
_model_loaded: bool = False      # True after Phase 2 completes
_sft_local: str | None = None    # Local path set by Phase 1
_dpo_local: str | None = None
_preload_error: str | None = None  # Phase-1 error message


# ──────────────────────────────────────────────────────────────────────────────
# Phase 1: runs at module import time (no GPU)
# ──────────────────────────────────────────────────────────────────────────────

def _preload() -> None:
    """
    Phase 1: load tokenizer and download adapter files.
    No weight tensors are loaded here — safe to run outside a GPU context.
    """
    global _tokenizer, _sft_local, _dpo_local, _preload_error

    hf_token = os.environ.get("HF_TOKEN")
    print(f"[INFO] Phase 1 | HF_TOKEN present: {bool(hf_token)}")

    try:
        if hf_token:
            hf_login(token=hf_token, add_to_git_credential=False)
            print("[INFO] hf_login() succeeded ✓")
        else:
            print("[WARN] No HF_TOKEN — Llama base model download may fail if gated")

        print("[INFO] Loading tokenizer...")
        _tokenizer = AutoTokenizer.from_pretrained(
            BASE_MODEL_ID, token=hf_token, use_fast=True
        )
        if _tokenizer.pad_token_id is None:
            _tokenizer.pad_token_id = _tokenizer.eos_token_id
        print("[INFO] Tokenizer loaded ✓")

        # Adapter repos are public — HTTP download only, no tensor ops.
        print("[INFO] Downloading SFT adapter files...")
        _sft_local = snapshot_download(
            repo_id=SFT_ADAPTER_ID, local_dir="/tmp/npcalign_sft"
        )
        print("[INFO] Downloading DPO adapter files...")
        _dpo_local = snapshot_download(
            repo_id=DPO_ADAPTER_ID, local_dir="/tmp/npcalign_dpo"
        )
        print("[INFO] Phase 1 complete — adapter files ready ✓")

    except Exception:
        _preload_error = traceback.format_exc()
        print(f"[ERROR] Phase 1 failed:\n{_preload_error}")


_preload()


# ──────────────────────────────────────────────────────────────────────────────
# Phase 2: runs lazily on the first @spaces.GPU call
# ──────────────────────────────────────────────────────────────────────────────

def _load_model_weights() -> None:
    """
    Phase 2: load base model + LoRA adapters inside a GPU context.
    Called once from generate_response() on its first invocation.
    ZeroGPU's torch patches require CUDA to be available for safetensors loading.
    """
    global _model, _model_loaded

    hf_token = os.environ.get("HF_TOKEN")
    print("[INFO] Phase 2 | Loading base model (GPU context active)...")
    base = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL_ID,
        torch_dtype=torch.bfloat16,
        device_map="auto",
        token=hf_token,
    )

    print("[INFO] Attaching SFT + DPO LoRA adapters...")
    _model = PeftModel.from_pretrained(base, _sft_local, adapter_name="sft")
    _model.load_adapter(_dpo_local, adapter_name="dpo")
    _model.eval()
    _model_loaded = True
    print("[INFO] Phase 2 complete — model ready ✓")


# ── Public helpers ─────────────────────────────────────────────────────────────

def build_npc_fields(
    name: str, personality: str, experience: str,
    occupation: str, location: str, quest: str,
) -> dict:
    """Combine form inputs into npc_fields dict for SYSTEM_TEMPLATE."""
    parts = [personality.strip()]
    if occupation.strip():
        parts.append(occupation.strip())
    parts.append(experience.strip())
    return {
        "name":       name.strip(),
        "background": " ".join(p for p in parts if p),
        "location":   location.strip(),
        "quest":      quest.strip(),
    }


def build_messages(npc_fields: dict, history: list[dict], turn: int) -> list[dict]:
    """
    Assemble the messages list for one inference call with sliding-window history.
    Injects a soft-hint system message from turn SOFT_HINT_TURN onward.
    """
    msgs = [{"role": "system", "content": SYSTEM_TEMPLATE.format(**npc_fields)}]

    trimmed = list(history)
    while len(trimmed) > 1 and _tokenizer is not None:
        n = len(_tokenizer.encode(
            " ".join(m["content"] for m in trimmed), add_special_tokens=False
        ))
        if n <= _HISTORY_BUDGET:
            break
        trimmed = trimmed[2:]

    msgs.extend(trimmed)

    if turn >= SOFT_HINT_TURN:
        msgs.append({
            "role":    "system",
            "content": _SOFT_HINT.format(turn=turn, max_turns=MAX_TURNS),
        })

    return msgs


def detect_end(text: str) -> tuple[str, bool]:
    """Strip <End> tag and return (clean_text, conversation_ended)."""
    ended   = bool(re.search(r"<[Ee]nd>", text))
    cleaned = re.sub(r"\s*<[Ee]nd>", "", text).strip()
    return cleaned, ended


def _truncate_at_sentence(text: str) -> str:
    """
    If the text ends mid-sentence (no terminal punctuation), truncate back to
    the last complete sentence so the response never feels cut off.
    Leaves the text unchanged if it already ends with . ! ? or a closing *.
    """
    # Already ends cleanly
    if re.search(r'[.!?*]["\']?\s*$', text):
        return text
    # Find the last sentence-ending punctuation followed by whitespace or end
    match = None
    for m in re.finditer(r'[.!?]["\']?(?=\s|$)', text):
        match = m
    if match and match.end() > len(text) // 3:
        return text[:match.end()].strip()
    return text  # fallback: return as-is


@spaces.GPU(duration=60)
def generate_response(adapter_name: str, messages: list[dict]) -> str:
    """
    Generate one NPC response using the specified LoRA adapter.

    On first call: loads base model + LoRA adapters (Phase 2, ~60 s).
    Subsequent calls: model already in CPU RAM, just move to GPU and infer.

    Args:
        adapter_name : "sft" or "dpo"
        messages     : output of build_messages()
    """
    if _preload_error:
        raise RuntimeError(
            f"Startup failed (Phase 1):\n{_preload_error}"
        )

    if not _model_loaded:
        try:
            _load_model_weights()
        except Exception:
            err = traceback.format_exc()
            print(f"[ERROR] Phase 2 failed:\n{err}")
            raise RuntimeError(f"Model weight loading failed:\n{err}")

    _model.set_adapter(adapter_name)

    device = "cuda" if torch.cuda.is_available() else "cpu"
    _model.to(device)

    # apply_chat_template may return BatchEncoding or a plain tensor depending
    # on the transformers version; extract the raw tensor in either case.
    encoded = _tokenizer.apply_chat_template(
        messages,
        tokenize=True,
        add_generation_prompt=True,
        return_tensors="pt",
    )
    input_ids = (encoded["input_ids"] if hasattr(encoded, "__getitem__") and not isinstance(encoded, torch.Tensor)
                 else encoded).to(device)

    with torch.no_grad():
        output_ids = _model.generate(
            input_ids=input_ids,
            max_new_tokens=_MAX_NEW_TOKENS,
            do_sample=True,
            temperature=0.7,
            top_p=0.9,
            pad_token_id=_tokenizer.eos_token_id,
            eos_token_id=_tokenizer.eos_token_id,
        )

    new_tokens = output_ids[0][input_ids.shape[1]:]
    result = _tokenizer.decode(new_tokens, skip_special_tokens=True).strip()
    result = _truncate_at_sentence(result)

    # Return model to CPU so ZeroGPU can cleanly release the GPU allocation.
    _model.to("cpu")
    return result