Update handler.py

handler.py

@@ -6,10 +6,9 @@ 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")
+        self.model = AutoModelForCausalLM.from_pretrained("Ozgur98/pushed_model_mosaic_small", trust_remote_code=True).to(device='cuda:0', dtype=torch.bfloat16)
         self.tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b")
         # Load the Lora model
         
@@ -18,13 +17,13 @@ class EndpointHandler():
         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')
+        tokenized_example = self.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)
+        answer = self.tokenizer.batch_decode(outputs, skip_special_tokens=True)
         prompt = answer[0].rstrip()
         return prompt