zerogpu_backend.py

"""In-Space ZeroGPU fallback for when Modal isn't used.

Only active when USE_ZEROGPU_FALLBACK=1. Requires the Space to run on ZeroGPU
hardware (HF PRO) and the heavy deps (torch, transformers) added to
requirements.txt. Kept out of the default path so the Space stays a light CPU
container. Loads MiniCPM4.1-8B locally and mimics llm.chat_json's JSON contract.

NOTE: without vLLM's guided decoding we can't *force* schema-valid JSON, so we
prompt firmly for JSON and defensively parse. Treat this as break-glass, not the
primary backend.
"""
from __future__ import annotations

import json

import config

_model = None
_tokenizer = None


def _load():
    global _model, _tokenizer
    if _model is not None:
        return
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer

    _tokenizer = AutoTokenizer.from_pretrained(config.MODEL_ID, trust_remote_code=True)
    _model = AutoModelForCausalLM.from_pretrained(
        config.MODEL_ID, trust_remote_code=True, torch_dtype=torch.bfloat16,
    ).eval().cuda()


def chat_json(messages: list[dict], schema: dict) -> dict:
    import spaces  # noqa: F401  (import guarded for non-ZeroGPU envs)

    @spaces.GPU(duration=120)
    def _run():
        _load()
        prompt = "\n\n".join(m["content"] for m in messages)
        prompt += "\n\nRespond with ONLY a JSON object matching: " + json.dumps(schema)
        out = _model.chat(_tokenizer, prompt, temperature=0.4, max_new_tokens=900)
        text = out[0] if isinstance(out, tuple) else out
        return _coerce(text)

    return _run()


def _coerce(raw: str) -> dict:
    s = (raw or "").strip().removeprefix("```json").removeprefix("```").removesuffix("```")
    start, end = s.find("{"), s.rfind("}")
    if start != -1 and end != -1:
        try:
            return json.loads(s[start:end + 1])
        except (json.JSONDecodeError, ValueError):
            pass
    return {}