Upload app.py with huggingface_hub

app.py

@@ -107,7 +107,20 @@ def _load() -> dict:
     ckpt = read_checkpoint(CKPT)
     saved = ckpt["config"] if isinstance(ckpt["config"], dict) else dict(vars(ckpt["config"]))
     config = TransformerConfig.from_saved(saved)
-    params = {k: v for k, v in ckpt["params"].items() if not k.startswith("mtp_")}
+    raw_params = {k: v for k, v in ckpt["params"].items() if not k.startswith("mtp_")}
+
+    # The published checkpoint is bf16. JAX on this GPU aborts with
+    #   "Unsupported conversion from bf16 to f16 / LLVM ERROR: Unsupported
+    #    rounding mode for conversion."
+    # so cast everything to float32 up front. NEEDLE_CAST controls this
+    # ("float32" default, "bfloat16" to keep the original, "none" to skip).
+    _cast = os.environ.get("NEEDLE_CAST", "float32").lower()
+    if _cast in ("float32", "float", "f32"):
+        params = {k: (v.astype("float32") if hasattr(v, "astype") else v)
+                  for k, v in raw_params.items()}
+    else:
+        params = raw_params
+
     model = SimpleAttentionNetwork(config=config)
     tokenizer = get_tokenizer()
 
@@ -117,7 +130,7 @@ def _load() -> dict:
 
     _CACHE.update({"jax": jax, "jnp": jnp, "params": params, "model": model,
                    "tokenizer": tokenizer, "decode_fn": decode_fn, "config": config,
-                   "backend": jax.default_backend(),
+                   "backend": jax.default_backend(), "cast": _cast,
                    "devices": [str(d) for d in jax.devices()]})
     return _CACHE
 
@@ -229,7 +242,7 @@ def _do_score(state: str, instructions: str, levels: list[str]) -> dict:
             "confidence": max(dist.values()) if dist else 0.0}
 
 
-@spaces.GPU(duration=60)
+@spaces.GPU(duration=120)
 def answer_systemone(state: str, questions: dict) -> dict:
     """Run a Jev-shaped request. This is the GPU-decorated entry point."""
     with _LOCK:
@@ -250,7 +263,7 @@ def answer_systemone(state: str, questions: dict) -> dict:
             "usage": {"input_tokens": None, "output_tokens": None}}
 
 
-@spaces.GPU(duration=60)
+@spaces.GPU(duration=120)
 def run_method_comparison() -> dict:
     """The scoring-method comparison on three unambiguous cases.
 
@@ -309,7 +322,7 @@ def health() -> str:
             f"JAX_PLATFORMS={os.environ.get('JAX_PLATFORMS')}")
 
 
-@spaces.GPU(duration=60)
+@spaces.GPU(duration=120)
 def gpu_probe() -> dict:
     """DECISIVE TEST: does JAX initialise on CUDA *inside* a @spaces.GPU function?
 
@@ -337,7 +350,7 @@ def gpu_probe() -> dict:
     return out
 
 
-@spaces.GPU(duration=60)
+@spaces.GPU(duration=120)
 def run_selftest() -> dict:
     """Known-answer probes, comparing positive vs negative on matched pairs."""
     with _LOCK: