Create handler.py

handler.py

@@ -0,0 +1,56 @@
+from typing import Any, Dict, List
+import torch
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+
+
+class EndpointHandler:
+    def __init__(self, path: str = ""):
+        self.tokenizer = AutoTokenizer.from_pretrained(path)
+        self.model = AutoModelForSeq2SeqLM.from_pretrained(path)
+
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.model.to(self.device)
+        self.model.eval()
+
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        text = data.get("inputs")
+        if text is None:
+            return {"error": "Missing required field: inputs"}
+
+        parameters = data.get("parameters", {})
+
+        encoded = self.tokenizer(
+            text,
+            return_tensors="pt",
+            truncation=True,
+            max_length=int(parameters.get("max_input_length", 1024)),
+        )
+        encoded = {k: v.to(self.device) for k, v in encoded.items()}
+
+        with torch.inference_mode():
+            output_ids = self.model.generate(
+                **encoded,
+                max_new_tokens=int(parameters.get("max_new_tokens", 48)),
+                num_beams=int(parameters.get("num_beams", 4)),
+                do_sample=bool(parameters.get("do_sample", False)),
+                temperature=float(parameters.get("temperature", 1.0)),
+                no_repeat_ngram_size=int(parameters.get("no_repeat_ngram_size", 3)),
+                early_stopping=True,
+            )
+
+        raw_text = self.tokenizer.decode(output_ids[0], skip_special_tokens=True).strip()
+
+        keywords = [x.strip() for x in raw_text.split(";") if x.strip()]
+
+        seen = set()
+        deduped: List[str] = []
+        for kw in keywords:
+            k = kw.lower()
+            if k not in seen:
+                seen.add(k)
+                deduped.append(kw)
+
+        return {
+            "generated_text": raw_text,
+            "keywords": deduped,
+        }