File size: 1,156 Bytes
a35c2e5 16db4fc a35c2e5 16db4fc a35c2e5 16db4fc a35c2e5 16db4fc a35c2e5 16db4fc |
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 |
from typing import Dict, Any
import logging
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch.cuda
device = "cuda" if torch.cuda.is_available() else "cpu"
LOGGER = logging.getLogger(__name__)
class EndpointHandler():
def __init__(self, path=""):
self.model = AutoModelForCausalLM.from_pretrained("Ozgur98/pushed_model_mosaic_small", load_in_8bit=True, device_map='auto')
self.tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b")
# Load the Lora model
def __call__(self, data):
"""
Args:
data (Dict): The payload with the text prompt and generation parameters.
"""
print("CALLED")
LOGGER.info(data)
# Forward
LOGGER.info(f"Start generation.")
tokenized_example = tokenizer(data, return_tensors='pt')
outputs = self.model.generate(tokenized_example['input_ids'].to('cuda:0'), max_new_tokens=100, do_sample=True, top_k=10, top_p = 0.95)
# Postprocess
answer = tokenizer.batch_decode(outputs, skip_special_tokens=True)
prompt = answer[0].rstrip()
return prompt |