src/model.py

"""gpt-oss-20b on ZeroGPU — the Producer's brain.

Loaded at module import (ZeroGPU pages weights in when the GPU attaches).
On machines without CUDA/transformers the import degrades gracefully and
src.lyrics falls back to the stub anchors.
"""

from . import GPU

MODEL_ID = "openai/gpt-oss-20b"

_model = None
_tokenizer = None
load_error: Exception | None = None

try:
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer, Mxfp4Config

    _tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
    # Explicit bf16 dequant: the MXFP4 kernels path breaks on ZeroGPU
    # ("Either a revision or a version must be specified"). MoE keeps
    # generation fast in bf16; H200 slice has the VRAM headroom.
    _model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID,
        torch_dtype=torch.bfloat16,
        device_map="cuda",
        quantization_config=Mxfp4Config(dequantize=True),
    )
except Exception as e:  # local dev without GPU, or download failure
    load_error = e


def _extract_final(text: str) -> str:
    """gpt-oss emits harmony channels; keep only the 'final' channel.
    If generation was cut off before reaching 'final' (all budget spent in
    'analysis'), return "" rather than leaking raw reasoning to the user."""
    marker = "<|channel|>final<|message|>"
    if marker in text:
        text = text.split(marker)[-1]
    elif "<|channel|>analysis" in text or "<|message|>" in text:
        return ""  # never reached final — incomplete, don't leak the thinking
    for tok in ("<|return|>", "<|end|>", "<|endoftext|>", "<|start|>"):
        text = text.replace(tok, "")
    return text.strip()


def generate_core(messages: list[dict], max_new_tokens: int = 1200,
                  temperature: float = 0.9) -> str:
    """Raw generation, NO @GPU decorator — safe to call from within another
    GPU context (e.g. the photo roast's wit stage)."""
    if _model is None:
        raise RuntimeError(f"gpt-oss-20b unavailable: {load_error}")
    inputs = _tokenizer.apply_chat_template(
        messages,
        add_generation_prompt=True,
        return_tensors="pt",
        return_dict=True,
        reasoning_effort="low",
    ).to(_model.device)
    out = _model.generate(
        **inputs,
        max_new_tokens=max_new_tokens,
        do_sample=True,
        temperature=temperature,
        top_p=0.95,
    )
    text = _tokenizer.decode(
        out[0][inputs["input_ids"].shape[-1] :], skip_special_tokens=False
    )
    return _extract_final(text)


@GPU(duration=120)
def generate_lyrics(messages: list[dict], max_new_tokens: int = 1200) -> str:
    return generate_core(messages, max_new_tokens)


@GPU(duration=300)
def run_producer(
    resume_text: str, job_description: str, genre: str, level: int, zone_desc: str,
    voice: str | None = None,
) -> dict:
    """The Producer's full agent loop in ONE GPU context (Best Agent):
    draft -> self-critique against the slider level -> revise. Returns the
    raw text of every stage; src.lyrics parses them and builds the trace.

    generate_core has no @GPU decorator, so the three calls share this one
    GPU attachment. Never raises: a stage that fails comes back as "" and
    the caller falls back to the draft (or the stub anchors)."""
    from .prompts import build_critique_messages, build_messages, build_revise_messages

    base = build_messages(resume_text, job_description, genre, level, zone_desc, voice)
    original_user = base[1]["content"]

    draft = generate_core(base, max_new_tokens=1200, temperature=0.9)

    critique = ""
    revised = ""
    if draft:  # only critique/revise something we actually have
        critique = generate_core(
            build_critique_messages(draft, level, zone_desc),
            max_new_tokens=300, temperature=0.6,
        )
        if critique:
            revised = generate_core(
                build_revise_messages(original_user, draft, critique),
                max_new_tokens=1200, temperature=0.85,
            )
    return {"draft": draft, "critique": critique, "revised": revised}