Update handler.py

handler.py

@@ -1,62 +1,38 @@
-from transformers import LongformerTokenizer, LongformerForSequenceClassification
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import torch
-from ts.torch_handler.base_handler import BaseHandler
-from safetensors.torch import load_file
-import os
-
-class LongformerRegressionHandler(BaseHandler):
-    def __init__(self):
-        super().__init__()
-        self.initialized = False
-
-    def initialize(self, ctx):
-        """Load model and tokenizer"""
-        properties = ctx.system_properties
-        model_dir = properties.get("model_dir")
-
-        # Load tokenizer and config
-        self.tokenizer = LongformerTokenizer.from_pretrained(model_dir)
-        self.model = LongformerForSequenceClassification.from_pretrained(model_dir)
-
-        # Load safetensors weights
-        weights_path = os.path.join(model_dir, "model.safetensors")
-        state_dict = load_file(weights_path)
-        self.model.load_state_dict(state_dict)
-
-        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        self.model.to(self.device)
-        self.model.eval()
-        self.initialized = True
-
-    def preprocess(self, requests):
-        """Convert raw text into model-ready inputs"""
-        inputs = []
-        for req in requests:
-            text = req.get("data") or req.get("body")
-            if isinstance(text, (bytes, bytearray)):
-                text = text.decode("utf-8")
-            tokens = self.tokenizer(
-                text,
-                padding="max_length",
-                truncation=True,
-                max_length=512,
-                return_tensors="pt"
-            )
-            tokens = {k: v.to(self.device) for k, v in tokens.items()}
-            inputs.append(tokens)
-        return inputs
-
-    def inference(self, inputs):
-        """Run forward pass and return clipped regression output"""
-        results = []
+
+# Load model and tokenizer once at module level
+MODEL_NAME = "open-paws/text_performance_prediction_longform"
+
+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
+
+    results = []
+    for text in texts:
+        # Tokenize
+        inputs = tokenizer(
+            text,
+            return_tensors="pt",
+            truncation=True,
+            padding="max_length",
+            max_length=512
+        )
+        inputs = {k: v.to(device) for k, v in inputs.items()}
+
         with torch.no_grad():
-            for inp in inputs:
-                output = self.model(**inp)
-                score = output.logits.squeeze().item()
-                clipped_score = min(max(score, 0.0), 1.0)
-                results.append(clipped_score)
-        return results
-
-    def postprocess(self, inference_output):
-        """Convert scores to response-friendly format"""
-        return [{"score": float(out)} for out in inference_output]
+            outputs = model(**inputs)
+            score = outputs.logits.squeeze().item()
+            clipped_score = min(max(score, 0.0), 1.0)
+
+        results.append({"score": clipped_score})
+
+    return results