File size: 2,803 Bytes
baa316f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 |
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):
# self.manifest = ctx.manifest
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):
# self.context = context
data = self.preprocess(data)
data = self.inference(data)
data = self.postprocess(data)
return data
|