handler.py

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