File size: 1,292 Bytes
5277582 da5a2e2 5277582 33a4c63 5277582 56e1a5f 5277582 |
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 |
from typing import Dict, List, Any
from transformers import AutoModel, AutoTokenizer
class EndpointHandler():
def __init__(self, path=""):
self.tokenizer = AutoTokenizer.from_pretrained("Wellcome/WellcomeBertMesh")
self.model = AutoModel.from_pretrained("Wellcome/WellcomeBertMesh", trust_remote_code=True)
def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
"""
Args:
data (:obj:):
includes the input data and the parameters for the inference.
Return:
A :obj:`list`:. The object returned should be a list of one list like [[{"label": 0.9939950108528137}]] containing :
- "label": A string representing what the label/class is. There can be multiple labels.
- "score": A score between 0 and 1 describing how confident the model is for this label/class.
"""
text = data.pop("inputs", data)
inputs = self.tokenizer(text, padding="max_length")
preds = self.model(input_ids=[inputs["input_ids"]])
id2label = self.model.config.id2label
prediction = [
{"label": id2label[label_id], "score": p}
for label_id, p in enumerate(preds[0].tolist()) if p > 0.5
]
return prediction
|