| 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 |
|
|
