|
import torch |
|
|
|
from typing import Any, Dict, Union |
|
from transformers import AutoModelForSequenceClassification, AutoTokenizer |
|
|
|
|
|
class EndpointHandler: |
|
def __init__(self, path=""): |
|
|
|
self.tokenizer = AutoTokenizer.from_pretrained(path) |
|
self.model = AutoModelForSequenceClassification.from_pretrained( |
|
path, device_map="auto", trust_remote_code=True |
|
) |
|
self.device = "cuda" if torch.cuda.is_available() else "cpu" |
|
|
|
def __call__(self, data: Dict[str, Any]) -> Dict[str, Union[str, float]]: |
|
|
|
inputs = data.pop("inputs", data) |
|
|
|
|
|
inputs = self.tokenizer(inputs, return_tensors="pt").to(self.device) |
|
|
|
|
|
logits = self.model(**inputs)[0] |
|
|
|
|
|
predicted_class_id = int(torch.argmax(logits, dim=-1)) |
|
predicted_score = float(logits[0, predicted_class_id]) |
|
predicted_label = str(self.model.config.id2label[predicted_class_id]) |
|
|
|
return {'label': predicted_label, 'score': predicted_score} |
|
|
