| 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) |
| top_predictions = [] |
| input_ids = input_tokens["input_ids"] |
| for idx, prob in zip(top_k.indices[0], top_k.values[0]): |
| |
| 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() |
| |
| last_word = text.split()[-1] |
| |
| 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}] |
