Update model_inference.py

model_inference.py

@@ -13,32 +13,42 @@ def load_json(prompt_path, endpoint_path):
 
     return prompt_dict, endpoint_dict
 
-def construct_message(agents, instruction, idx):
-    if len(agents) == 0:
-        prompt = "Can you double check that your answer is correct. Please reiterate your answer, making sure to state your answer at the end of the response."
-        return prompt
-    
-    contexts = [agents[0][idx]['content'], agents[1][idx]['content'], agents[2][idx]['content']]
+def construct_message(agent_context, instruction, idx):
+    prefix_string = "Here are a list of opinions from different agents: "
+
+    prefix_string = prefix_string + agent_context + "\n\n Write a summary of the different opinions from each of the individual agent."
+
+    message = [{"role": "user", "content": prefix_string}]
+
+    try:
+        completion = openai.ChatCompletion.create(
+            model="gpt-3.5-turbo-0613",
+            messages=message,
+            max_tokens=256,
+            n=1
+        )['choices'][0]['message']['content']
+    except:
+        print("retrying ChatGPT due to an error......")
+        time.sleep(5)
+        return construct_message(agent_context, instruction, idx)
+
+    prefix_string = f"Here is a summary of responses from other agents: {completion}"
+    prefix_string = prefix_string + "\n\n Use this summarization carefully as additional advice, can you provide an updated answer? Make sure to state your answer at the end of the response." + instruction
+    return prefix_string
 
-    # system prompt & user prompt for gpt-3.5-turbo
-    sys_prompt = f"I want you to act as a summarizer. You can look at multiple responses and summarize the main points of them so that the meaning is not lost. Multiple responses will be given, which are responses from several different models to a single question. And you should use your excellent summarizing skills to output the best summary."
-    summarize_prompt = f"[Response 1]: {contexts[0]}\n[Response 2]: {contexts[1]}\nResponse 3: {contexts[2]}\n\nThese are response of each model to a certain question. Summarize comprehensively without compromising the meaning of each response."
+def summarize_message(agent_contexts, instruction, idx):
+    prefix_string = "Here are a list of opinions from different agents: "
 
-    message = [
-        {"role": "system", "content": sys_prompt},
-        {"role": "user", "content": summarize_prompt},
-    ]
+    for agent in agent_contexts:
+        agent_response = agent[-1]["content"]
+        response = "\n\n One agent response: ```{}```".format(agent_response)
 
-    completion = openai.ChatCompletion.create(
-        model="gpt-3.5-turbo-16k-0613",
-        messages=message,
-        max_tokens=256,
-        n=1
-    )
+        prefix_string = prefix_string + response
 
-    prefix_string = f"This is the summarization of recent/updated opinions from other agents: {completion}"
-    prefix_string = prefix_string + "\n\n Use this summarization carefully as additional advice, can you provide an updated answer? Make sure to state your answer at the end of the response." + instruction
-    return prefix_string
+    prefix_string = prefix_string + "\n\n Write a summary of the different opinions from each of the individual agent."
+    completion = construct_message(prefix_string, instruction, idx)
+
+    return completion
 
 def generate_question(agents, question):
     agent_contexts = [[{"model": agent, "content": question}] for agent in agents]
@@ -47,7 +57,7 @@ def generate_question(agents, question):
 
     return agent_contexts, content
 
-def Inference(model_list, question, API_KEY, auth_token, round, cot):
+def Inference(model_list, question, API_KEY, cot):
     if len(model_list) != 3:
         raise ValueError("Please choose just '3' models! Neither more nor less!")
         
@@ -58,16 +68,21 @@ def Inference(model_list, question, API_KEY, auth_token, round, cot):
     def generate_answer(model, formatted_prompt):
         API_URL = endpoint_dict[model]["API_URL"]
         headers = endpoint_dict[model]["headers"]
-        payload = {"inputs": formatted_prompt}
+        payload = {
+            "inputs": formatted_prompt,
+            "parameters": {
+                "max_new_tokens": 256
+            }
+        }
         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 generate_answer(model, formatted_prompt)
         
-        return {"model": model, "content": response[0]["generated_text"].split(prompt_dict[model]["response_split"])[-1]}
+        return {"model": model, "content": response[0]["generated_text"]}
     
     def prompt_formatting(model, instruction, cot):
         if model == "alpaca" or model == "orca":
@@ -77,37 +92,37 @@ def Inference(model_list, question, API_KEY, auth_token, round, cot):
         
         if cot:
             instruction += "Let's think step by step."
-    
-        return {"model": model, "content": prompt.format(instruction)}
-    
+
+        return {"model": model, "content": prompt.format(instruction=instruction)}
+
     agents = len(model_list)
-    rounds = round
+    rounds = 2
 
-    generated_description = []
+    agent_contexts, content = generate_question(agents=model_list, question=args.question)
 
-    agent_contexts, content = generate_question(agents=model_list, question=question)
+    message = []
 
     # Debate
     for debate in range(rounds+1):
         # Refer to the summarized previous response
         if debate != 0:
-            message = construct_message(agent_contexts, content, 2 * debate - 1)
-            for i in range(agent_contexts):
+            message.append(summarize_message(agent_contexts, content, 2 * debate - 1))
+            for i in range(len(agent_contexts)):
                 agent_contexts[i].append(prompt_formatting(agent_contexts[i][-1]["model"], message, args.cot))
 
         # Generate new response based on summarized response
         for agent_context in agent_contexts:
-            completion = generate_answer(agent_context[-1]["model"], agent_context[-1]["content"] if debate != 0 else prompt_formatting(agent_context[-1]["model"], agent_context[-1]["content"], args.cot)["content"])
+            completion = generate_answer(agent_context[-1]["model"], agent_context[-1]["content"])
             agent_context.append(completion)
 
     models_response = {
-        f"{args.m1}": [agent_contexts[0][1]["content"], agent_contexts[0][3]["content"], agent_contexts[0][-1]["content"]],
-        f"{args.m2}": [agent_contexts[1][1]["content"], agent_contexts[1][3]["content"], agent_contexts[1][-1]["content"]],
-        f"{args.m3}": [agent_contexts[2][1]["content"], agent_contexts[2][3]["content"], agent_contexts[2][-1]["content"]]
+        f"{model_list[0]}": [agent_contexts[0][1]["content"], agent_contexts[0][3]["content"], agent_contexts[0][-1]["content"]],
+        f"{model_list[1]}": [agent_contexts[1][1]["content"], agent_contexts[1][3]["content"], agent_contexts[1][-1]["content"]],
+        f"{model_list[2]}": [agent_contexts[2][1]["content"], agent_contexts[2][3]["content"], agent_contexts[2][-1]["content"]]
     }
     response_summarization = [
-        agent_contexts[0][2], agent_contexts[0][4]
+        message[0], message[1]
     ]
-    generated_description.append({"question": content, "agent_response": models_response, "summarization": response_summarization})
+    generated_description = {"question": content, "agent_response": models_response, "summarization": response_summarization})
 
     return generated_description