Create reward_fn.py

reward_fn.py

@@ -0,0 +1,36 @@
+import torch
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+
+# Load reward model once
+rm_name = "OpenAssistant/reward-model-deberta-v3-large-v2"
+
+rm_tokenizer = AutoTokenizer.from_pretrained(rm_name)
+rm_model = AutoModelForSequenceClassification.from_pretrained(
+    rm_name,
+    torch_dtype=torch.float32
+).eval()
+
+# For Spaces: keep reward model on CPU to save VRAM
+device = torch.device("cpu")
+rm_model.to(device)
+
+
+def reward_fn(text: str) -> float:
+    """
+    Returns a scalar helpfulness reward from a trained reward model.
+    Higher = more helpful, clearer, more aligned response.
+    """
+    inputs = rm_tokenizer(
+        text,
+        return_tensors="pt",
+        truncation=True,
+        max_length=512
+    ).to(device)
+
+    with torch.no_grad():
+        scores = rm_model(**inputs).logits
+
+    # Reward = model's score for "helpful" class
+    reward = float(scores[0].item())
+
+    return reward