File size: 1,160 Bytes
7d3a25a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 | 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]
|