| from transformers import Pipeline |
| import torch.nn.functional as F |
| import torch |
| from ts.torch_handler.base_handler import BaseHandler |
| import logging |
| import os |
| import transformers |
| from transformers import AutoTokenizer |
| logger = logging.getLogger(__name__) |
| logger.info("Transformers version %s", transformers.__version__) |
| from optimum.onnxruntime import ORTModelForFeatureExtraction |
|
|
| def mean_pooling(model_output, attention_mask): |
| token_embeddings = model_output[0] |
| input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float() |
| return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9) |
| |
| class SentenceEmbeddingHandler(BaseHandler): |
| def __init__(self): |
| super(SentenceEmbeddingHandler, self).__init__() |
| self._context = None |
| self.initialized = False |
| class SentenceEmbeddingPipeline(Pipeline): |
| def initialize(self, context): |
| """ |
| Initialize function loads the model and the tokenizer |
| |
| Args: |
| context (context): It is a JSON Object containing information |
| pertaining to the model artifacts parameters. |
| |
| Raises: |
| RuntimeError: Raises the Runtime error when the model or |
| tokenizer is missing |
| """ |
|
|
| properties = context.system_properties |
| self.manifest = context.manifest |
| model_dir = properties.get("model_dir") |
|
|
| |
| self.device = torch.device( |
| "cuda:" + str(properties.get("gpu_id")) |
| if torch.cuda.is_available() and properties.get("gpu_id") is not None |
| else "cpu" |
| ) |
| logger.info(f'Using device {self.device}') |
|
|
| |
| model_file = self.manifest['model']['modelFile'] |
| model_path = os.path.join(model_dir, model_file) |
|
|
| if os.path.isfile(model_path): |
| |
| self.model = ORTModelForFeatureExtraction.from_pretrained(model_dir, file_name="model_optimized.onnx") |
| self.model.to(self.device) |
| |
| logger.info(f'Successfully loaded model from {model_file}') |
| else: |
| raise RuntimeError('Missing the model file') |
|
|
| |
| self.tokenizer = AutoTokenizer.from_pretrained(model_dir) |
| if self.tokenizer is not None: |
| logger.info('Successfully loaded tokenizer') |
| else: |
| raise RuntimeError('Missing tokenizer') |
|
|
| self.initialized = True |
| def _sanitize_parameters(self, **kwargs): |
| |
| preprocess_kwargs = {} |
| return preprocess_kwargs, {}, {} |
| |
| def preprocess_text(self, inputs): |
| encoded_inputs = self.tokenizer(inputs, padding=True, truncation=True, return_tensors='pt') |
| return encoded_inputs |
| |
| def preprocess(self, requests): |
| """ |
| Tokenize the input text using the suitable tokenizer and convert |
| it to tensor |
| |
| If token_ids is provided, the json must be of the form |
| {'input_ids': [[101, 102]], 'token_type_ids': [[0, 0]], 'attention_mask': [[1, 1]]} |
| |
| Args: |
| requests: A list containing a dictionary, might be in the form |
| of [{'body': json_file}] or [{'data': json_file}] or [{'token_ids': json_file}] |
| Returns: |
| the tensor containing the batch of token vectors. |
| """ |
|
|
| |
| data = requests[0].get('body') |
| if data is None: |
| data = requests[0].get('data') |
| |
| texts = data.get('input') |
| if texts is not None: |
| logger.info('Text provided') |
| return self.preprocess_text(texts) |
|
|
| encodings = data.get('encodings') |
| if encodings is not None: |
| logger.info('Encodings provided') |
| return transformers.BatchEncoding(data={k: torch.tensor(v) for k, v in encodings.items()}) |
|
|
| raise Exception("unsupported payload") |
| def inference(self, model_inputs): |
| outputs = self.model(**model_inputs) |
| sentence_embeddings = mean_pooling(outputs, model_inputs['attention_mask']) |
| sentence_embeddings = F.normalize(sentence_embeddings, p=2, dim=1) |
| return sentence_embeddings |
| |
| def postprocess(self, outputs): |
| formatted_outputs = [] |
| data=[outputs.tolist()] |
| for dat in data: |
| formatted_outputs.append({"status":"success","data":dat}) |
| return formatted_outputs |
