|
| 1 | +import json |
| 2 | +import logging |
| 3 | +import os |
| 4 | +import zipfile |
| 5 | +from abc import ABC |
| 6 | + |
| 7 | +import torch |
| 8 | +import transformers |
| 9 | +from transformers import BloomForCausalLM, BloomTokenizerFast |
| 10 | + |
| 11 | +from ts.torch_handler.base_handler import BaseHandler |
| 12 | + |
| 13 | +logger = logging.getLogger(__name__) |
| 14 | +logger.info("Transformers version %s", transformers.__version__) |
| 15 | + |
| 16 | + |
| 17 | +TORCH_DTYPES = { |
| 18 | + "float16": torch.float16, |
| 19 | + "float32": torch.float32, |
| 20 | + "float64": torch.float64, |
| 21 | +} |
| 22 | + |
| 23 | + |
| 24 | +class TransformersSeqClassifierHandler(BaseHandler, ABC): |
| 25 | + """ |
| 26 | + Transformers handler class for sequence, token classification and question answering. |
| 27 | + """ |
| 28 | + |
| 29 | + def __init__(self): |
| 30 | + super(TransformersSeqClassifierHandler, self).__init__() |
| 31 | + self.initialized = False |
| 32 | + |
| 33 | + def initialize(self, ctx): |
| 34 | + """In this initialize function, the BERT model is loaded and |
| 35 | + the Layer Integrated Gradients Algorithm for Captum Explanations |
| 36 | + is initialized here. |
| 37 | + Args: |
| 38 | + ctx (context): It is a JSON Object containing information |
| 39 | + pertaining to the model artifacts parameters. |
| 40 | + """ |
| 41 | + self.manifest = ctx.manifest |
| 42 | + properties = ctx.system_properties |
| 43 | + model_dir = properties.get("model_dir") |
| 44 | + |
| 45 | + self.device = torch.device( |
| 46 | + "cuda:" + str(properties.get("gpu_id")) |
| 47 | + if torch.cuda.is_available() and properties.get("gpu_id") is not None |
| 48 | + else "cpu" |
| 49 | + ) |
| 50 | + # Loading the model and tokenizer from checkpoint and config files based on the user's choice of mode |
| 51 | + # further setup config can be added. |
| 52 | + with zipfile.ZipFile(model_dir + "/model.zip", "r") as zip_ref: |
| 53 | + zip_ref.extractall(model_dir + "/model") |
| 54 | + |
| 55 | + # read configs for the mode, model_name, etc. from setup_config.json |
| 56 | + setup_config_path = os.path.join(model_dir, "setup_config.json") |
| 57 | + if os.path.isfile(setup_config_path): |
| 58 | + with open(setup_config_path) as setup_config_file: |
| 59 | + self.setup_config = json.load(setup_config_file) |
| 60 | + else: |
| 61 | + logger.warning("Missing the setup_config.json file.") |
| 62 | + |
| 63 | + self.model = BloomForCausalLM.from_pretrained( |
| 64 | + model_dir + "/model", |
| 65 | + revision=self.setup_config["revision"], |
| 66 | + max_memory={ |
| 67 | + int(key) if key.isnumeric() else key: value |
| 68 | + for key, value in self.setup_config["max_memory"].items() |
| 69 | + }, |
| 70 | + low_cpu_mem_usage=self.setup_config["low_cpu_mem_usage"], |
| 71 | + device_map=self.setup_config["device_map"], |
| 72 | + offload_folder=self.setup_config["offload_folder"], |
| 73 | + offload_state_dict=self.setup_config["offload_state_dict"], |
| 74 | + torch_dtype=TORCH_DTYPES[self.setup_config["torch_dtype"]], |
| 75 | + ) |
| 76 | + |
| 77 | + self.tokenizer = BloomTokenizerFast.from_pretrained( |
| 78 | + model_dir + "/model", return_tensors="pt" |
| 79 | + ) |
| 80 | + |
| 81 | + self.model.eval() |
| 82 | + logger.info("Transformer model from path %s loaded successfully", model_dir) |
| 83 | + |
| 84 | + self.initialized = True |
| 85 | + |
| 86 | + def preprocess(self, requests): |
| 87 | + """Basic text preprocessing, based on the user's chocie of application mode. |
| 88 | + Args: |
| 89 | + requests (str): The Input data in the form of text is passed on to the preprocess |
| 90 | + function. |
| 91 | + Returns: |
| 92 | + list : The preprocess function returns a list of Tensor for the size of the word tokens. |
| 93 | + """ |
| 94 | + input_ids_batch = None |
| 95 | + attention_mask_batch = None |
| 96 | + for idx, data in enumerate(requests): |
| 97 | + input_text = data.get("data") |
| 98 | + if input_text is None: |
| 99 | + input_text = data.get("body") |
| 100 | + if isinstance(input_text, (bytes, bytearray)): |
| 101 | + input_text = input_text.decode("utf-8") |
| 102 | + |
| 103 | + max_length = self.setup_config["max_length"] |
| 104 | + logger.info("Received text: '%s'", input_text) |
| 105 | + |
| 106 | + inputs = self.tokenizer.encode_plus( |
| 107 | + input_text, |
| 108 | + max_length=int(max_length), |
| 109 | + pad_to_max_length=True, |
| 110 | + add_special_tokens=True, |
| 111 | + return_tensors="pt", |
| 112 | + ) |
| 113 | + |
| 114 | + input_ids = inputs["input_ids"].to(self.device) |
| 115 | + attention_mask = inputs["attention_mask"].to(self.device) |
| 116 | + # making a batch out of the recieved requests |
| 117 | + # attention masks are passed for cases where input tokens are padded. |
| 118 | + if input_ids.shape is not None: |
| 119 | + if input_ids_batch is None: |
| 120 | + input_ids_batch = input_ids |
| 121 | + attention_mask_batch = attention_mask |
| 122 | + else: |
| 123 | + input_ids_batch = torch.cat((input_ids_batch, input_ids), 0) |
| 124 | + attention_mask_batch = torch.cat( |
| 125 | + (attention_mask_batch, attention_mask), 0 |
| 126 | + ) |
| 127 | + return (input_ids_batch, attention_mask_batch) |
| 128 | + |
| 129 | + def inference(self, input_batch): |
| 130 | + """Predict the class (or classes) of the received text using the |
| 131 | + serialized transformers checkpoint. |
| 132 | + Args: |
| 133 | + input_batch (list): List of Text Tensors from the pre-process function is passed here |
| 134 | + Returns: |
| 135 | + list : It returns a list of the predicted value for the input text |
| 136 | + """ |
| 137 | + (input_ids_batch, _) = input_batch |
| 138 | + inferences = [] |
| 139 | + input_ids_batch = input_ids_batch.to(self.device) |
| 140 | + outputs = self.model.generate( |
| 141 | + input_ids_batch, |
| 142 | + do_sample=True, |
| 143 | + max_new_tokens=int(self.setup_config["max_length"]), |
| 144 | + top_p=0.95, |
| 145 | + top_k=60, |
| 146 | + ) |
| 147 | + for i, _ in enumerate(outputs): |
| 148 | + inferences.append( |
| 149 | + self.tokenizer.decode(outputs[i], skip_special_tokens=True) |
| 150 | + ) |
| 151 | + |
| 152 | + logger.info("Generated text: '%s'", inferences) |
| 153 | + |
| 154 | + print("Generated text", inferences) |
| 155 | + return inferences |
| 156 | + |
| 157 | + def postprocess(self, inference_output): |
| 158 | + """Post Process Function converts the predicted response into Torchserve readable format. |
| 159 | + Args: |
| 160 | + inference_output (list): It contains the predicted response of the input text. |
| 161 | + Returns: |
| 162 | + (list): Returns a list of the Predictions and Explanations. |
| 163 | + """ |
| 164 | + return inference_output |
0 commit comments