Update model_inference.py

model_inference.py

@@ -47,31 +47,6 @@ def generate_question(agents, question):
 
     return agent_contexts, content
 
-def generate_answer(model, formatted_prompt):
-    API_URL = endpoint_dict[model]
-    headers = {"Authorization": f"Bearer {args.auth_token}"}
-    payload = {"inputs": formatted_prompt}
-    try:
-        resp = requests.post(API_URL, json=payload, headers=headers)
-        response = resp.json()
-    except:
-        print("retrying due to an error......")
-        time.sleep(5)
-        return generate_answer(API_URL, headers, payload)
-    
-    return {"model": model, "content": response[0]["generated_text"].split(prompt_dict[model]["response_split"])[-1]}
-
-def prompt_formatting(model, instruction, cot):
-    if model == "alpaca" or model == "orca":
-        prompt = prompt_dict[model]["prompt_no_input"]
-    else:
-        prompt = prompt_dict[model]["prompt"]
-    
-    if cot:
-        instruction += "Let's think step by step."
-
-    return {"model": model, "content": prompt.format(instruction)}
-
 def Inference(model_list, question, API_KEY, auth_token, round, cot):
     if len(model_list) != 3:
         raise ValueError("Please choose just '3' models! Neither more nor less!")
@@ -80,6 +55,31 @@ def Inference(model_list, question, API_KEY, auth_token, round, cot):
 
     prompt_dict, endpoint_dict = load_json("src/prompt_template.json", "src/inference_endpoint.json")
 
+    def generate_answer(model, formatted_prompt):
+        API_URL = endpoint_dict[model]["API_URL"]
+        headers = endpoint_dict[model]["headers"]
+        payload = {"inputs": formatted_prompt}
+        try:
+            resp = requests.post(API_URL, json=payload, headers=headers)
+            response = resp.json()
+        except:
+            print("retrying due to an error......")
+            time.sleep(5)
+            return generate_answer(API_URL, headers, payload)
+        
+        return {"model": model, "content": response[0]["generated_text"].split(prompt_dict[model]["response_split"])[-1]}
+    
+    def prompt_formatting(model, instruction, cot):
+        if model == "alpaca" or model == "orca":
+            prompt = prompt_dict[model]["prompt_no_input"]
+        else:
+            prompt = prompt_dict[model]["prompt"]
+        
+        if cot:
+            instruction += "Let's think step by step."
+    
+        return {"model": model, "content": prompt.format(instruction)}
+    
     agents = len(model_list)
     rounds = round