| from typing import Any, Dict, List |
| from transformers import AutoModelForSequenceClassification, AutoTokenizer |
| import torch |
|
|
| class EndpointHandler: |
| def __init__(self, path=""): |
| self.tokenizer = AutoTokenizer.from_pretrained(path) |
| self.model = AutoModelForSequenceClassification.from_pretrained(path) |
| self.model.eval() |
|
|
| def __call__(self, data: Dict[str, Any]) -> List[Dict]: |
| inputs_text = data.pop("inputs", data) |
| if isinstance(inputs_text, str): |
| inputs_text = [inputs_text] |
| encoded = self.tokenizer( |
| inputs_text, return_tensors="pt", padding=True, truncation=True, |
| return_token_type_ids=False, |
| ) |
| with torch.no_grad(): |
| logits = self.model(**encoded).logits |
| scores = torch.softmax(logits, dim=-1) |
| id2label = self.model.config.id2label |
| results = [] |
| for row in scores: |
| results.append(sorted( |
| [{"label": id2label[i], "score": float(row[i])} for i in range(len(row))], |
| key=lambda x: -x["score"], |
| )) |
| return results if len(results) > 1 else results[0] |
|
|
