Delete jev_core.py

jev_core.py

@@ -1,187 +0,0 @@
-"""JEV-9B inference core, ported 1:1 from the model authors' documented `transformers` +
-`peft` path (https://huggingface.co/autotrust/JEV-9B -> "Other ways to run it").
-
-The Blocks-of-Experts architecture:
-  * backbone  = frozen Qwen3.5-9B text tower (bit-identical to the release)
-  * System 1  = LoRA r=16 (40.1M, shipped unmerged in `adapter/`) + a 24-slot fp32
-               decision head (`head.safetensors`) initialised from `lm_head` verbalizer rows
-  * System 2  = the original `lm_head`; used with the adapter off (pristine Qwen3.5-9B)
-
-System 1 answers a typed question (`noul` yes/no, `choice` over 2-16 options, `score` on a
-0-5 scale) in ONE prefill pass and returns a calibrated probability distribution. The
-decision prompt follows the authors' "bare-v1" template; the head is applied to the last
-token's post-norm hidden state, divided by the per-kind calibration temperature.
-"""
-
-import json
-import os
-
-import torch
-from huggingface_hub import hf_hub_download
-from peft import PeftModel
-from safetensors.torch import load_file
-from transformers import AutoModelForCausalLM, AutoTokenizer
-
-REPO = "autotrust/JEV-9B"
-LETTERS = "ABCDEFGHIJKLMNOP"
-KINDS = ("noul", "choice", "score")
-MAX_STATE_TOKENS = 1024  # the authors truncate long inputs to 1,024 tokens (head 60% / tail 40%)
-DEVICE = os.environ.get("JEV_DEVICE", "cuda")
-# Optional local-dir override (used only for offline verification of the wiring).
-_LOCAL = os.environ.get("JEV_MODEL_PATH")
-
-
-def _file(name: str) -> str:
-    """Resolve a model-repo file to a local path, falling back to the Hub."""
-    if _LOCAL:
-        p = os.path.join(_LOCAL, name)
-        if os.path.exists(p):
-            return p
-    return hf_hub_download(REPO, name)
-
-
-_SRC = _LOCAL or REPO
-
-# ---------------------------------------------------------------------------
-# Model loading. On ZeroGPU the final `_peft.to(DEVICE)` is intercepted at
-# startup: tensors are packed to disk and streamed into VRAM on the first
-# @spaces.GPU entry, so the load cost is never paid inside a request.
-#
-# Everything is built on CPU first, then moved once:
-#   * peft's default would be `torch_device=infer_device()` -> "cuda", making
-#     safetensors read the adapter straight onto the GPU. That bypasses the
-#     packer and crashes the GPU-less main process ("No CUDA GPUs are
-#     available"). `torch_device="cpu"` keeps the read on the CPU.
-#   * one `.to(DEVICE)` at the end moves every tensor (backbone, LoRA, head)
-#     through the standard machinery.
-# ---------------------------------------------------------------------------
-_tok = AutoTokenizer.from_pretrained(_SRC)
-_base = AutoModelForCausalLM.from_pretrained(_SRC, dtype=torch.bfloat16)
-
-# Keep the adapter UNMERGED so System 2 (generation) runs with the adapter
-# disabled -- `with _peft.disable_adapter():` -- i.e. pristine Qwen3.5-9B.
-_peft = PeftModel.from_pretrained(_base, _SRC, subfolder="adapter", torch_device="cpu")
-
-# System 1 head: fp32 linear, 4096 -> 24 slots.
-_HEAD = torch.nn.Linear(4096, 24, dtype=torch.float32)
-_head_sd = load_file(_file("head.safetensors"))
-with torch.no_grad():
-    _HEAD.weight.copy_(_head_sd["proj.weight"])
-    _HEAD.bias.copy_(_head_sd["proj.bias"])
-_HEAD.eval()
-
-_JUDGE_CFG = json.load(open(_file("judge_config.json")))
-_CALIB = json.load(open(_file("calibration.json")))["per_kind"]
-SLOT_RANGES = {k: tuple(v) for k, v in _JUDGE_CFG["slots"]["ranges"].items()}
-TEMPERATURES = {k: float(_CALIB[k]) for k in KINDS}
-
-# Module handles: PeftModel -> LoraModel -> Qwen3_5ForCausalLM -> Qwen3_5TextModel.
-_CAUSAL = _peft.base_model.model          # the Qwen3_5ForCausalLM with LoRA modules inside
-_BACKBONE = _CAUSAL.model                 # the text tower (LoRA active by default)
-
-# Single move of the whole tree (backbone + LoRA + head) to the compute device.
-_peft.to(DEVICE)
-_HEAD.to(DEVICE)
-
-# ---------------------------------------------------------------------------
-# Prompt template -- exactly the authors' `bare-v1` decision format.
-# ---------------------------------------------------------------------------
-
-
-def build_decision_prompt(kind: str, state: str, question: str, options) -> tuple:
-    """Return (prompt_text, option_labels) in the authors' `bare-v1` format."""
-    if kind == "noul":
-        options = ["false", "true"]
-    elif kind == "score":
-        options = [str(i) for i in range(6)]
-    options = list(options)
-    if kind == "choice":
-        lines = [f"{LETTERS[i]}) {o}" for i, o in enumerate(options)]
-    else:
-        lines = options
-    prompt = (
-        f"[kind] {kind}\n[state] {state}\n[question] {question}\n[options]\n"
-        + "\n".join(lines)
-        + "\n[decision]:"
-    )
-    return prompt, options
-
-
-def truncate_state(state: str) -> tuple:
-    """Head-60% / tail-40% truncation of the state to `MAX_STATE_TOKENS` tokens.
-
-    Returns (state_used, was_truncated, n_tokens).
-    """
-    ids = _tok(state, add_special_tokens=False)["input_ids"]
-    if len(ids) <= MAX_STATE_TOKENS:
-        return state, False, len(ids)
-    keep_head = int(MAX_STATE_TOKENS * 0.6)
-    keep_tail = MAX_STATE_TOKENS - keep_head
-    return _tok.decode(ids[:keep_head] + ids[-keep_tail:], skip_special_tokens=True), True, len(ids)
-
-
-def decide(kind: str, state: str, question: str, options=None) -> dict:
-    """Run ONE System 1 typed decision; return the calibrated probability distribution.
-
-    This mirrors the model card's reference `decide()` byte for byte: one prefill pass
-    through the backbone with the LoRA active, the 24-slot head applied to the last
-    token's hidden state, per-kind temperature, softmax over the active slots.
-    """
-    if kind not in KINDS:
-        raise ValueError(f"unknown kind {kind!r}")
-    if kind == "choice":
-        if not isinstance(options, (list, tuple)) or not 2 <= len(options) <= 16:
-            raise ValueError("`choice` needs between 2 and 16 options")
-        if len(set(options)) != len(options):
-            raise ValueError("`choice` options must be unique")
-    prompt, labels = build_decision_prompt(kind, state, question, options)
-    device = _BACKBONE.embed_tokens.weight.device   # resolve inside the GPU call
-    ids = _tok(prompt, return_tensors="pt", add_special_tokens=False).to(device)
-    with torch.no_grad():
-        h = _BACKBONE(**ids).last_hidden_state[0, -1].float()
-        z = _HEAD(h.to(_HEAD.weight.device))       # [24] fp32 logits
-    lo = SLOT_RANGES[kind][0]
-    z = z[lo : lo + len(labels)] / TEMPERATURES[kind]
-    p = torch.softmax(z, dim=0)
-    return {label: float(v) for label, v in zip(labels, p.tolist())}
-
-
-def n_state_tokens(state: str) -> int:
-    return len(_tok(state, add_special_tokens=False)["input_ids"])
-
-
-# ---------------------------------------------------------------------------
-# System 2 -- ordinary generation from the pristine base model (adapter off).
-# ---------------------------------------------------------------------------
-
-
-def build_escalation_prompt(state: str, question: str, options) -> str:
-    """The model card's confidence-gated escalation prompt (System 1 -> System 2)."""
-    opts = [str(o) for o in options]
-    if opts == ["false", "true"]:
-        tail = "Answer `true` or `false` on the last line."
-    else:
-        tail = "Options: " + "; ".join(opts) + "\nThink it through, then give exactly one option on the last line."
-    return f"{state}\n\nQuestion: {question}\n{tail}"
-
-
-def generate(state: str, question: str, options, max_new_tokens: int = 512,
-             enable_thinking: bool = True) -> tuple:
-    """System 2: generate a reasoned answer with the adapter OFF -- pristine Qwen3.5-9B.
-
-    Returns (text, completed): `completed` is False when generation hit the token
-    budget mid-sentence (the model is verbose in thinking mode), which the UI needs
-    to avoid presenting a cut-off trace as a finished answer.
-    """
-    prompt = build_escalation_prompt(state, question, options)
-    msgs = [{"role": "user", "content": prompt}]
-    enc = _tok.apply_chat_template(
-        msgs, add_generation_prompt=True, enable_thinking=enable_thinking,
-        return_tensors="pt", return_dict=True,
-    ).to(_BACKBONE.embed_tokens.weight.device)
-    n_prompt = enc["input_ids"].shape[1]
-    with torch.no_grad(), _peft.disable_adapter():
-        out = _CAUSAL.generate(**enc, max_new_tokens=int(max_new_tokens), do_sample=False)
-    text = _tok.decode(out[0, n_prompt:], skip_special_tokens=True)
-    completed = out.shape[1] - n_prompt < int(max_new_tokens)  # stopped before the cap
-    return text.strip(), completed