| | import torch |
| | import gc |
| | from ts.torch_handler.base_handler import BaseHandler |
| | from transformers import GPT2LMHeadModel |
| |
|
| | import logging |
| |
|
| | logger = logging.getLogger(__name__) |
| |
|
| |
|
| | class SampleTransformerModel(BaseHandler): |
| | def __init__(self): |
| | super(SampleTransformerModel, self).__init__() |
| | self.model = None |
| | self.device = None |
| | self.initialized = False |
| |
|
| | def load_model(self, model_dir): |
| | self.model = GPT2LMHeadModel.from_pretrained(model_dir, return_dict=True) |
| | self.model.to(self.device) |
| |
|
| | def initialize(self, ctx): |
| | |
| | properties = ctx.system_properties |
| | model_dir = properties.get("model_dir") |
| | self.device = torch.device("cuda:" + str(properties.get("gpu_id")) if torch.cuda.is_available() else "cpu") |
| |
|
| | self.load_model(model_dir) |
| |
|
| | self.model.eval() |
| | self.initialized = True |
| |
|
| | def preprocess(self, requests): |
| | input_batch = {} |
| | for idx, data in enumerate(requests): |
| | input_ids = torch.tensor([data.get("body").get("text")]).to(self.device) |
| | input_batch["input_ids"] = input_ids |
| | input_batch["num_samples"] = data.get("body").get("num_samples") |
| | input_batch["length"] = data.get("body").get("length") + len(data.get("body").get("text")) |
| | del requests |
| | gc.collect() |
| | return input_batch |
| |
|
| | def inference(self, input_batch): |
| | input_ids = input_batch["input_ids"] |
| | length = input_batch["length"] |
| |
|
| | inference_output = self.model.generate(input_ids, |
| | bos_token_id=self.model.config.bos_token_id, |
| | eos_token_id=self.model.config.eos_token_id, |
| | pad_token_id=self.model.config.eos_token_id, |
| | do_sample=True, |
| | max_length=length, |
| | top_k=50, |
| | top_p=0.95, |
| | no_repeat_ngram_size=2, |
| | num_return_sequences=input_batch["num_samples"]) |
| |
|
| | if torch.cuda.is_available(): |
| | torch.cuda.empty_cache() |
| | del input_batch |
| | gc.collect() |
| | return inference_output |
| |
|
| | def postprocess(self, inference_output): |
| | output = inference_output.cpu().numpy().tolist() |
| | del inference_output |
| | gc.collect() |
| | return [output] |
| |
|
| | def handle(self, data, context): |
| | |
| | data = self.preprocess(data) |
| | data = self.inference(data) |
| | data = self.postprocess(data) |
| | return data |
| |
|
