| from typing import Dict, List, Any |
| from transformers import AutoModelForMaskedLM, AutoTokenizer |
| import torch |
|
|
|
|
| class EndpointHandler(): |
| def __init__(self, path=""): |
| tokenizer = AutoTokenizer.from_pretrained(path) |
| model = AutoModelForMaskedLM.from_pretrained(path) |
| self.tokenizer = tokenizer |
| self.model = model |
|
|
| def __call__(self, data: Dict[str, Any]) -> List[Dict[Any, Any]]: |
| """ |
| data args: |
| inputs (:obj: `str`) |
| date (:obj: `str`) |
| Return: |
| A :obj:`list` | `dict`: will be serialized and returned |
| """ |
| |
| text = data.pop("text", data) |
| tokens = self.tokenizer(text, return_tensors='pt') |
| output = self.model(**tokens) |
| vec = torch.max( |
| torch.log( |
| 1 + torch.relu(output.logits) |
| ) * tokens.attention_mask.unsqueeze(-1), |
| dim=1)[0].squeeze() |
| cols = vec.nonzero().squeeze().cpu().tolist() |
| |
| weights = vec[cols].cpu().tolist() |
| |
| sparse_dict = dict(zip(map(str, cols), weights)) |
| return sparse_dict |
|
|
