File size: 1,870 Bytes
1377356 3b7c54f 1377356 3b7c54f 1377356 3b7c54f 1377356 3b7c54f | 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 41 42 43 44 45 46 47 48 49 50 51 52 | from typing import Dict, List, Any
from transformers import AutoTokenizer, AutoModelForMaskedLM
import torch
class EndpointHandler:
def __init__(self, path: str = ""):
# Device to use for inference. "cuda" if CUDA is available, else "cpu"
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Load tokenizer and model
self.tokenizer = AutoTokenizer.from_pretrained(path)
self.model = AutoModelForMaskedLM.from_pretrained(path).to(self.device)
def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
"""
data args:
inputs (:obj: `str`)
Return:
A :obj:`list` | `dict`: will be serialized and returned
"""
# Get inputs
inputs = data.pop("inputs", data)
if isinstance(inputs, str):
inputs = [inputs]
# Run the model to get logits
inputs = self.tokenizer(
inputs,
return_tensors="pt",
padding=True,
truncation=True,
max_length=self.model.config.max_position_embeddings, # Max length according to the model configuration
).to(self.device)
with torch.no_grad():
logits = self.model(**inputs).logits
# Apply transformations to get the sparse vectors
inter = torch.log1p(torch.relu(logits))
token_max = torch.max(inter, dim=1)
nz_tokens_i, nz_tokens_j = torch.where(token_max.values > 0)
output = []
for i in range(token_max.values.shape[0]):
nz_tokens = nz_tokens_j[nz_tokens_i == i]
nz_weights = token_max.values[i, nz_tokens]
output.append(
{"indices": nz_tokens.tolist(), "values": nz_weights.tolist()}
)
return output[0] if isinstance(inputs, str) else output
|