Sync Capability Decision Lab from GitHub

app.py

@@ -22,7 +22,6 @@ import torch
 import laya
 from kev.api import SystemOneRequest, output_tokens, to_answers, to_record
 from kev.checkpoint import Checkpoint, LoadOptions
-from transformers import AutoModelForMultimodalLM, AutoProcessor
 
 NONE_ID = "__none__"
 INFER_MAX_STATE = 8192
@@ -32,7 +31,6 @@ MODEL_IDS = {
     "kev-0.8b": "jaredpalmer/kev-0.8b",
     "kev-4b": "jaredpalmer/kev-4b",
     "laya": "convaiinnovations/laya-typed-decisions",
-    "gemma-4-e4b": "google/gemma-4-E4B",
 }
 
 
@@ -54,12 +52,6 @@ def _load_kev(repo_id: str):
 KEV_08_CK, KEV_08_TOK, KEV_08 = _load_kev(MODEL_IDS["kev-0.8b"])
 KEV_4_CK, KEV_4_TOK, KEV_4 = _load_kev(MODEL_IDS["kev-4b"])
 LAYA = laya.load(MODEL_IDS["laya"], device="cuda")
-GEMMA4_PROCESSOR = AutoProcessor.from_pretrained(MODEL_IDS["gemma-4-e4b"])
-GEMMA4 = AutoModelForMultimodalLM.from_pretrained(
-    MODEL_IDS["gemma-4-e4b"],
-    torch_dtype=torch.bfloat16,
-).to("cuda").eval()
-
 KEV_MODELS = {
     "kev-0.8b": (KEV_08_CK, KEV_08_TOK, KEV_08),
     "kev-4b": (KEV_4_CK, KEV_4_TOK, KEV_4),
@@ -212,97 +204,6 @@ def _laya_predict(
     }
 
 
-def _extract_json_object(text: str) -> dict[str, Any]:
-    stripped = (text or "").strip()
-    try:
-        value = json.loads(stripped)
-        if isinstance(value, dict):
-            return value
-    except Exception:
-        pass
-    start = stripped.find("{")
-    end = stripped.rfind("}")
-    if start >= 0 and end > start:
-        return json.loads(stripped[start:end + 1])
-    raise ValueError(f"Gemma 4 response did not contain a JSON object: {stripped[:500]}")
-
-
-def _gemma4_predict(case: dict[str, Any], reverse: bool = False):
-    candidates = list(case["candidates"])
-    if reverse:
-        candidates.reverse()
-    system = (
-        "You are a strict capability resolver. "
-        "Choose one supplied candidate only if using it exactly as described would fully satisfy the user's original intent. "
-        "Semantic relatedness is not enough and a nearby substitute does not count. "
-        "If none of the candidates fully satisfies the intent, choose null. "
-        "Do not consider permission or authorization. "
-        "Return JSON only with selected and reason. "
-        "selected must be one supplied candidate id or null. "
-        "Keep reason to one short sentence."
-    )
-    user_payload = {
-        "original_intent": case["intent"],
-        "requested_signature": case["signature"],
-        "candidates": [
-            {
-                "id": candidate["id"],
-                "name": candidate.get("name"),
-                "summary": candidate.get("summary"),
-                "provides": candidate.get("provides") or [],
-                "matching_operation_signatures": candidate.get("matching_operation_signatures") or [],
-            }
-            for candidate in candidates
-        ],
-        "question": "Which candidate, if any, fully satisfies the original intent?",
-    }
-    messages = [
-        {"role": "system", "content": [{"type": "text", "text": system}]},
-        {"role": "user", "content": [{"type": "text", "text": json.dumps(user_payload)}]},
-    ]
-    inputs = GEMMA4_PROCESSOR.apply_chat_template(
-        messages,
-        add_generation_prompt=True,
-        tokenize=True,
-        return_dict=True,
-        return_tensors="pt",
-    )
-    inputs = {k: v.to("cuda") if hasattr(v, "to") else v for k, v in inputs.items()}
-    torch.cuda.synchronize()
-    started = time.perf_counter()
-    with torch.inference_mode():
-        generated = GEMMA4.generate(
-            **inputs,
-            max_new_tokens=96,
-            do_sample=False,
-        )
-    torch.cuda.synchronize()
-    elapsed_ms = (time.perf_counter() - started) * 1000
-    input_len = inputs["input_ids"].shape[-1]
-    text = GEMMA4_PROCESSOR.decode(generated[0][input_len:], skip_special_tokens=True)
-    parsed = _extract_json_object(text)
-    selected = parsed.get("selected")
-    if selected is None or str(selected).strip().upper() == "NONE":
-        selected = None
-    else:
-        selected = str(selected).strip()
-        allowed = {candidate["id"] for candidate in candidates}
-        if selected not in allowed:
-            raise ValueError(f"Gemma 4 selected unknown candidate: {selected}")
-    return {
-        "order": [candidate["id"] for candidate in candidates],
-        "answer": {
-            "selected": selected,
-            "confidence": None,
-            "probabilities": {},
-            "probability_margin": None,
-        },
-        "elapsed_ms": round(elapsed_ms, 3),
-        "raw_text": text,
-        "reason": str(parsed.get("reason") or "").strip(),
-    }
-
-
 def _score_case(case: dict[str, Any], selected: str | None) -> bool:
     expected = set(case.get("expected") or [])
     return selected in expected if expected else selected is None
@@ -367,11 +268,6 @@ def _run_model(
                     answer = raw.pop("answer")
                     answer["correct"] = _score_case(case, answer["selected"])
                     runs[label] = {"order": order, **answer, **raw}
-                elif model_name == "gemma-4-e4b":
-                    raw = _gemma4_predict(case, reverse=reverse)
-                    answer = raw.pop("answer")
-                    answer["correct"] = _score_case(case, answer["selected"])
-                    runs[label] = {**answer, **raw}
                 else:
                     raw = _laya_predict(state, questions, stream=stream)
                     answer = raw.pop("answer")
@@ -440,8 +336,6 @@ def _run_suite_impl(payload_json: str, parallel: bool = True, batch_size: int =
     if parallel and len(requested) > 1:
         # Give each resident model its own CUDA stream so independent kernels can
         # overlap on the same ZeroGPU allocation where the GPU scheduler permits.
-        if "gemma-4-e4b" in requested:
-            raise ValueError("gemma-4-e4b is benchmarked in serial mode; do not mix it into the CUDA-stream stress test")
         streams = {name: torch.cuda.Stream() for name in requested}
         with ThreadPoolExecutor(max_workers=len(requested)) as pool:
             futures = {
@@ -482,7 +376,7 @@ def run_suite_api(payload_json: str, parallel: bool = True, batch_size: int = 1)
 with gr.Blocks(title="Capability Decision Lab") as demo:
     gr.Markdown(
         "# Capability Decision Lab\n"
-        "ZeroGPU benchmark service for Kev-0.8B, Kev-4B, Laya typed-decisions, and Gemma 4 E4B."
+        "ZeroGPU benchmark service for Kev-0.8B, Kev-4B, and Laya typed-decisions."
     )
     payload_box = gr.Textbox(
         label="Benchmark payload (JSON)",