File size: 1,818 Bytes
e53226a 3a46ab7 e53226a | 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 | from transformers import GPT2LMHeadModel, GPT2Tokenizer, pipeline
import torch
import json
class EndpointHandler:
def __init__(self, path=""):
# Load tokenizer and model
self.tokenizer = GPT2Tokenizer.from_pretrained(path)
self.model = GPT2LMHeadModel.from_pretrained(
path,
torch_dtype=torch.float16,
device_map="auto"
)
# Set pad token if not set
if self.tokenizer.pad_token is None:
self.tokenizer.pad_token = self.tokenizer.eos_token
def __call__(self, data):
"""
Args:
data (dict): The payload sent to the endpoint
Returns:
dict: prediction response
"""
inputs = data.pop("inputs", "")
parameters = data.pop("parameters", {})
# Set default parameters
max_length = parameters.get("max_length", 100)
temperature = parameters.get("temperature", 0.7)
do_sample = parameters.get("do_sample", True)
top_p = parameters.get("top_p", 0.9)
# Tokenize input
input_ids = self.tokenizer.encode(inputs, return_tensors="pt")
# Generate
with torch.no_grad():
outputs = self.model.generate(
input_ids,
max_length=max_length,
temperature=temperature,
do_sample=do_sample,
top_p=top_p,
pad_token_id=self.tokenizer.pad_token_id,
eos_token_id=self.tokenizer.eos_token_id,
attention_mask=torch.ones_like(input_ids)
)
# Decode output
generated_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
return {"generated_text": generated_text} |