File size: 1,242 Bytes
0d9299d | 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 31 32 33 34 35 36 37 38 39 40 | 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:
label_scores = sorted(
[{"label": id2label[i], "score": float(row[i])} for i in range(len(row))],
key=lambda x: -x["score"],
)
results.append(label_scores)
return results if len(results) > 1 else results[0]
|