models.py

"""Swappable text-model interface for Spend Elegy.

One seam for "which LLM normalizes the receipt + how to call it", selected by the
NORMALIZER_MODEL_ID env var:

  - Space / Modal (CUDA): nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16 — loaded via
    transformers' NATIVE NemotronH class (transformers>=5.8) + the `kernels`
    library (prebuilt Hub Mamba kernels); no trust_remote_code.
  - Local (Apple Silicon / CPU): Qwen2.5-3B-Instruct — a drop-in, since Nano's
    Mamba kernels are CUDA-only and can't run on MPS.

Both are plain causal LMs loaded and called identically, so swapping is just the
model id. The model is loaded at import (module scope) so a ZeroGPU @spaces.GPU
fork shares the already-resident weights.
"""

from __future__ import annotations

import os

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

NORMALIZER_MODEL_ID = os.environ.get(
    "NORMALIZER_MODEL_ID", "nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16"
)
MAX_NEW_TOKENS = int(os.environ.get("MAX_NEW_TOKENS", "1024"))


def _device() -> str:
    if torch.cuda.is_available():  # True on ZeroGPU (CUDA emulation at module scope)
        return "cuda"
    if torch.backends.mps.is_available():
        return "mps"
    return "cpu"


DEVICE = _device()
DTYPE = torch.float32 if DEVICE == "cpu" else torch.bfloat16


class TextModel:
    """A loaded causal-LM normalizer behind a single ``generate(messages)`` call."""

    def __init__(self, model_id: str = NORMALIZER_MODEL_ID):
        self.model_id = model_id
        print(f"[spend-elegy] loading text model {model_id} on {DEVICE} ({DTYPE})")
        self.tokenizer = AutoTokenizer.from_pretrained(model_id)
        self.model = (
            AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=DTYPE)
            .to(DEVICE)
            .eval()
        )

    def generate(self, messages: list[dict[str, str]]) -> str:
        """Run chat ``messages`` and return only the newly generated text."""
        template_kwargs = dict(
            tokenize=True,
            add_generation_prompt=True,
            return_tensors="pt",
            return_dict=True,  # -> BatchEncoding (input_ids + attention_mask)
        )
        try:
            inputs = self.tokenizer.apply_chat_template(
                messages,
                enable_thinking=False,  # Nemotron Nano: reasoning off for clean JSON
                **template_kwargs,
            )
        except TypeError:
            # Templates that don't accept enable_thinking (e.g. Qwen2.5).
            inputs = self.tokenizer.apply_chat_template(messages, **template_kwargs)
        inputs = inputs.to(self.model.device)
        input_len = inputs["input_ids"].shape[-1]
        pad_id = self.tokenizer.pad_token_id or self.tokenizer.eos_token_id
        outputs = self.model.generate(
            **inputs,
            max_new_tokens=MAX_NEW_TOKENS,
            do_sample=False,
            eos_token_id=self.tokenizer.eos_token_id,
            pad_token_id=pad_id,
        )
        new_tokens = outputs[0][input_len:]
        return self.tokenizer.decode(new_tokens, skip_special_tokens=True)


# Module-scope singleton: loaded once, shared across ZeroGPU forks.
text_model = TextModel()