| import os |
| from tqdm import tqdm |
| from loguru import logger |
| import json |
| from dataclasses import asdict |
| from agents.Reflexion import Reflexion |
| from utils.utils import extract_function_signatures, clear_code, extract_function_calls |
| from prompts import prompt_for_reflection |
| from memories.Memory import MemoryClassMeta |
| from models.Base import BaseModel |
| from retrievers.retriever import BM25Retriever |
| from prompts import prompt_for_generation |
| from concurrent.futures import ThreadPoolExecutor, as_completed |
|
|
|
|
|
|
| class Reflexion_Oneshot(Reflexion): |
|
|
| def __init__(self, model: BaseModel, dataset, corpus_path, mem_file=None, descendant_num=1): |
| self.model = model |
| self.dataset = dataset |
| self.memories = [] |
|
|
| self.instruction_retriever = BM25Retriever() |
| self.instruction_retriever.process(content_input_path=corpus_path) |
| self.code_retriever = BM25Retriever(mode="code") |
| self.code_retriever.process(content_input_path=corpus_path) |
|
|
| self.memory_init(mem_file, descendant_num) |
|
|
| def memory_init(self, mem_file=None, descendant_num=1): |
| class Memory(metaclass=MemoryClassMeta, field_names=["ps", |
| "err_msg", |
| "reflection", |
| "function_signatures", |
| "oneshot", |
| "pass_call", |
| ]): |
| pass |
| |
| if mem_file is not None: |
| assert mem_file.endswith(".json"), f"expect a json file, but got {mem_file} instead" |
| with open(mem_file, "r") as f: |
| input_mems = json.load(f) |
| assert len(input_mems) == len(self.dataset), f"expect {len(self.dataset)} samples, but got {len(input_mems)} instead" |
|
|
| for ps in self.dataset.problem_states: |
| if ps.label: |
| fs_mem = extract_function_signatures(ps.label) |
| else: |
| fs_mem = None |
| if mem_file is None: |
| os_mem = self.instruction_retriever.query(ps.instruction)[0] |
| tmp_mem = Memory(ps=ps, |
| err_msg=None, |
| reflection=None, |
| function_signatures=fs_mem, |
| oneshot=os_mem["code"], |
| pass_call=False, |
| ) |
| else: |
| input_mem = input_mems[ps.filename] |
| tmp_mem = Memory(ps=ps, |
| err_msg=input_mem["err_msg"], |
| reflection=input_mem["reflection"], |
| function_signatures=fs_mem, |
| oneshot=input_mem["oneshot"], |
| pass_call=input_mem["pass_call"], |
| ) |
| self.memories.append(tmp_mem) |
|
|
| def run(self, output_path=None, multi_thread=True, verbose=False, datalen=None, iteration_num=0, temperature=0): |
| data_len = datalen if datalen else len(self.dataset) |
| for iter in range(iteration_num): |
| logger.info(f"\n=== Iteration {iter} ===") |
| if output_path is not None: |
| root, extension = os.path.splitext(output_path) |
| iter_path = f"{root}_{iter}{extension}" |
|
|
| if multi_thread: |
| thread_num = 3 |
| |
| |
| logger.info(f"\ngenerate solution") |
| with tqdm(total=data_len) as pbar: |
| if multi_thread: |
| |
| with ThreadPoolExecutor(max_workers=thread_num) as executor: |
| futures = {executor.submit(self.generate_solution, mem, temperature): mem for mem in self.memories[:data_len]} |
| for future in as_completed(futures): |
| pbar.update(1) |
| else: |
| for mem in self.memories[:data_len]: |
| self.generate_solution(mem, temperature=temperature) |
| pbar.update(1) |
| |
| |
| logger.info(f"\nrun scripts on gpu") |
| for mem in tqdm(self.memories[:data_len]): |
| if mem.pass_call: |
| continue |
| is_pass, err_msg = self.dataset.run_single_call(mem.ps) |
| if not is_pass: |
| mem.err_msg = err_msg |
|
|
| |
| logger.info(f"\ngenerate reflections") |
| with tqdm(total=data_len) as pbar: |
| if multi_thread: |
| with ThreadPoolExecutor(max_workers=thread_num) as executor: |
| futures = {executor.submit(self.generate_reflexion, mem, temperature): mem for mem in self.memories[:data_len]} |
| for future in as_completed(futures): |
| pbar.update(1) |
| else: |
| for mem in self.memories[:data_len]: |
| self.generate_reflexion(mem, temperature=temperature) |
| pbar.update(1) |
| |
| if output_path is not None: |
| self.dataset.write_file(iter_path) |
| |
|
|
| |
| def generate_solution(self, mem, temperature=0): |
| if mem.pass_call: |
| return |
| |
| tab = "\n" |
| fss_text = "".join(f"* {sig}{tab}" for sig in mem.function_signatures) |
| text = prompt_for_generation.prompt.format( |
| instruction=mem.ps.instruction, |
| function_signatures=fss_text |
| ) |
|
|
| if not mem.ps.solution: |
| text += f"\nHere is an example snippet of code: {mem.oneshot}" |
| else: |
| one_shot = self.code_retriever.query(mem.ps.solution)[0]["code"] |
| text += f"\nHere is an example snippet of code: {one_shot}" |
| text += f"\nPrevious attempt implementation:{mem.ps.solution}" |
| |
| |
| if mem.err_msg: |
| text += f"\nTest messages for previous attempt:{mem.err_msg}" |
| |
| if mem.reflection: |
| text += f"\nReflection on previous attempt:{mem.reflection}" |
|
|
| text += "Please output the codes only without explanation, which we can run directly." |
| msg = [ |
| {"role": "user", "content": text}, |
| ] |
| response = self.model.generate(msg, temperature=temperature) |
| mem.ps.solution = clear_code(response) |
|
|
| return |
|
|
|
|
|
|
| def generate_reflexion(self, mem, temperature): |
| if mem.pass_call: |
| return |
| reflect_txt = prompt_for_reflection.prompt.format( |
| problem=mem.ps.instruction, |
| solution=mem.ps.solution, |
| test_result=mem.err_msg |
| ) |
| reflect_msg = [ |
| { |
| "role": "user", |
| "content": reflect_txt |
| } |
| ] |
| mem.reflection = self.model.generate(reflect_msg, temperature=temperature) |
