Upload app.py with huggingface_hub

app.py

@@ -41,6 +41,7 @@ from transformers import AutoConfig, AutoModel, AutoProcessor  # noqa: E402
 # --------------------------------------------------------------------------- #
 
 MODEL_ID = "Alibaba-DAMO-Academy/RynnValue-4B"
+DTYPE = torch.bfloat16
 
 _config = AutoConfig.from_pretrained(MODEL_ID, trust_remote_code=True)
 # The exported checkpoint doesn't persist the attention implementation, so force
@@ -51,9 +52,12 @@ model = AutoModel.from_pretrained(
     MODEL_ID,
     config=_config,
     trust_remote_code=True,
-    torch_dtype=torch.bfloat16,
+    torch_dtype=DTYPE,
 )
-model = model.eval().to("cuda")
+# The value heads are built in fp32 regardless of `torch_dtype`; the reference
+# CLI casts the whole model with `.to(device=..., dtype=...)`, so do the same or
+# the head matmuls hit "mat1 and mat2 must have the same dtype".
+model = model.eval().to(device="cuda", dtype=DTYPE)
 
 processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
 tokenizer = processor.tokenizer