Create handler.py

handler.py

@@ -0,0 +1,62 @@
+from transformers import LongformerTokenizer, LongformerForSequenceClassification
+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 = []
+        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]