Spaces:

pngwn
/

open-jev

Running on Zero

pngwn HF Staff
PR loader: drop changeset files from the file list and counts too; area options match the eval set
779479b verified about 1 month ago
26.7 kB
"""System One - parallel, guaranteed-output decisions.
The model's contribution is the numbers: a 4B scorer fine-tuned and
temperature-scaled to give calibrated distributions over the caller's options.
This Space shows the inference path that makes many of those cheap: (1) the
shared state is prefilled once, (2) every (question, option) branch is scored
from that cache in one batched forward, (3) softmax over exactly the caller's
options. It is compared with the instruct model on the same base generating a
JSON answer token by token.
"""
import copy
import json
import math
import os
import re
import time

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # noqa: E402 -- must precede torch: it monkey-patches cuda availability

import torch  # noqa: E402
from fastapi.responses import HTMLResponse  # noqa: E402
from gradio import Server as App  # noqa: E402

BASE_MODEL = "Qwen/Qwen3.5-4B-Base"
INSTRUCT_MODEL = "Qwen/Qwen3.5-4B"
DEFAULT_ADAPTER = "pngwn/system-one-qwen3.5-4b-scorer"
ADAPTER_REPO = os.environ.get("SYSTEM_ONE_ADAPTER", DEFAULT_ADAPTER)
TEMPERATURE = float(os.environ.get("SYSTEM_ONE_TEMPERATURE", "1.75"))
# Held-out ECE from the model card (0.135 raw head -> 0.044 at T=1.75). Only
# quoted for the adapter and temperature it was measured with.
HELD_OUT_ECE = 0.044 if (ADAPTER_REPO, TEMPERATURE) == (DEFAULT_ADAPTER, 1.75) else None
TRAIN_MAX_LEN = 384  # the scorer was trained with sequences truncated here
TRAIN_MAX_OPTIONS = 16
MAX_STATE_TOKENS = 16384
MAX_QUESTIONS = 24
MAX_OPTIONS = 32
# The training format never truncates the question or the option, so over-long
# ones are refused rather than silently changed.
MAX_QUESTION_TOKENS = 96
MAX_OPTION_TOKENS = 64
BRANCH_CHUNK = 96  # most branches in one batched forward
# Branching materializes the state's KV once per branch (~32 KB per token per
# branch on this model), so long states get fewer branches per forward.
BRANCH_TOKEN_BUDGET = 200_000
NAIVE_FORWARD_TOKENS = 48_000  # padded tokens per naive forward
NAIVE_MAX_TOKENS = 400_000  # beyond this the naive check is skipped, not run
GEN_PROMPT_MAX = 20_000  # beyond this the instruct lane is skipped, not truncated

# SYSTEM_ONE_MOCK=1 skips the models and fakes scores/generation so the
# frontend and the streaming plumbing can be exercised without a GPU.
MOCK = os.environ.get("SYSTEM_ONE_MOCK") == "1"
DEVICE = "cuda" if torch.cuda.is_available() and not MOCK else "cpu"
DTYPE = torch.bfloat16


# ------------------------------------------------------------------ rendering
# Identical to system_one.py: head and tail are tokenized separately and
# concatenated, and the head pools the last non-pad token, so the option stays
# at the end. That split is also exactly the cache boundary: the head is the
# shared prefix, each tail is one branch.

def head_text(state):
    return "State:\n" + state


def tail_text(question, option):
    return "\n\nQuestion:\n" + question + "\n\nOption:\n" + option


def pad_right(seqs, pad_id):
    L = max(len(s) for s in seqs)
    ids = torch.full((len(seqs), L), pad_id, dtype=torch.long)
    attn = torch.zeros((len(seqs), L), dtype=torch.long)
    for i, s in enumerate(seqs):
        ids[i, : len(s)] = torch.tensor(s, dtype=torch.long)
        attn[i, : len(s)] = 1
    return ids, attn


def softmax(xs):
    m = max(xs)
    es = [math.exp(x - m) for x in xs]
    s = sum(es)
    return [e / s for e in es]


def _sync():
    if DEVICE == "cuda":
        torch.cuda.synchronize()


# ---------------------------------------------------------------------- models
# Module scope: the ZeroGPU pack step captures module-scope .to("cuda") tensors
# at import time. The vision tower of the multimodal checkpoint is ignored by
# the text-only classes.

from transformers import AutoTokenizer  # noqa: E402

tok = AutoTokenizer.from_pretrained(BASE_MODEL)
if tok.pad_token is None:
    tok.pad_token = tok.eos_token

scorer = None
baseline = None
EOS_IDS = {tok.eos_token_id}

if not MOCK:
    from peft import PeftModel
    from transformers import Qwen3_5ForCausalLM, Qwen3_5TextForSequenceClassification

    def _load_scorer():
        model = Qwen3_5TextForSequenceClassification.from_pretrained(
            BASE_MODEL, num_labels=1, dtype=DTYPE, low_cpu_mem_usage=True, attn_implementation="sdpa")
        # pad/eos ids are None on this checkpoint; force them onto both configs.
        model.config.pad_token_id = tok.pad_token_id
        model.config.eos_token_id = tok.eos_token_id
        tc = model.config.get_text_config()
        tc.pad_token_id = tok.pad_token_id
        tc.eos_token_id = tok.eos_token_id
        # torch_device='cpu' is required on ZeroGPU: `import spaces` monkey-patches
        # torch.cuda.is_available() to True, so peft's infer_device() resolves 'cuda'
        # and safetensors dispatches a real CUDA op at module scope, where no GPU is
        # attached. Load on CPU; the following .to('cuda') is what the pack step captures.
        model = PeftModel.from_pretrained(model, ADAPTER_REPO, torch_device="cpu")
        # Merged so the cached path can call the text model and the score head
        # directly; the LoRA delta is folded into the base weights, `score` is
        # restored from modules_to_save.
        model = model.merge_and_unload()
        return model.to(DEVICE).eval()

    def _load_baseline():
        model = Qwen3_5ForCausalLM.from_pretrained(
            INSTRUCT_MODEL, dtype=DTYPE, low_cpu_mem_usage=True, attn_implementation="sdpa")
        return model.to(DEVICE).eval()

    scorer = _load_scorer()
    baseline = _load_baseline()
    gen_eos = baseline.generation_config.eos_token_id
    EOS_IDS |= set(gen_eos if isinstance(gen_eos, (list, tuple)) else [gen_eos])
    EOS_IDS.discard(None)


# ------------------------------------------------------------------------ plan

def normalize_questions(questions):
    out = []
    for q in (questions or [])[:MAX_QUESTIONS]:
        qtype = str(q.get("type", "choice")).lower()
        text = str(q.get("question", "")).strip()
        opts = [str(o).strip() for o in q.get("options", []) if str(o).strip()]
        if qtype not in ("noul", "choice", "score"):
            qtype = "choice"
        if qtype == "noul":
            opts = ["yes", "no"]
        opts = list(dict.fromkeys(opts))[:MAX_OPTIONS]
        if text and len(opts) >= 2:
            out.append({"id": f"q{len(out) + 1}", "type": qtype, "question": text, "options": opts})
    return out


def branch_chunk(P):
    return max(1, min(BRANCH_CHUNK, BRANCH_TOKEN_BUDGET // max(P, 1)))


def naive_rows(P, max_tail):
    return max(1, NAIVE_FORWARD_TOKENS // (P + max_tail))


def check_lengths(qs):
    """The first question or option that is too long to score, as a message, or None."""
    for q in qs:
        n = len(tok(q["question"], add_special_tokens=False)["input_ids"])
        if n > MAX_QUESTION_TOKENS:
            return f'{q["id"]}: the question is {n} tokens; the limit is {MAX_QUESTION_TOKENS}.'
        for i, o in enumerate(q["options"]):
            n = len(tok(o, add_special_tokens=False)["input_ids"])
            if n > MAX_OPTION_TOKENS:
                return f'{q["id"]}: option {i + 1} is {n} tokens; the limit is {MAX_OPTION_TOKENS}.'
    return None


def build_plan(state, qs):
    """Tokenize once and account for every token each path would process."""
    head_full = tok(head_text(state), add_special_tokens=False)["input_ids"]
    head = head_full[:MAX_STATE_TOKENS]
    branches = []
    for qi, q in enumerate(qs):
        for oi, o in enumerate(q["options"]):
            ids = tok(tail_text(q["question"], o), add_special_tokens=False)["input_ids"]
            branches.append({"q": qi, "o": oi, "ids": ids})
    P = len(head)
    suffix = sum(len(b["ids"]) for b in branches)
    max_tail = max(len(b["ids"]) for b in branches)
    cached = P + suffix
    naive = P * len(branches) + suffix
    info = {
        "state_tokens": P,
        "state_tokens_full": len(head_full),
        "truncated": len(head_full) > P,
        "n_questions": len(qs),
        "groups": [len(q["options"]) for q in qs],
        "n_branches": len(branches),
        "suffix_tokens": suffix,
        "cached_tokens": cached,
        "naive_tokens": naive,
        "ratio": round(naive / cached, 2),
        "naive_forwards": sum(math.ceil(len(q["options"]) / naive_rows(P, max_tail)) for q in qs),
        "cached_forwards": 1 + math.ceil(len(branches) / branch_chunk(P)),
        "branch_chunk": branch_chunk(P),
        "max_tail": max_tail,
        "beyond_train_len": P + max_tail > TRAIN_MAX_LEN,
        "train_max_len": TRAIN_MAX_LEN,
        "over_option_cap": [q["id"] for q in qs if len(q["options"]) > TRAIN_MAX_OPTIONS],
    }
    return head, branches, info


# --------------------------------------------------------------------- scorer

@torch.no_grad()
def score_cached(head, branches):
    """Prefill the state once, then score every branch from the cached state."""
    P = len(head)
    x = torch.tensor([head], dtype=torch.long, device=DEVICE)
    _sync()
    t0 = time.perf_counter()
    prefix = scorer.model(input_ids=x, use_cache=True).past_key_values
    _sync()
    prefill_ms = (time.perf_counter() - t0) * 1000.0

    scores = []
    step = branch_chunk(P)
    chunks = [branches[i: i + step] for i in range(0, len(branches), step)]
    t1 = time.perf_counter()
    for chunk in chunks:
        n = len(chunk)
        # Batch 1 -> n. reorder_cache index-selects the KV of the full-attention
        # layers and the conv/recurrent state of the linear-attention layers alike.
        cache = prefix if len(chunks) == 1 else copy.deepcopy(prefix)
        cache.reorder_cache(torch.zeros(n, dtype=torch.long, device=DEVICE))
        ids, sattn = pad_right([b["ids"] for b in chunk], tok.pad_token_id)
        attn = torch.cat([torch.ones((n, P), dtype=torch.long), sattn], dim=1)
        h = scorer.model(input_ids=ids.to(DEVICE), attention_mask=attn.to(DEVICE),
                         past_key_values=cache, use_cache=True).last_hidden_state
        last = (sattn.sum(1) - 1).to(DEVICE)
        pooled = h[torch.arange(n, device=DEVICE), last]
        scores += scorer.score(pooled).squeeze(-1).float().cpu().tolist()
        del cache, h
    _sync()
    branch_ms = (time.perf_counter() - t1) * 1000.0
    return scores, prefill_ms, branch_ms, 1 + len(chunks)


@torch.no_grad()
def score_naive(head, qs, branches):
    """What the published demo does: one forward per question, state re-encoded per option."""
    scores, forwards = [], 0
    rows = naive_rows(len(head), max(len(b["ids"]) for b in branches))
    _sync()
    t0 = time.perf_counter()
    for qi in range(len(qs)):
        seqs = [head + b["ids"] for b in branches if b["q"] == qi]
        for i in range(0, len(seqs), rows):  # one forward per question unless the state is long
            ids, attn = pad_right(seqs[i: i + rows], tok.pad_token_id)
            out = scorer(input_ids=ids.to(DEVICE), attention_mask=attn.to(DEVICE), use_cache=False)
            scores += out.logits.squeeze(-1).float().cpu().tolist()
            forwards += 1
    _sync()
    return scores, (time.perf_counter() - t0) * 1000.0, forwards


def distributions(qs, branches, scores):
    out = []
    for qi, q in enumerate(qs):
        logits = [s for b, s in zip(branches, scores) if b["q"] == qi]
        out.append(softmax([x / TEMPERATURE for x in logits]))
    return out


def scorer_questions(qs, probs):
    blocks = []
    for q, p in zip(qs, probs):
        i = max(range(len(p)), key=p.__getitem__)
        b = {"id": q["id"], "type": q["type"], "question": q["question"], "options": q["options"],
             "probs": [round(x, 4) for x in p], "chosen_index": i, "chosen": q["options"][i],
             "conf": round(p[i], 4)}
        if q["type"] == "score":
            b["expected"] = round(sum((k + 1) * x for k, x in enumerate(p)), 2)
        blocks.append(b)
    return blocks


# ------------------------------------------------------------------- instruct

SYSTEM_PROMPT = (
    "You are a decision function. Read the state, then answer every question. "
    "Reply with one JSON object and nothing else: each key is a question id and each "
    "value is exactly one of that question's allowed options, copied verbatim.")


def build_json_prompt(state, qs):
    lines = ["State:", state, "", "Questions:"]
    for q in qs:
        lines.append(f'{q["id"]}: {q["question"]}')
        lines.append("    allowed options: " + json.dumps(q["options"], ensure_ascii=False))
    lines += ["", "Reply format: {" + ", ".join(f'"{q["id"]}": "..."' for q in qs) + "}"]
    messages = [{"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": "\n".join(lines)}]
    # Thinking disabled: this arm should answer directly, and a reasoning
    # preamble would only add sequential tokens before the answer appears.
    try:
        text = tok.apply_chat_template(messages, tokenize=False,
                                       add_generation_prompt=True, enable_thinking=False)
    except TypeError:
        text = tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    return tok(text, add_special_tokens=False)["input_ids"]


@torch.no_grad()
def generate_stream(prompt_ids, max_new):
    """Greedy decode, one forward per token, yielding after every token."""
    x = torch.tensor([prompt_ids], dtype=torch.long, device=DEVICE)
    _sync()
    t0 = time.perf_counter()
    out = baseline(input_ids=x, use_cache=True, logits_to_keep=1)
    cache = out.past_key_values
    nxt = out.logits[0, -1].argmax()
    _sync()
    prefill_ms = (time.perf_counter() - t0) * 1000.0
    gen, stopped = [], "length"
    while len(gen) < max_new:
        tid = int(nxt)
        if tid in EOS_IDS:
            stopped = "eos"
            break
        gen.append(tid)
        yield {"ids": list(gen), "prefill_ms": prefill_ms, "stopped": None,
               "ms": (time.perf_counter() - t0) * 1000.0}
        out = baseline(input_ids=nxt.view(1, 1), past_key_values=cache, use_cache=True, logits_to_keep=1)
        nxt = out.logits[0, -1].argmax()
    yield {"ids": list(gen), "prefill_ms": prefill_ms, "stopped": stopped,
           "ms": (time.perf_counter() - t0) * 1000.0}


def _mock_generate_stream(qs, max_new):
    answers = {q["id"]: q["options"][0] for q in qs}
    if len(qs) > 1:  # exercise the frontend's "not verbatim" state
        answers[qs[1]["id"]] = qs[1]["options"][0].upper() + "!"
    ids = tok(json.dumps(answers), add_special_tokens=False)["input_ids"][:max_new]
    t0 = time.perf_counter()
    time.sleep(0.12)
    for k in range(1, len(ids) + 1):
        time.sleep(0.03)
        yield {"ids": ids[:k], "prefill_ms": 120.0, "stopped": "eos" if k == len(ids) else None,
               "ms": (time.perf_counter() - t0) * 1000.0}


def validate_json(text, qs):
    """Strict check of the generated text against the schema the caller asked for."""
    body = re.sub(r"^```(?:json)?|```$", "", text.strip(), flags=re.MULTILINE).strip()
    m = re.search(r"\{.*\}", body, flags=re.DOTALL)
    obj = None
    if m:
        try:
            obj = json.loads(m.group(0))
        except json.JSONDecodeError:
            obj = None
    if not isinstance(obj, dict):
        return {"parsed": False, "answers": [{"id": q["id"], "status": "unparsed", "value": None} for q in qs],
                "n_valid": 0, "extra_keys": []}
    answers = []
    for q in qs:
        v = obj.get(q["id"])
        if v is None:
            answers.append({"id": q["id"], "status": "missing", "value": None})
            continue
        v = v if isinstance(v, str) else json.dumps(v, ensure_ascii=False)
        if v in q["options"]:
            status = "valid"
        elif v.strip().lower() in [o.lower() for o in q["options"]]:
            status = "not_verbatim"
        else:
            status = "not_allowed"
        answers.append({"id": q["id"], "status": status, "value": v})
    ids = {q["id"] for q in qs}
    return {"parsed": True, "answers": answers,
            "n_valid": sum(a["status"] == "valid" for a in answers),
            "extra_keys": [k for k in obj if k not in ids]}


# ------------------------------------------------------------------------ run

_WARM = False


@torch.no_grad()
def _warmup():
    """Pay lazy CUDA init before anything is timed; every path once, tiny shapes."""
    head, tails = [1, 2, 3, 4, 5, 6], [{"q": 0, "o": i, "ids": [7, 8, 9 + i]} for i in range(2)]
    score_cached(head, tails)
    score_naive(head, [None], tails)
    for i, _ in enumerate(generate_stream(head, 3)):
        pass


@torch.no_grad()
def _touch():
    """One untimed forward per model, so the first timed op does not absorb GPU wake-up."""
    x = torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8]], dtype=torch.long, device=DEVICE)
    scorer.model(input_ids=x)
    baseline(input_ids=x, logits_to_keep=1)
    _sync()


def _stream_gen(snap, qs, prompt_ids):
    max_new = min(768, 32 + 28 * len(qs))
    stream = _mock_generate_stream(qs, max_new) if MOCK else generate_stream(prompt_ids, max_new)
    last_emit = 0.0
    for part in stream:
        final = part["stopped"] is not None
        now = time.perf_counter()
        if not final and now - last_emit < 0.05:
            continue
        last_emit = now
        text = tok.decode(part["ids"], skip_special_tokens=True)
        n = len(part["ids"])
        snap["gen"] = {"text": text, "n_tokens": n, "prompt_tokens": len(prompt_ids),
                       "prefill_ms": round(part["prefill_ms"], 1), "ms": round(part["ms"], 1),
                       # the prefill forward emits token 1; each further token is one more
                       # forward, plus the one that emitted the end-of-sequence token
                       "forwards": n + (1 if part["stopped"] == "eos" else 0) if final else max(n, 1),
                       "tokens": len(prompt_ids) + n, "stopped": part["stopped"],
                       "validation": validate_json(text, qs) if final else None}
        yield copy.deepcopy(snap)


def _run(state, questions, compare, verify):
    global _WARM
    qs = normalize_questions(questions)
    if not qs:
        yield {"error": "Add at least one question with a question text and two or more options."}
        return
    too_long = check_lengths(qs)
    if too_long:
        yield {"error": too_long}
        return
    state = str(state or "")
    head, branches, info = build_plan(state, qs)
    snap = {"plan": info, "scorer": None, "naive": None, "gen": None, "done": False,
            "mock": MOCK, "cold_start": not _WARM,
            "models": {"scorer": ADAPTER_REPO, "instruct": INSTRUCT_MODEL, "temperature": TEMPERATURE,
                       "ece": HELD_OUT_ECE}}
    yield copy.deepcopy(snap)

    if not MOCK:
        _touch() if _WARM else _warmup()
    _WARM = True

    # 1. cached + parallel
    if MOCK:
        time.sleep(0.15)
        scores = [math.sin(7.3 * (i + 1)) * 3 for i in range(len(branches))]
        prefill_ms, branch_ms, forwards = 61.0, 88.0, info["cached_forwards"]
    else:
        scores, prefill_ms, branch_ms, forwards = score_cached(head, branches)
    probs = distributions(qs, branches, scores)
    snap["scorer"] = {"questions": scorer_questions(qs, probs), "prefill_ms": round(prefill_ms, 1),
                      "branch_ms": round(branch_ms, 1), "total_ms": round(prefill_ms + branch_ms, 1),
                      "forwards": forwards, "tokens": info["cached_tokens"]}
    yield copy.deepcopy(snap)

    # 2. the same scorer, run naively, as a live equivalence check
    if verify and info["naive_tokens"] > NAIVE_MAX_TOKENS:
        snap["naive"] = {"skipped": f'the naive path would push {info["naive_tokens"]:,} tokens through the model '
                                    f'(limit {NAIVE_MAX_TOKENS:,}), so the check was not run',
                         "forwards": info["naive_forwards"], "tokens": info["naive_tokens"]}
        yield copy.deepcopy(snap)
    elif verify:
        if MOCK:
            time.sleep(0.05 * len(qs))
            nscores, naive_ms, nforwards = [s + 0.001 for s in scores], 110.0 * len(qs), info["naive_forwards"]
        else:
            nscores, naive_ms, nforwards = score_naive(head, qs, branches)
        nprobs = distributions(qs, branches, nscores)
        dp = max(abs(a - b) for pa, pb in zip(probs, nprobs) for a, b in zip(pa, pb))
        agree = sum(max(range(len(a)), key=a.__getitem__) == max(range(len(b)), key=b.__getitem__)
                    for a, b in zip(probs, nprobs))
        snap["naive"] = {"total_ms": round(naive_ms, 1), "forwards": nforwards,
                         "tokens": info["naive_tokens"], "max_abs_dp": round(dp, 5),
                         "top1_agree": agree, "n": len(qs)}
        yield copy.deepcopy(snap)

    # 3. the instruct model on the same base, generating the same decisions as JSON
    if compare:
        if info["truncated"]:
            state = tok.decode(head)[len(head_text("")):]
        prompt_ids = build_json_prompt(state, qs)
        if len(prompt_ids) > GEN_PROMPT_MAX:
            snap["gen"] = {"skipped": f"the prompt is {len(prompt_ids):,} tokens (limit {GEN_PROMPT_MAX:,}); "
                                      "it is not truncated, because that would cut the instructions"}
            yield copy.deepcopy(snap)
        else:
            yield from _stream_gen(snap, qs, prompt_ids)

    snap["done"] = True
    yield copy.deepcopy(snap)


app = App()


# Decorate the function Gradio binds, in this order: the startup scan walks
# Gradio's registered handlers looking for a @spaces.GPU function, and
# @spaces.GPU outside @app.api(...) raises 'No @spaces.GPU function detected
# during startup'. The return annotation is the type of each yielded snapshot;
# gr.api reads it to build the output schema.
def _gpu_seconds(state: str, questions: list, compare: bool = True, verify: bool = True) -> int:
    """Short states finish in a few seconds; only long ones ask ZeroGPU for a long slot."""
    return 60 if len(str(state or "")) < 12_000 else 150


@spaces.GPU(duration=_gpu_seconds)
def run(state: str, questions: list, compare: bool = True, verify: bool = True) -> dict:
    try:
        yield from _run(state, questions, compare, verify)
    except Exception as e:
        import traceback

        traceback.print_exc()
        yield {"error": repr(e)}


app.api(name="run")(run)


# ---------------------------------------------------------------- pull requests
# A pull request is a natural long state: one diff, many typed questions about it.

PR_URL = re.compile(r"github\.com/([\w.-]+)/([\w.-]+)/pull/(\d+)")
LOCKFILES = re.compile(r"(^|/)(package-lock\.json|pnpm-lock\.yaml|yarn\.lock|poetry\.lock|uv\.lock|Cargo\.lock|go\.sum)$")
# Changesets state the change type and the version bump outright, which are two of the questions.
# Even listing one says "this needs a bump", so they are dropped from the file list and the counts.
CHANGESETS = re.compile(r"^\.changeset/.+\.md$")
STATUS = {"added": "A", "removed": "D", "modified": "M", "renamed": "R"}


def render_pr(repo, pr, files):
    """Summary first, diff last: the state is truncated from the end. Labels and changesets are left out
    on purpose: they are the answers to some of the questions."""
    hidden = [f for f in files if CHANGESETS.search(f["filename"])]
    files = [f for f in files if not CHANGESETS.search(f["filename"])]
    additions = pr["additions"] - sum(f["additions"] for f in hidden)
    deletions = pr["deletions"] - sum(f["deletions"] for f in hidden)
    body = re.sub(r"<!--.*?-->", "", pr.get("body") or "", flags=re.DOTALL).strip() or "(none)"
    lines = [f"Repository: {repo}", f'Pull request #{pr["number"]}: {pr["title"]}',
             f'Base branch: {pr["base"]["ref"]}', "", "Description:", body, "",
             f"Changed files: {len(files)} (+{additions} -{deletions})"]
    lines += [f'  {STATUS.get(f["status"], "M")} {f["filename"]} (+{f["additions"]} -{f["deletions"]})' for f in files]
    lines += ["", "Diff:"]
    for f in files:
        lines.append(f'--- {f["filename"]}')
        if LOCKFILES.search(f["filename"]):
            lines.append("(lockfile diff omitted)")
        else:
            lines.append(f.get("patch") or "(binary or too large to show)")
    return "\n".join(lines)


def fetch_pr(url):
    import httpx

    m = PR_URL.search(url or "")
    if not m:
        raise ValueError("Expected a URL like https://github.com/owner/repo/pull/123")
    owner, name, number = m.groups()  # only these three parts of the input reach the request
    api = f"https://api.github.com/repos/{owner}/{name}/pulls/{number}"
    headers = {"Accept": "application/vnd.github+json", "User-Agent": "open-jev-space",
               "X-GitHub-Api-Version": "2022-11-28"}
    if os.environ.get("GITHUB_TOKEN"):
        headers["Authorization"] = "Bearer " + os.environ["GITHUB_TOKEN"]
    with httpx.Client(timeout=20, headers=headers) as c:
        pr, files = c.get(api), c.get(api + "/files", params={"per_page": 100})
    for r in (pr, files):
        if r.status_code in (403, 429):
            raise RuntimeError("GitHub rate limit reached for this Space. Paste the PR text into the state instead.")
        if r.status_code == 404:
            raise RuntimeError("That pull request was not found (private repositories cannot be loaded).")
        r.raise_for_status()
    return f"{owner}/{name}", pr.json(), files.json()


def load_pr(url: str) -> dict:
    try:
        repo, pr, files = fetch_pr(url)
    except Exception as e:
        return {"error": str(e)}
    state = render_pr(repo, pr, files)
    n = len(tok(head_text(state), add_special_tokens=False)["input_ids"])
    return {"state": state, "title": pr["title"], "repo": repo, "number": pr["number"], "files": len(files),
            "additions": pr["additions"], "deletions": pr["deletions"], "state_tokens": n,
            "max_state_tokens": MAX_STATE_TOKENS}


app.api(name="load_pr")(load_pr)


@app.get("/pr_preset.json")
async def pr_preset():
    from fastapi.responses import FileResponse

    return FileResponse(os.path.join(os.path.dirname(os.path.abspath(__file__)), "pr_preset.json"))


_INDEX_HTML = None


@app.get("/", response_class=HTMLResponse)
async def homepage():
    global _INDEX_HTML
    if _INDEX_HTML is None or MOCK:
        here = os.path.dirname(os.path.abspath(__file__))
        with open(os.path.join(here, "index.html")) as f:
            _INDEX_HTML = f.read()
    return _INDEX_HTML


if __name__ == "__main__":
    app.launch()