Update handler.py

handler.py

@@ -1,38 +1,43 @@
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import torch
 
-# Load model and tokenizer once at module level
-MODEL_NAME = "open-paws/text_performance_prediction_longform"
+# Load once when the endpoint starts
+model_name = "open-paws/text_performance_prediction_longform"
 
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForSequenceClassification.from_pretrained(model_name)
 model.eval()
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model.to(device)
 
-def predict(texts):
-    """Hugging Face expects a `predict()` function for custom handlers"""
-    if isinstance(texts, str):
-        texts = [texts]  # Ensure batch input
+def predict(inputs):
+    """
+    Hugging Face Inference Endpoints will call this function.
+    `inputs` can be a single string or a list of strings.
+    """
+    if isinstance(inputs, str):
+        inputs = [inputs]
 
     results = []
-    for text in texts:
-        # Tokenize
-        inputs = tokenizer(
+    for text in inputs:
+        encoded = tokenizer(
             text,
             return_tensors="pt",
             truncation=True,
             padding="max_length",
-            max_length=512
+            max_length=4096,
         )
-        inputs = {k: v.to(device) for k, v in inputs.items()}
+        encoded = {k: v.to(device) for k, v in encoded.items()}
 
         with torch.no_grad():
-            outputs = model(**inputs)
-            score = outputs.logits.squeeze().item()
-            clipped_score = min(max(score, 0.0), 1.0)
+            outputs = model(**encoded)
 
-        results.append({"score": clipped_score})
+        raw_score = outputs.logits.squeeze().item()
+        clipped_score = min(max(raw_score, 0.0), 1.0)
+
+        results.append({
+            "score": round(clipped_score, 4),
+        })
 
     return results