Upload app.py with huggingface_hub

app.py

@@ -109,13 +109,17 @@ def _load() -> dict:
     config = TransformerConfig.from_saved(saved)
     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).
+    # The published checkpoint declares dtype=bfloat16, and TransformerConfig builds
+    # every layer with that dtype. JAX on this GPU aborts on bf16->f16 conversions:
+    #   "Unsupported conversion from bf16 to f16 /
+    #    LLVM ERROR: Unsupported rounding mode for conversion."
+    # Casting the PARAMS alone does not help -- the model's own dtype must change too,
+    # so set it on the config BEFORE constructing the network.
+    # NEEDLE_CAST: "float32" (default) | "bfloat16" (original) | "none"
     _cast = os.environ.get("NEEDLE_CAST", "float32").lower()
     if _cast in ("float32", "float", "f32"):
+        # TransformerConfig is a plain @dataclass (no .replace()), so assign directly.
+        config.dtype = "float32"
         params = {k: (v.astype("float32") if hasattr(v, "astype") else v)
                   for k, v in raw_params.items()}
     else: