| import os |
| import time |
| from typing import Any |
|
|
| import numpy as np |
| import triton_python_backend_utils as pb_utils |
| from transformers import AutoTokenizer |
|
|
|
|
| class TritonPythonModel: |
| @staticmethod |
| def _decode_text(value: Any) -> str: |
| if isinstance(value, np.bytes_): |
| return value.tobytes().decode("utf-8") |
| if isinstance(value, bytes): |
| return value.decode("utf-8") |
| return str(value) |
|
|
| def initialize(self, args): |
| os.environ["HF_HUB_OFFLINE"] = "1" |
| os.environ["TRANSFORMERS_OFFLINE"] = "1" |
|
|
| self.logger = pb_utils.Logger |
| self.max_length = int(os.environ.get("BGE_M3_MAX_LENGTH", "128")) |
|
|
| model_repo = args["model_repository"] |
| tokenizer_path = os.path.join(model_repo, "1", "tokenizer") |
| if not os.path.isdir(tokenizer_path): |
| raise pb_utils.TritonModelException(f"tokenizer directory not found: {tokenizer_path}") |
|
|
| self.tokenizer = AutoTokenizer.from_pretrained( |
| tokenizer_path, |
| local_files_only=True, |
| ) |
| self.logger.log_info(f"bge_m3_preprocess initialized. tokenizer_path={tokenizer_path}") |
|
|
| def execute(self, requests): |
| responses = [] |
| st = time.time() |
|
|
| for request in requests: |
| input_tensor = pb_utils.get_input_tensor_by_name(request, "TEXTS") |
| texts_np = input_tensor.as_numpy().reshape(-1) |
| texts = [self._decode_text(v) for v in texts_np] |
|
|
| if len(texts) == 0: |
| input_ids = np.empty((0, 1), dtype=np.int64) |
| attention_mask = np.empty((0, 1), dtype=np.int64) |
| else: |
| encoded = self.tokenizer( |
| texts, |
| return_tensors="np", |
| padding=True, |
| truncation=True, |
| max_length=self.max_length, |
| ) |
| input_ids = np.asarray(encoded["input_ids"], dtype=np.int64) |
| attention_mask = np.asarray(encoded["attention_mask"], dtype=np.int64) |
|
|
| responses.append( |
| pb_utils.InferenceResponse( |
| output_tensors=[ |
| pb_utils.Tensor("input_ids", input_ids), |
| pb_utils.Tensor("attention_mask", attention_mask), |
| ] |
| ) |
| ) |
|
|
| self.logger.log_info(f"bge_m3_preprocess execute duration : {int((time.time() - st) * 1000)} ms") |
| return responses |
|
|
