fixed reflexion pipeline

Author: benshi34

USACOBench/prompts.py

@@ -32,7 +32,7 @@ def retrieval_prompt_fn(query, retrieval_type=RetrievalType.EPISODIC):
 [END PROBLEM]"""
 
 def reflexion_prompt_fn(query, retrieval=False):
-    retrieval = ""
+    retrieval_text = ""
     if retrieval:
         retrieval_text = "You were also given a couple of similar problems to the problem above along with their solutions to aid you in solving the problem at hand. Here are the similar problems you were given:\n" + query['retrieval_text'] 
 

models.py

@@ -105,7 +105,7 @@ def chatgpt(messages, model="gpt-4", temperature=0.7, max_tokens=2000, n=1, stop
 def chatgpt_raw(messages, model="gpt-4", temperature=0.7, max_tokens=2000, n=1, stop=None, **kwargs) -> list:
     return chatgpts_raw([messages] * n, model=model, temperature=temperature, max_tokens=max_tokens, stop=stop, **kwargs)[0]
 
-def chatgpts(messages_list, model="gpt-4", temperature=0.7, max_tokens=2000, stop=None, max_messages=200, **kwargs) -> list:
+def chatgpts(messages_list, model="gpt-4", temperature=0.7, max_tokens=2000, stop=None, max_messages=400, **kwargs) -> list:
     texts = []
     for i in range(0, len(messages_list), max_messages):
         responses = asyncio.run(generate_from_openai_chat_completion(model=model, messages_list=messages_list[i: i + max_messages], temperature=temperature, max_tokens=max_tokens, top_p=1, stop=stop, **kwargs))
@@ -115,7 +115,7 @@ def chatgpts(messages_list, model="gpt-4", temperature=0.7, max_tokens=2000, sto
         # prompt_tokens[model] += sum(x["usage"]["prompt_tokens"] for x in responses if "usage" in x and "prompt_tokens" in x["usage"])
     return texts
 
-def chatgpts_raw(messages_list, model="gpt-4", temperature=0.7, max_tokens=2000, stop=None, max_messages=200, **kwargs) -> list:
+def chatgpts_raw(messages_list, model="gpt-4", temperature=0.7, max_tokens=2000, stop=None, max_messages=400, **kwargs) -> list:
     '''
     Returns raw response messages, not just the text content
     '''
@@ -128,7 +128,7 @@ def chatgpts_raw(messages_list, model="gpt-4", temperature=0.7, max_tokens=2000,
         # prompt_tokens[model] += sum(x["usage"]["prompt_tokens"] for x in responses if "usage" in x and "prompt_tokens" in x["usage"])
     return responses_all
 
-def claude(prompts, model="claude-3-sonnet-20240229", temperature=0.7, max_tokens=3000, stop=None, max_messages=200, system_prompt=None, **kwargs) -> list:
+def claude(prompts, model="claude-3-sonnet-20240229", temperature=0.7, max_tokens=3000, stop=None, max_messages=400, system_prompt=None, **kwargs) -> list:
     texts = []
     if system_prompt is not None:
         messages_list = [[{'role': 'system', 'content': system_prompt},

run_usaco.py

@@ -25,7 +25,7 @@
 parser.add_argument('-s', '--semantic_retrieval', help='whether to use semantic retrieval', action="store_true", default=False)
 parser.add_argument('-r', '--reflexion', help='whether to use reflexion', action="store_true", default=False)
 parser.add_argument('-a', '--attempts', help='number of attempts', default=1)
-parser.add_argument('-n', '--num_reflexion', help='number of reflexion iterations', default=3)
+parser.add_argument('-n', '--num_reflexion', help='number of reflexion iterations', default=2)
 args = parser.parse_args()
 
 model_name = args.model_name
@@ -39,9 +39,10 @@
 problem_dict = load_problem_dict('usaco_subset307')
 model_fn = partial(model_fn, model=model_name)
 
+# A little redundant but it does the job and it's readable...
 if not args.episodic_retrieval and not args.semantic_retrieval and not args.reflexion:
     rdict, sdict, rs, ss = run_solve(model_fn, model_name, problem_dict, args.attempts)
-    
+
 elif args.episodic_retrieval and not args.semantic_retrieval and not args.reflexion:
     rdict, sdict, rs, ss = run_solve(model_fn, model_name, problem_dict, args.attempts)
     rdict, sdict, rs, ss = run_retrieval(model_fn, model_name, problem_dict, args.attempts, ss, args.num_retrieved, RetrievalType.EPISODIC)
@@ -55,11 +56,48 @@
     rdict, sdict, rs, ss = run_retrieval(model_fn, model_name, problem_dict, args.attempts, ss, args.num_retrieved, RetrievalType.EPISODIC_SEMANTIC)
 
 elif not args.episodic_retrieval and not args.semantic_retrieval and args.reflexion:
-    rdict, sdict, rs, ss, queries = run_solve(model_fn, model_name, problem_dict, args.attempts, return_queries=True)
-    reflexions = []
+    rdict, sdict, rs, ss = run_solve(model_fn, model_name, problem_dict, args.attempts)
+    reflexions = [rdict]
+    query_dict = None
+    for i in range(args.num_reflexion):
+        rdict, sdict, rs, ss, query_dict = run_reflexion(model_fn, model_name, problem_dict, args.attempts, rdict, sdict, query_dict, i, return_queries=True)
+        reflexions.append(rdict)
+
+    rs = calculate_final_rs(reflexions, problem_dict)
+
+elif args.episodic_retrieval and not args.semantic_retrieval and args.reflexion:
+    rdict, sdict, rs, ss = run_solve(model_fn, model_name, problem_dict, args.attempts)
+    rdict, sdict, rs, ss = run_retrieval(model_fn, model_name, problem_dict, args.attempts, ss, args.num_retrieved, RetrievalType.EPISODIC)
+
+    reflexions = [rdict]
+    query_dict = None
+    for i in range(args.num_reflexion):
+        rdict, sdict, rs, ss, query_dict = run_reflexion(model_fn, model_name, problem_dict, args.attempts, rdict, sdict, query_dict, i, return_queries=True, retrieval=True)
+        reflexions.append(rdict)
+
+    rs = calculate_final_rs(reflexions, problem_dict)
+    
+elif not args.episodic_retrieval and args.semantic_retrieval and args.reflexion:
+    rdict, sdict, rs, ss = run_solve(model_fn, model_name, problem_dict, args.attempts)
+    rdict, sdict, rs, ss = run_retrieval(model_fn, model_name, problem_dict, args.attempts, ss, args.num_retrieved, RetrievalType.SEMANTIC)
+
+    reflexions = [rdict]
+    query_dict = None
+    for i in range(args.num_reflexion):
+        rdict, sdict, rs, ss, query_dict = run_reflexion(model_fn, model_name, problem_dict, args.attempts, rdict, sdict, query_dict, i, return_queries=True, retrieval=True)
+        reflexions.append(rdict)
+
+    rs = calculate_final_rs(reflexions, problem_dict)
+
+elif args.episodic_retrieval and args.semantic_retrieval and args.reflexion:
+    rdict, sdict, rs, ss = run_solve(model_fn, model_name, problem_dict, args.attempts)
+    rdict, sdict, rs, ss = run_retrieval(model_fn, model_name, problem_dict, args.attempts, ss, args.num_retrieved, RetrievalType.EPISODIC_SEMANTIC)
+
+    reflexions = [rdict]
+    query_dict = None
     for i in range(args.num_reflexion):
-        rdict, sdict, rs, ss, queries = run_reflexion(model_fn, model_name, problem_dict, args.attempts, rdict, sdict, queries, i, return_queries=True)
-        reflexions.append(rs)
+        rdict, sdict, rs, ss, query_dict = run_reflexion(model_fn, model_name, problem_dict, args.attempts, rdict, sdict, query_dict, i, return_queries=True, retrieval=True)
+        reflexions.append(rdict)
 
     rs = calculate_final_rs(reflexions, problem_dict)
 

utils.py

@@ -221,33 +221,32 @@ def generate_episodic_semantic_retrieval_queries(num_problems_fetched, problem_d
     save_json(final_queries, 'queries_firstsolve_episodic_semantic')
     return final_queries
 
-def generate_reflexion_queries(rdict, sdict, problem_dict, model_name, prev_queries, retrieval=False):
+def generate_reflexion_queries(rdict, sdict, problem_dict, model_name, iteration, prev_queries_dict=None, retrieval=False):
     reflection_queries_dict = dict()
-    prev_queries_dict = dict()
-    for query in prev_queries:
-        prev_queries_dict[query['problem_id']] = query
 
-    # Extracting Original Response
     for problem_id in sdict.keys():
         if problem_id in problem_dict.keys():
             for solution in sdict[problem_id][:1]:
-                original_response = solution['solution']
+                prev_buffer = ''
+                if prev_queries_dict:
+                    prev_buffer = prev_queries_dict[problem_id]['reflection_buffer']
+                current_response = solution['solution']
+                current_execution_output = rdict[solution['problem_id']][0]['result_list']
                 num_samples = problem_dict[problem_id]['description'].count("SAMPLE INPUT")
-                unparsed_execution_output = rdict[problem_id][0]['result_list']
                 execution_output = ""
-                if unparsed_execution_output:
-                    unparsed_execution_output = unparsed_execution_output[:num_samples]
-                    for i, result in enumerate(unparsed_execution_output):
+                if current_execution_output:
+                    current_execution_output = current_execution_output[:num_samples]
+                    for i, result in enumerate(current_execution_output):
                         execution_output += f"Test Case {i}\n" + result['status'] + "\n"
                 else:
                     execution_output = "No submission, formatting error during judging."
+
                 if retrieval:
                     retrieval_text = prev_queries_dict[problem_id]['retrieval_text']
                     retrieval_problem_ids = prev_queries_dict[problem_id]['retrieval_problem_ids']
-                    reflection_queries_dict[problem_id] = {'problem_id': problem_id, 'original_response': original_response, 'execution_response': execution_output, 'retrieval_text': retrieval_text, 'retrieval_problem_ids': retrieval_problem_ids, 'problem_description': problem_dict[problem_id]['description']}
+                    reflection_queries_dict[problem_id] = {'problem_id': problem_id, 'reflection_buffer': prev_buffer + f'\n Reflection Response Number {iteration+1}: \n' + current_response + f'\n Reflection Response Execution Output Number {iteration+1}:\n' + execution_output, 'retrieval_text': retrieval_text, 'retrieval_problem_ids': retrieval_problem_ids, 'problem_description': problem_dict[problem_id]['description']}
                 else:
-                    reflection_queries_dict[problem_id] = {'problem_id': problem_id, 'original_response': original_response, 'execution_response': execution_output, 'problem_description': problem_dict[problem_id]['description']}
-    
+                    reflection_queries_dict[problem_id] = {'problem_id': problem_id, 'reflection_buffer': prev_buffer + f'\n Reflection Response Number {iteration+1} \n' + current_response + f'\n Reflection Response Execution Output Number {iteration+1}:\n' + execution_output, 'problem_description': problem_dict[problem_id]['description']}
     if retrieval:
         name = f'queries_dict_{model_name}_retrieval_reflexion'
     else:
@@ -257,7 +256,7 @@ def generate_reflexion_queries(rdict, sdict, problem_dict, model_name, prev_quer
 
 def calculate_final_rs(reflexions, problem_dict):
     rs = []
-    for problem_id in problem_dict.keys():
+    for problem_id in reflexions[0].keys():
         num_samples = problem_dict[problem_id]['description'].count('SAMPLE INPUT')
         for i, reflexion_result in enumerate(reflexions):
             if reflexion_result[problem_id][0]['result_list']:
@@ -288,7 +287,7 @@ def run_solve(model_fn, model_name, problem_dict, attempts, return_queries=False
     for problem_id in problem_dict.keys():
         queries.append({'problem_id': problem_id, 'problem_description': problem_dict[problem_id]['description']})
 
-    rdict, sdict, rs, ss = evaluate_model(model_fn, solve_prompt_fn, queries=queries, verbose=True, attempts=attempts, problem_ids=list(problem_dict.keys()))
+    rdict, sdict, rs, ss = evaluate_model(model_fn, solve_prompt_fn, queries=queries, verbose=True, attempts=attempts, problem_ids=list(problem_dict.keys())[:2])
     save_json([rdict, sdict, rs, ss], f'results/results_{model_name}_solve_{attempts}attempts')
     return (rdict, sdict, rs, ss) if not return_queries else (rdict, sdict, rs, ss, queries)
 
@@ -306,9 +305,9 @@ def run_retrieval(model_fn, model_name, problem_dict, attempts, solution_sets, n
 
     return (rdict, sdict, rs, ss) if not return_queries else (rdict, sdict, rs, ss, queries)
 
-def run_reflexion(model_fn, model_name, problem_dict, attempts, prev_result_dict, prev_solution_dict, prev_queries, iteration, return_queries=True):
-    new_reflexion_queries = generate_reflexion_queries(prev_result_dict, prev_solution_dict, problem_dict, model_name, prev_queries)
-    rdict, sdict, rs, ss = evaluate_model(model_fn, reflexion_prompt_fn, queries=new_reflexion_queries, verbose=True, attempts=attempts, problem_ids=list(problem_dict.keys()))
+def run_reflexion(model_fn, model_name, problem_dict, attempts, prev_result_dict, prev_solution_dict, prev_queries_dict, iteration, return_queries=True, retrieval=False):
+    new_reflexion_queries_dict = generate_reflexion_queries(prev_result_dict, prev_solution_dict, problem_dict, model_name, iteration, prev_queries_dict=prev_queries_dict, retrieval=retrieval)
+    rdict, sdict, rs, ss = evaluate_model(model_fn, reflexion_prompt_fn, queries=list(new_reflexion_queries_dict.values()), verbose=True, attempts=attempts)
     save_json([rdict, sdict, rs, ss], f'results_{model_name}_reflexion_{str(iteration)}iteration')
 
-    return (rdict, sdict, rs, ss) if not return_queries else (rdict, sdict, rs, ss, new_reflexion_queries)
+    return (rdict, sdict, rs, ss) if not return_queries else (rdict, sdict, rs, ss, new_reflexion_queries_dict)