Update handler to return meaningful words

handler.py

@@ -0,0 +1,35 @@
+from typing import Dict, List, Any
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+import re
+
+class EndpointHandler:
+    def __init__(self, path=""):
+        self.model = AutoModelForCausalLM.from_pretrained("EleutherAI/pythia-410m")
+        self.tokenizer = AutoTokenizer.from_pretrained("EleutherAI/pythia-410m")
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.model.to(self.device)
+
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+        inputs = data.pop("inputs", data).strip()
+        input_tokens = self.tokenizer(inputs, return_tensors="pt").to(self.device)
+        with torch.no_grad():
+            outputs = self.model(**input_tokens)
+            logits = outputs.logits[:, -1, :]
+            probs = torch.softmax(logits, dim=-1)
+            top_k = torch.topk(probs, k=100)  # Check top 100 tokens
+        top_predictions = []
+        input_ids = input_tokens["input_ids"]
+        for idx, prob in zip(top_k.indices[0], top_k.values[0]):
+            # Append token to input and decode
+            next_token_id = idx.item()
+            test_ids = torch.cat([input_ids, torch.tensor([[next_token_id]], device=self.device)], dim=-1)
+            text = self.tokenizer.decode(test_ids[0], skip_special_tokens=True).strip()
+            # Extract last word
+            last_word = text.split()[-1]
+            # Filter for meaningful words (3+ letters, alphabetic)
+            if re.match(r"^[a-zA-Z]{3,}$", last_word) and last_word not in [p["text"] for p in top_predictions]:
+                top_predictions.append({"text": last_word, "probability": prob.item()})
+            if len(top_predictions) >= 5:
+                break
+        return top_predictions if top_predictions else [{"text": "No valid words found", "probability": 0.0}]