Update handler.py

handler.py

@@ -9,30 +9,22 @@ LOGGER = logging.getLogger(__name__)
 
 class EndpointHandler():
     def __init__(self, path=""):
-        self.model = AutoModelForCausalLM.from_pretrained("bigscience/bloom-3b", load_in_8bit=True, device_map='auto')
-        self.tokenizer = AutoTokenizer.from_pretrained("bigscience/bloom-3b")
+        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: Dict[str, Any]) -> Dict[str, Any]:
+    def __call__(self, data):
         """
         Args:
             data (Dict): The payload with the text prompt and generation parameters.
         """
-        LOGGER.info(f"Received data: {data}")
-        # Get inputs
-        prompt = data.pop("inputs", None)
-        parameters = data.pop("parameters", None)
-        if prompt is None:
-            raise ValueError("Missing prompt.")
-        # Preprocess
-        input_ids = self.tokenizer(prompt, return_tensors="pt").input_ids.to(device)
+        print("CALLED")
+        LOGGER.info(data)
         # Forward
         LOGGER.info(f"Start generation.")
-        if parameters is not None:
-            output = self.model.generate(input_ids=input_ids, **parameters)
-        else:
-            output = self.model.generate(input_ids=input_ids)
+        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
-        prediction = self.tokenizer.decode(output[0])
-        LOGGER.info(f"Generated text: {prediction}")
-        return {"generated_text": prediction}
+        answer = tokenizer.batch_decode(outputs, skip_special_tokens=True)
+        prompt = answer[0].rstrip()
+        return prompt