Fixing return structure

handler.py

@@ -27,11 +27,13 @@ class EndpointHandler:
 
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
 
-        self.tokenizer = AutoTokenizer.from_pretrained(path)
-        self.model = AutoModelForCausalLM.from_pretrained(path, device_map="auto", 
-                                                          offload_folder='offload',
-                                                          trust_remote_code=True,
-                                                          load_in_8bit=True)
+        # self.tokenizer = AutoTokenizer.from_pretrained(path)
+        # self.model = AutoModelForCausalLM.from_pretrained(path, device_map="auto", 
+        #                                                   offload_folder='offload',
+        #                                                   trust_remote_code=True,
+        #                                                   load_in_8bit=True)
+        self.tokenizer = tokenizer
+        self.model = model
 
     def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
 
@@ -39,6 +41,7 @@ class EndpointHandler:
         if 'prompt' in data.keys():
             text = data['prompt']
         else:
+            print(data.keys())
             user_data = data.pop('query',data)
             text = self.prompt_ar.format_map({'Question':user_data})
             inputs = data.pop("inputs", data)
@@ -71,10 +74,10 @@ class EndpointHandler:
         response = self.tokenizer.batch_decode(generate_ids, 
                                                skip_special_tokens=True, 
                                                clean_up_tokenization_spaces=True)[0]
-        final_response = response.split("### Response: [|AI|]")
-        turn = [f'[|Human|] {query}', f'[|AI|] {final_response[-1]}']
-        chat_history.extend(turn)
         if 'prompt' in data.keys():
             return response
         else:
+            final_response = response.split("### Response: [|AI|]")
+            turn = [f'[|Human|] {query}', f'[|AI|] {final_response[-1]}']
+            chat_history.extend(turn)
             return {"response": final_response, "chat_history": chat_history}