rewriter.py

"""Production rewriter — same code that will land in animoflow-web/nodes/prompt_rewriter/."""
from __future__ import annotations

import json
import os
import time
from pathlib import Path

import numpy as np
import torch
from sentence_transformers import SentenceTransformer
from transformers import AutoModelForCausalLM, AutoTokenizer

try:
    import spaces  # HF ZeroGPU decorator
    _HAS_SPACES = True
except ImportError:
    _HAS_SPACES = False

# Switchable via MODEL_REPO env var. Default Qwen2.5-1.5B-Instruct (Apache 2.0, ungated).
# Gemma-3-1B-it can be selected once HF Space has license acceptance via HF_TOKEN.
MODEL_ID = os.environ.get("MODEL_REPO", "Qwen/Qwen2.5-1.5B-Instruct")
RETRIEVER_ID = "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
ON_ZEROGPU = os.environ.get("SPACES_ZERO_GPU") == "true"

SYSTEM_PROMPT = """You are a prompt rewriter for a text-to-motion model that was trained on the HumanML3D dataset. Convert the user's input (in any language, any phrasing) into ONE short English caption that describes the motion only, in the exact style of HumanML3D captions.

HumanML3D style rules:
- English only.
- Present tense or "a person ..." phrasing.
- Describe motion only — strip greetings, polite framing, "please show me", emotional context.
- Use motion verbs (walk, run, jump, kick, bend, sit, stand, raise, lower, turn, step).
- 8 to 30 tokens.
- No quotes, no commentary, no preamble. Output only the caption."""

USER_TEMPLATE = """Examples of HumanML3D-style captions in the right style:
{examples}

Now rewrite this user input as ONE HumanML3D-style caption:
{user_input}"""


class Retriever:
    def __init__(self, data_dir: Path, device: str):
        self.model = SentenceTransformer(RETRIEVER_ID, device=device)
        self.captions: list[str] = json.loads((data_dir / "captions.json").read_text())
        self.embeddings: np.ndarray = np.load(data_dir / "embeddings.npy").astype(np.float32)

    def topk(self, query: str, k: int = 3) -> list[str]:
        q = self.model.encode([query], normalize_embeddings=True, convert_to_numpy=True)[0]
        sims = self.embeddings @ q
        idxs = np.argpartition(-sims, k)[:k]
        idxs = idxs[np.argsort(-sims[idxs])]
        return [self.captions[i] for i in idxs]


class Rewriter:
    def __init__(self, data_dir: Path, device: str | None = None, dtype: str | None = None):
        # On ZeroGPU, load the model on CPU at boot (GPU is only available inside @spaces.GPU calls).
        # We move it to CUDA on first inference inside the decorated method.
        if ON_ZEROGPU:
            device = "cpu"
            dtype = "bfloat16"
        elif device is None:
            device = "cuda" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu")
        if dtype is None:
            dtype = "bfloat16" if device != "cpu" else "float32"
        self.device = device
        self.model_id = MODEL_ID
        self._on_zerogpu = ON_ZEROGPU
        self._moved_to_cuda = False
        t0 = time.time()
        self.tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
        self.model = AutoModelForCausalLM.from_pretrained(
            MODEL_ID,
            dtype=getattr(torch, dtype),
            device_map=device,
        ).eval()
        self.load_latency = time.time() - t0
        self.retriever = Retriever(data_dir, device="cpu")  # MiniLM on CPU is fine

    def _generate_core(self, prompt_ids: torch.Tensor, max_new_tokens: int) -> torch.Tensor:
        with torch.no_grad():
            return self.model.generate(
                prompt_ids,
                max_new_tokens=max_new_tokens,
                do_sample=False,
                temperature=None,
                top_p=None,
            )

    def _generate_on_gpu(self, prompt_ids: torch.Tensor, max_new_tokens: int) -> torch.Tensor:
        # ZeroGPU: move model + inputs to CUDA inside this call, run, optionally leave on GPU.
        if not self._moved_to_cuda:
            self.model = self.model.to("cuda")
            self._moved_to_cuda = True
        return self._generate_core(prompt_ids.to("cuda"), max_new_tokens)

    def rewrite(self, user_input: str, max_new_tokens: int = 64) -> dict:
        t0 = time.time()
        examples = self.retriever.topk(user_input, k=3)
        examples_block = "\n".join(f"- {e}" for e in examples)
        user_message = USER_TEMPLATE.format(examples=examples_block, user_input=user_input)
        if "gemma-3" in self.model_id:
            messages = [{"role": "user", "content": [{"type": "text", "text": SYSTEM_PROMPT + "\n\n" + user_message}]}]
        else:
            messages = [
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": user_message},
            ]
        encoded = self.tokenizer.apply_chat_template(
            messages, add_generation_prompt=True, return_tensors="pt", return_dict=False
        )
        if hasattr(encoded, "shape"):
            prompt_ids = encoded
        else:
            prompt_ids = encoded["input_ids"]
        input_len = prompt_ids.shape[-1]
        if self._on_zerogpu:
            out = self._generate_on_gpu(prompt_ids, max_new_tokens)
        else:
            out = self._generate_core(prompt_ids.to(self.device), max_new_tokens)
        latency = time.time() - t0
        new_tokens = out[0, input_len:]
        decoded = self.tokenizer.decode(new_tokens, skip_special_tokens=True).strip()
        decoded = decoded.strip().strip('"').strip("'").lstrip("-").strip()
        decoded = decoded.split("\n")[0].strip()
        return {
            "rewritten": decoded,
            "examples": examples,
            "latency_s": latency,
            "input_tokens": int(input_len),
            "output_tokens": int(new_tokens.shape[0]),
        }