From e95ba070697ad7a16f4b3e2f428b100c1c546695 Mon Sep 17 00:00:00 2001 From: Jianlin Shi Date: Sat, 11 Jan 2025 00:10:01 -0700 Subject: [PATCH] support local llm server --- README.md | 5 +++ agents.py | 38 ++++++++--------- ai_lab_repo.py | 27 +++++++++---- inference.py | 43 +++++++++++++++----- mlesolver.py | 103 ++++++++++++++++++++++++++++------------------- papersolver.py | 14 ++++--- requirements.txt | 2 +- tools.py | 81 +++++++++++++++++++++++++++---------- 8 files changed, 207 insertions(+), 106 deletions(-) diff --git a/README.md b/README.md index e9758f9..e390954 100755 --- a/README.md +++ b/README.md @@ -78,6 +78,11 @@ To run Agent Laboratory in copilot mode, simply set the copilot-mode flag to `"t `python ai_lab_repo.py --api-key "API_KEY_HERE" --llm-backend "o1-mini" --research-topic "YOUR RESEARCH IDEA" --copilot-mode "true"` +### Run with local LLM server + +`python ai_lab_repo.py --api-key "ollama" --llm-backend "qwen2.5-coder" --research-topic "YOUR RESEARCH IDEA" --copilot-mode "true" --base-url http://localhost:11434/v1` +or +`python ai_lab_repo.py --api-key "llama.cpp" --llm-backend "qwen2.5-coder" --research-topic "YOUR RESEARCH IDEA" --copilot-mode "true" --base-url http://localhost:8080/v1` ----- ## Tips for better research outcomes diff --git a/agents.py b/agents.py index 1822ef3..e67f8a5 100755 --- a/agents.py +++ b/agents.py @@ -32,7 +32,7 @@ def extract_json_between_markers(llm_output): -def get_score(outlined_plan, latex, reward_model_llm, reviewer_type=None, attempts=3, openai_api_key=None): +def get_score(outlined_plan, latex, reward_model_llm, reviewer_type=None, attempts=3, openai_api_key=None, base_url=''): e = str() for _attempt in range(attempts): try: @@ -144,7 +144,7 @@ def get_score(outlined_plan, latex, reward_model_llm, reviewer_type=None, attemp openai_api_key=openai_api_key, prompt=( f"Outlined in the following text is the research plan that the machine learning engineer was tasked with building: {outlined_plan}\n\n" - f"The following text is the research latex that the model produced: \n{latex}\n\n"), temp=0.0) + f"The following text is the research latex that the model produced: \n{latex}\n\n"), temp=0.0, base_url=base_url) review_json = extract_json_between_markers(scoring) overall = int(review_json["Overall"]) / 10 @@ -181,27 +181,28 @@ def get_score(outlined_plan, latex, reward_model_llm, reviewer_type=None, attemp class ReviewersAgent: - def __init__(self, model="gpt-4o-mini", notes=None, openai_api_key=None): + def __init__(self, model="gpt-4o-mini", notes=None, openai_api_key=None, base_url=''): if notes is None: self.notes = [] else: self.notes = notes self.model = model self.openai_api_key = openai_api_key + self.base_url=base_url def inference(self, plan, report): reviewer_1 = "You are a harsh but fair reviewer and expect good experiments that lead to insights for the research topic." - review_1 = get_score(outlined_plan=plan, latex=report, reward_model_llm=self.model, reviewer_type=reviewer_1, openai_api_key=self.openai_api_key) + review_1 = get_score(outlined_plan=plan, latex=report, reward_model_llm=self.model, reviewer_type=reviewer_1, openai_api_key=self.openai_api_key, base_url=self.base) reviewer_2 = "You are a harsh and critical but fair fair reviewer who is looking for idea that would be impactful in the field." - review_2 = get_score(outlined_plan=plan, latex=report, reward_model_llm=self.model, reviewer_type=reviewer_2, openai_api_key=self.openai_api_key) + review_2 = get_score(outlined_plan=plan, latex=report, reward_model_llm=self.model, reviewer_type=reviewer_2, openai_api_key=self.openai_api_key, base_url=self.base) reviewer_3 = "You are a harsh but fair open-minded reviewer that is looking for novel ideas that have not been proposed before." - review_3 = get_score(outlined_plan=plan, latex=report, reward_model_llm=self.model, reviewer_type=reviewer_3, openai_api_key=self.openai_api_key) + review_3 = get_score(outlined_plan=plan, latex=report, reward_model_llm=self.model, reviewer_type=reviewer_3, openai_api_key=self.openai_api_key, base_url=self.base) return f"Reviewer #1:\n{review_1}, \nReviewer #2:\n{review_2}, \nReviewer #3:\n{review_3}" class BaseAgent: - def __init__(self, model="gpt-4o-mini", notes=None, max_steps=100, openai_api_key=None): + def __init__(self, model="gpt-4o-mini", notes=None, max_steps=100, openai_api_key=None, base_url=''): if notes is None: self.notes = [] else: self.notes = notes self.max_steps = max_steps @@ -222,6 +223,7 @@ def __init__(self, model="gpt-4o-mini", notes=None, max_steps=100, openai_api_ke self.prev_results_code = str() self.prev_interpretation = str() self.openai_api_key = openai_api_key + self.base_url=base_url self.second_round = False self.max_hist_len = 15 @@ -251,7 +253,7 @@ def inference(self, research_topic, phase, step, feedback="", temp=None): f"Current Step #{step}, Phase: {phase}\n{complete_str}\n" f"[Objective] Your goal is to perform research on the following topic: {research_topic}\n" f"Feedback: {feedback}\nNotes: {notes_str}\nYour previous command was: {self.prev_comm}. Make sure your new output is very different.\nPlease produce a single command below:\n") - model_resp = query_model(model_str=self.model, system_prompt=sys_prompt, prompt=prompt, temp=temp, openai_api_key=self.openai_api_key) + model_resp = query_model(model_str=self.model, system_prompt=sys_prompt, prompt=prompt, temp=temp, openai_api_key=self.openai_api_key, base_url=self.base_url) print("^"*50, phase, "^"*50) model_resp = self.clean_text(model_resp) self.prev_comm = model_resp @@ -291,8 +293,8 @@ def example_command(self, phase): class ProfessorAgent(BaseAgent): - def __init__(self, model="gpt4omini", notes=None, max_steps=100, openai_api_key=None): - super().__init__(model, notes, max_steps, openai_api_key) + def __init__(self, model="gpt4omini", notes=None, max_steps=100, openai_api_key=None, base_url=''): + super().__init__(model, notes, max_steps, openai_api_key, base_url) self.phases = ["report writing"] def generate_readme(self): @@ -301,7 +303,7 @@ def generate_readme(self): prompt = ( f"""History: {history_str}\n{'~' * 10}\n""" f"Please produce the readme below in markdown:\n") - model_resp = query_model(model_str=self.model, system_prompt=sys_prompt, prompt=prompt, openai_api_key=self.openai_api_key) + model_resp = query_model(model_str=self.model, system_prompt=sys_prompt, prompt=prompt, openai_api_key=self.openai_api_key, base_url=self.base_url) return model_resp.replace("```markdown", "") def context(self, phase): @@ -357,8 +359,8 @@ def role_description(self): class PostdocAgent(BaseAgent): - def __init__(self, model="gpt4omini", notes=None, max_steps=100, openai_api_key=None): - super().__init__(model, notes, max_steps, openai_api_key) + def __init__(self, model="gpt4omini", notes=None, max_steps=100, openai_api_key=None, base_url=''): + super().__init__(model, notes, max_steps, openai_api_key, base_url) self.phases = ["plan formulation", "results interpretation"] def context(self, phase): @@ -433,8 +435,8 @@ def role_description(self): class MLEngineerAgent(BaseAgent): - def __init__(self, model="gpt4omini", notes=None, max_steps=100, openai_api_key=None): - super().__init__(model, notes, max_steps, openai_api_key) + def __init__(self, model="gpt4omini", notes=None, max_steps=100, openai_api_key=None, base_url=''): + super().__init__(model, notes, max_steps, openai_api_key, base_url) self.phases = [ "data preparation", "running experiments", @@ -498,8 +500,8 @@ def role_description(self): class PhDStudentAgent(BaseAgent): - def __init__(self, model="gpt4omini", notes=None, max_steps=100, openai_api_key=None): - super().__init__(model, notes, max_steps, openai_api_key) + def __init__(self, model="gpt4omini", notes=None, max_steps=100, openai_api_key=None, base_url=''): + super().__init__(model, notes, max_steps, openai_api_key,base_url) self.phases = [ "literature review", "plan formulation", @@ -579,7 +581,7 @@ def requirements_txt(self): prompt = ( f"""History: {history_str}\n{'~' * 10}\n""" f"Please produce the requirements.txt below in markdown:\n") - model_resp = query_model(model_str=self.model, system_prompt=sys_prompt, prompt=prompt, openai_api_key=self.openai_api_key) + model_resp = query_model(model_str=self.model, system_prompt=sys_prompt, prompt=prompt, openai_api_key=self.openai_api_key, base_url=self.base_url) return model_resp def example_command(self, phase): diff --git a/ai_lab_repo.py b/ai_lab_repo.py index 58f10cd..8c8805a 100755 --- a/ai_lab_repo.py +++ b/ai_lab_repo.py @@ -11,7 +11,8 @@ class LaboratoryWorkflow: - def __init__(self, research_topic, openai_api_key, max_steps=100, num_papers_lit_review=5, agent_model_backbone=f"{DEFAULT_LLM_BACKBONE}", notes=list(), human_in_loop_flag=None, compile_pdf=True, mlesolver_max_steps=3, papersolver_max_steps=5): + def __init__(self, research_topic, openai_api_key, max_steps=100, num_papers_lit_review=5, agent_model_backbone=f"{DEFAULT_LLM_BACKBONE}", + notes=list(), human_in_loop_flag=None, compile_pdf=True, mlesolver_max_steps=3, papersolver_max_steps=5, base_url=''): """ Initialize laboratory workflow @param research_topic: (str) description of research idea to explore @@ -25,6 +26,7 @@ def __init__(self, research_topic, openai_api_key, max_steps=100, num_papers_lit self.max_steps = max_steps self.compile_pdf = compile_pdf self.openai_api_key = openai_api_key + self.base_url=base_url self.research_topic = research_topic self.model_backbone = agent_model_backbone self.num_papers_lit_review = num_papers_lit_review @@ -79,11 +81,11 @@ def __init__(self, research_topic, openai_api_key, max_steps=100, num_papers_lit self.save = True self.verbose = True - self.reviewers = ReviewersAgent(model=self.model_backbone, notes=self.notes, openai_api_key=self.openai_api_key) - self.phd = PhDStudentAgent(model=self.model_backbone, notes=self.notes, max_steps=self.max_steps, openai_api_key=self.openai_api_key) - self.postdoc = PostdocAgent(model=self.model_backbone, notes=self.notes, max_steps=self.max_steps, openai_api_key=self.openai_api_key) - self.professor = ProfessorAgent(model=self.model_backbone, notes=self.notes, max_steps=self.max_steps, openai_api_key=self.openai_api_key) - self.ml_engineer = MLEngineerAgent(model=self.model_backbone, notes=self.notes, max_steps=self.max_steps, openai_api_key=self.openai_api_key) + self.reviewers = ReviewersAgent(model=self.model_backbone, notes=self.notes, openai_api_key=self.openai_api_key, base_url=self.base_url) + self.phd = PhDStudentAgent(model=self.model_backbone, notes=self.notes, max_steps=self.max_steps, openai_api_key=self.openai_api_key, base_url=self.base_url) + self.postdoc = PostdocAgent(model=self.model_backbone, notes=self.notes, max_steps=self.max_steps, openai_api_key=self.openai_api_key, base_url=self.base_url) + self.professor = ProfessorAgent(model=self.model_backbone, notes=self.notes, max_steps=self.max_steps, openai_api_key=self.openai_api_key, base_url=self.base_url) + self.ml_engineer = MLEngineerAgent(model=self.model_backbone, notes=self.notes, max_steps=self.max_steps, openai_api_key=self.openai_api_key, base_url=self.base_url) # remove previous files remove_figures() @@ -240,7 +242,8 @@ def report_writing(self): # instantiate mle-solver from papersolver import PaperSolver self.reference_papers = [] - solver = PaperSolver(notes=report_notes, max_steps=self.papersolver_max_steps, plan=lab.phd.plan, exp_code=lab.phd.results_code, exp_results=lab.phd.exp_results, insights=lab.phd.interpretation, lit_review=lab.phd.lit_review, ref_papers=self.reference_papers, topic=research_topic, openai_api_key=self.openai_api_key, llm_str=self.model_backbone["report writing"], compile_pdf=compile_pdf) + solver = PaperSolver(notes=report_notes, max_steps=self.papersolver_max_steps, plan=lab.phd.plan, exp_code=lab.phd.results_code, exp_results=lab.phd.exp_results, insights=lab.phd.interpretation, lit_review=lab.phd.lit_review, ref_papers=self.reference_papers, topic=research_topic, openai_api_key=self.openai_api_key, + llm_str=self.model_backbone["report writing"], compile_pdf=compile_pdf, base_url=self.base_url) # run initialization for solver solver.initial_solve() # run solver for N mle optimization steps @@ -603,17 +606,24 @@ def parse_arguments(): help='Total number of paper-solver steps' ) + parser.add_argument( + '--base-url', + type=str, + default="", + help='Set to different url if you are using a custom server, e.g. http://localhost:11434 if using ollama, or http://localhost:8080/v1 if using llama.cpp.' + ) + return parser.parse_args() if __name__ == "__main__": args = parse_arguments() - llm_backend = args.llm_backend human_mode = args.copilot_mode.lower() == "true" compile_pdf = args.compile_latex.lower() == "true" load_existing = args.load_existing.lower() == "true" + base_url=args.base_url try: num_papers_lit_review = int(args.num_papers_lit_review.lower()) except Exception: @@ -711,6 +721,7 @@ def parse_arguments(): num_papers_lit_review=num_papers_lit_review, papersolver_max_steps=papersolver_max_steps, mlesolver_max_steps=mlesolver_max_steps, + base_url=base_url ) lab.perform_research() diff --git a/inference.py b/inference.py index 4fec787..6d9b9c6 100755 --- a/inference.py +++ b/inference.py @@ -8,6 +8,7 @@ encoding = tiktoken.get_encoding("cl100k_base") + def curr_cost_est(): costmap_in = { "gpt-4o": 2.50 / 1000000, @@ -17,15 +18,18 @@ def curr_cost_est(): "claude-3-5-sonnet": 3.00 / 1000000, } costmap_out = { - "gpt-4o": 10.00/ 1000000, + "gpt-4o": 10.00 / 1000000, "gpt-4o-mini": 0.6 / 1000000, "o1-preview": 60.00 / 1000000, "o1-mini": 12.00 / 1000000, "claude-3-5-sonnet": 12.00 / 1000000, } - return sum([costmap_in[_]*TOKENS_IN[_] for _ in TOKENS_IN]) + sum([costmap_out[_]*TOKENS_OUT[_] for _ in TOKENS_OUT]) + return sum([costmap_in[_] * TOKENS_IN[_] for _ in TOKENS_IN]) + sum( + [costmap_out[_] * TOKENS_OUT[_] for _ in TOKENS_OUT]) + -def query_model(model_str, prompt, system_prompt, openai_api_key=None, anthropic_api_key=None, tries=5, timeout=5.0, temp=None, print_cost=True, version="1.5"): +def query_model(model_str, prompt, system_prompt, openai_api_key=None, anthropic_api_key=None, tries=5, timeout=5.0, + temp=None, print_cost=True, version="1.5", base_url=''): preloaded_api = os.getenv('OPENAI_API_KEY') if openai_api_key is None and preloaded_api is not None: openai_api_key = preloaded_api @@ -36,6 +40,7 @@ def query_model(model_str, prompt, system_prompt, openai_api_key=None, anthropic os.environ["OPENAI_API_KEY"] = openai_api_key if anthropic_api_key is not None: os.environ["ANTHROPIC_API_KEY"] = anthropic_api_key + encoding=None for _ in range(tries): try: if model_str == "gpt-4o-mini" or model_str == "gpt4omini" or model_str == "gpt-4omini" or model_str == "gpt4o-mini": @@ -124,17 +129,36 @@ def query_model(model_str, prompt, system_prompt, openai_api_key=None, anthropic completion = client.chat.completions.create( model="o1-preview", messages=messages) answer = completion.choices[0].message.content - - if model_str in ["o1-preview", "o1-mini", "claude-3.5-sonnet"]: - encoding = tiktoken.encoding_for_model("gpt-4o") - else: encoding = tiktoken.encoding_for_model(model_str) + else: + messages = [ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": prompt}] + client = OpenAI() + if base_url is not None and base_url != '': + client.base_url = base_url + completion = client.chat.completions.create( + model=f"{model_str}", messages=messages) + answer = completion.choices[0].message.content + try: + if model_str in ["o1-preview", "o1-mini", "claude-3.5-sonnet"]: + encoding = tiktoken.encoding_for_model("gpt-4o") + else: + encoding = tiktoken.encoding_for_model(model_str) + except Exception as e: + pass + if encoding is None: + # set a default encoding to by pass the error for now. + encoding = tiktoken.encoding_for_model('gpt-4') if model_str not in TOKENS_IN: TOKENS_IN[model_str] = 0 TOKENS_OUT[model_str] = 0 TOKENS_IN[model_str] += len(encoding.encode(system_prompt + prompt)) TOKENS_OUT[model_str] += len(encoding.encode(answer)) if print_cost: - print(f"Current experiment cost = ${curr_cost_est()}, ** Approximate values, may not reflect true cost") + try: + print(f"Current experiment cost = ${curr_cost_est()}, ** Approximate values, may not reflect true cost") + except Exception as e: + pass return answer except Exception as e: print("Inference Exception:", e) @@ -142,5 +166,4 @@ def query_model(model_str, prompt, system_prompt, openai_api_key=None, anthropic continue raise Exception("Max retries: timeout") - -#print(query_model(model_str="o1-mini", prompt="hi", system_prompt="hey")) \ No newline at end of file +# print(query_model(model_str="o1-mini", prompt="hi", system_prompt="hey")) diff --git a/mlesolver.py b/mlesolver.py index cfc4896..d357175 100755 --- a/mlesolver.py +++ b/mlesolver.py @@ -4,15 +4,14 @@ from common_imports import * from abc import abstractmethod - from tools import * from inference import * from pathlib import Path - from contextlib import contextmanager import sys, os + @contextmanager def suppress_stdout(): with open(os.devnull, "w") as devnull: @@ -29,8 +28,8 @@ def suppress_stdout(): warnings.filterwarnings("ignore") warnings.simplefilter(action='ignore', category=FutureWarning) import logging -logging.getLogger('sklearn.model_selection').setLevel(logging.WARNING) +logging.getLogger('sklearn.model_selection').setLevel(logging.WARNING) GLOBAL_REPAIR_ATTEMPTS = 2 @@ -62,6 +61,7 @@ def parse_command(self, cmd_str) -> tuple: @@@@@@@@@@@@@@@@@@ """ + class Replace(Command): def __init__(self): super().__init__() @@ -92,7 +92,6 @@ def parse_command(self, *args) -> tuple: return True, (new_code.split("\n"), code_ret) - class Edit(Command): def __init__(self): super().__init__() @@ -116,7 +115,7 @@ def execute_command(self, *args) -> str: args = args[0] current_code = args[2] lines_to_add = list(reversed(args[3])) - lines_to_replace = list(reversed(range(args[0], args[1]+1))) + lines_to_replace = list(reversed(range(args[0], args[1] + 1))) for _ln in lines_to_replace: current_code.pop(_ln) for _line in lines_to_add: @@ -148,7 +147,7 @@ def parse_command(self, *args) -> tuple: return False, (None, None, None, None, None) -def get_score(outlined_plan, code, code_return, REWARD_MODEL_LLM, attempts=3, openai_api_key=None): +def get_score(outlined_plan, code, code_return, REWARD_MODEL_LLM, attempts=3, openai_api_key=None, base_url=''): e = str() for _attempt in range(attempts): try: @@ -164,7 +163,7 @@ def get_score(outlined_plan, code, code_return, REWARD_MODEL_LLM, attempts=3, op prompt=( f"Outlined in the following text is the research plan that the machine learning engineer was tasked with building: {outlined_plan}\n\n" f"The following text is the research code that the model produced: \n{code}\n\n" - f"The following is the output from the model: {code_return}\n\n"), temp=0.6) + f"The following is the output from the model: {code_return}\n\n"), temp=0.6, base_url=base_url) performance = extract_prompt(text=scoring, word="SCORE") performance = float(performance) return performance, f"The performance of your submission is: {performance}", True @@ -173,7 +172,7 @@ def get_score(outlined_plan, code, code_return, REWARD_MODEL_LLM, attempts=3, op return 0, e -def code_repair(code, error, ctype, REPAIR_LLM, openai_api_key=None): +def code_repair(code, error, ctype, REPAIR_LLM, openai_api_key=None, base_url=''): if ctype == "replace": repair_sys = ( "You are an automated code repair tool.\n" @@ -186,14 +185,14 @@ def code_repair(code, error, ctype, REPAIR_LLM, openai_api_key=None): openai_api_key=openai_api_key, model_str=f"{REPAIR_LLM}", system_prompt=repair_sys, - prompt=f"Provided here is the error: {error}\n\nProvided below is the code:\n\n{code}", temp=0.8) + prompt=f"Provided here is the error: {error}\n\nProvided below is the code:\n\n{code}", temp=0.8, base_url=base_url) return extract_prompt(model_resp, "python") elif ctype == "edit": repair_sys = ( "You are an automated code repair tool.\n" "Your goal is to take in code and an error and repair the code to make sure the same error does not repeat itself, and also to remove any other potential errors from the code without affecting the code output.\n" "Your output should match the original code as closely as possible.\n" - + "============= CODE EDITING TOOL =============\n" "You have access to a code editing tool. \n" "This tool allows you to replace lines indexed n through m (n:m) of the current code with as many lines of new code as you want to add. This removal is inclusive meaning that line n and m and everything between n and m is removed. This will be the primary way that you interact with code. \n" @@ -206,17 +205,22 @@ def code_repair(code, error, ctype, REPAIR_LLM, openai_api_key=None): openai_api_key=openai_api_key, model_str=f"{REPAIR_LLM}", system_prompt=repair_sys, - prompt=f"Provided here is the error: {error}\n\nProvided below is the code:\n\n{code}", temp=0.2) + prompt=f"Provided here is the error: {error}\n\nProvided below is the code:\n\n{code}", temp=0.2, base_url=base_url) return model_resp class MLESolver: - def __init__(self, dataset_code, openai_api_key=None, notes=None, max_steps=10, insights=None, plan=None, llm_str=None): - if notes is None: self.notes = [] - else: self.notes = notes + def __init__(self, dataset_code, openai_api_key=None, notes=None, max_steps=10, insights=None, plan=None, + llm_str=None, base_url=''): + if notes is None: + self.notes = [] + else: + self.notes = notes self.dataset_code = dataset_code - if plan is None: self.plan = "" - else: self.plan = plan + if plan is None: + self.plan = "" + else: + self.plan = plan self.llm_str = llm_str self.verbose = False self.max_codes = 2 @@ -230,6 +234,7 @@ def __init__(self, dataset_code, openai_api_key=None, notes=None, max_steps=10, self.prev_code_ret = str() self.should_execute_code = True self.openai_api_key = openai_api_key + self.base_url = base_url def initial_solve(self): """ @@ -273,7 +278,8 @@ def gen_initial_code(self): openai_api_key=self.openai_api_key, model_str=self.model, system_prompt=self.system_prompt(), - prompt=f"{err_hist}\nYou should now use ```REPLACE to create initial code to solve the challenge. Now please enter the ```REPLACE command below:\n ", temp=1.0) + prompt=f"{err_hist}\nYou should now use ```REPLACE to create initial code to solve the challenge. Now please enter the ```REPLACE command below:\n ", + temp=1.0, base_url=self.base_url) model_resp = self.clean_text(model_resp) cmd_str, code_lines, prev_code_ret, should_execute_code, score = self.process_command(model_resp) print(f"@@@ INIT ATTEMPT: Command Exec // Attempt {num_attempts}: ", str(cmd_str).replace("\n", " | ")) @@ -289,13 +295,16 @@ def solve(self): self.prev_code_ret = None self.should_execute_code = False while True: - if len(self.commands) == 2: cmd_app_str = "You must output either the ```EDIT or ```REPLACE command immediately. " - else: cmd_app_str = "" + if len(self.commands) == 2: + cmd_app_str = "You must output either the ```EDIT or ```REPLACE command immediately. " + else: + cmd_app_str = "" model_resp = query_model( openai_api_key=self.openai_api_key, model_str=self.model, system_prompt=self.system_prompt(), - prompt=f"The following is your history:{self.history_str()}\n\n{cmd_app_str}Now please enter a command: ", temp=1.0) + prompt=f"The following is your history:{self.history_str()}\n\n{cmd_app_str}Now please enter a command: ", + temp=1.0, base_url=self.base_url) model_resp = self.clean_text(model_resp) self.code_lines = copy(random.choice(self.best_codes)[0]) cmd_str, code_lines, prev_code_ret, should_execute_code, score = self.process_command(model_resp) @@ -303,10 +312,12 @@ def solve(self): if len(self.st_history) > self.st_hist_len: self.st_history.pop(0) if score is not None: if top_score is None: - best_pkg = copy(code_lines), copy(prev_code_ret), copy(should_execute_code), copy(model_resp), copy(cmd_str) + best_pkg = copy(code_lines), copy(prev_code_ret), copy(should_execute_code), copy(model_resp), copy( + cmd_str) top_score = score elif score > top_score: - best_pkg = copy(code_lines), copy(prev_code_ret), copy(should_execute_code), copy(model_resp), copy(cmd_str) + best_pkg = copy(code_lines), copy(prev_code_ret), copy(should_execute_code), copy(model_resp), copy( + cmd_str) top_score = score print(f"@@@ Command Exec // Attempt {num_attempts}: ", str(cmd_str).replace("\n", " | ")) print(f"$$$ Score: {score}") @@ -329,10 +340,13 @@ def reflect_code(self): Provide a reflection on produced behavior for next execution @return: (str) language model-produced reflection """ - code_strs = ("$"*40 + "\n\n").join([self.generate_code_lines(_code[0]) + f"\nCode Return {_code[1]}" for _code in self.best_codes]) + code_strs = ("$" * 40 + "\n\n").join( + [self.generate_code_lines(_code[0]) + f"\nCode Return {_code[1]}" for _code in self.best_codes]) code_strs = f"Please reflect on the following sets of code: {code_strs} and come up with generalizable insights that will help you improve your performance on this benchmark." syst = self.system_prompt(commands=False) + code_strs - return query_model(prompt="Please reflect on ideas for how to improve your current code. Examine the provided code and think very specifically (with precise ideas) on how to improve performance, which methods to use, how to improve generalization on the test set with line-by-line examples below:\n", system_prompt=syst, model_str=f"{self.llm_str}", openai_api_key=self.openai_api_key) + return query_model( + prompt="Please reflect on ideas for how to improve your current code. Examine the provided code and think very specifically (with precise ideas) on how to improve performance, which methods to use, how to improve generalization on the test set with line-by-line examples below:\n", + system_prompt=syst, model_str=f"{self.llm_str}", openai_api_key=self.openai_api_key, base_url=self.base_url) def process_command(self, model_resp): """ @@ -349,7 +363,7 @@ def process_command(self, model_resp): should_execute_code = self.should_execute_code code_lines = copy(self.code_lines) remove_figures() - with suppress_stdout(): # shhh + with suppress_stdout(): # shhh for cmd in self.commands: if cmd.matches_command(model_resp): # attempt to execute the code edit command @@ -364,12 +378,16 @@ def process_command(self, model_resp): code_err = f"Return from executing code: {cmd_return[2]}" if cmd_return[0]: # if success code_lines = copy(cmd_return[1]) - score, cmd_str, is_valid = get_score(self.plan, "\n".join(code_lines), cmd_return[2], openai_api_key=self.openai_api_key, REWARD_MODEL_LLM=self.llm_str) + score, cmd_str, is_valid = get_score(self.plan, "\n".join(code_lines), + cmd_return[2], + openai_api_key=self.openai_api_key, + REWARD_MODEL_LLM=self.llm_str) if is_valid: failed = False break code_err += f"\nReturn from executing code on real test set {cmd_str}" - repaired_code = code_repair(model_resp, code_err, REPAIR_LLM=self.llm_str, ctype="edit", openai_api_key=self.openai_api_key) + repaired_code = code_repair(model_resp, code_err, REPAIR_LLM=self.llm_str, ctype="edit", + openai_api_key=self.openai_api_key, base_url=self.base_url) model_resp = repaired_code print(f" * Attempting repair // try {_tries}*") if failed: @@ -382,7 +400,7 @@ def process_command(self, model_resp): should_execute_code = True return cmd_str, code_lines, prev_code_ret, should_execute_code, score # attempt to execute the code replace command - elif cmd.cmd_type == "CODE-replace": # DONE + elif cmd.cmd_type == "CODE-replace": # DONE score = None failed = True code_err = str() @@ -391,12 +409,16 @@ def process_command(self, model_resp): code_err = f"Return from executing code: {args[1]}" if success: code_lines = copy(args[0]) - score, cmd_str, is_valid = get_score(self.plan, "\n".join(code_lines), args[1], openai_api_key=self.openai_api_key, REWARD_MODEL_LLM=self.llm_str) + score, cmd_str, is_valid = get_score(self.plan, "\n".join(code_lines), args[1], + openai_api_key=self.openai_api_key, + REWARD_MODEL_LLM=self.llm_str) if is_valid: failed = False break code_err += f"\nReturn from executing code on real test set {cmd_str}" - repaired_code = code_repair(extract_prompt(model_resp, "REPLACE", ), code_err, ctype="replace", openai_api_key=self.openai_api_key, REPAIR_LLM=self.llm_str) + repaired_code = code_repair(extract_prompt(model_resp, "REPLACE", ), code_err, + ctype="replace", openai_api_key=self.openai_api_key, + REPAIR_LLM=self.llm_str, base_url=self.base_url) repaired_code = f"```REPLACE\n{repaired_code}\n```" model_resp = repaired_code print(f" * Attempting repair // try {_tries}*") @@ -420,12 +442,13 @@ def history_str(self): """ hist_str = "" for _hist in range(len(self.st_history)): - hist_str += f"-------- History ({len(self.st_history)-_hist} steps ago) -----\n" - hist_str += f"Because of the following response: {self.st_history[_hist][0]}\n" if len(self.st_history[_hist][0]) > 0 else "" + hist_str += f"-------- History ({len(self.st_history) - _hist} steps ago) -----\n" + hist_str += f"Because of the following response: {self.st_history[_hist][0]}\n" if len( + self.st_history[_hist][0]) > 0 else "" hist_str += f"and the following COMMAND response output: {self.st_history[_hist][3]}\n" - hist_str += f"With the following code used: {'#'*20}\n{self.st_history[_hist][2]}\n{'#'*20}\n\n" + hist_str += f"With the following code used: {'#' * 20}\n{self.st_history[_hist][2]}\n{'#' * 20}\n\n" hist_str += f"The environment feedback and reflection was as follows: {self.st_history[_hist][1]}\n" - hist_str += f"-------- End of history ({len(self.st_history)-_hist} steps ago) -------\n" + hist_str += f"-------- End of history ({len(self.st_history) - _hist} steps ago) -------\n" return hist_str def system_prompt(self, commands=True): @@ -483,7 +506,8 @@ def feedback(self, code_return): reflect_prompt = f"This is your code: {code_str}\n\nYour code returned the following error {code_return}. Please provide a detailed reflection on why this error was returned, which lines in the code caused this error, and exactly (line by line) how you hope to fix this in the next update. This step is mostly meant to reflect in order to help your future self fix the error better. Do not provide entirely new code but provide suggestions on how to fix the bug using LINE EDITS." elif os.path.exists("submission.csv"): self.prev_working_code = copy(self.code_lines) - grade_return = get_score(self.plan, "\n".join(self.prev_working_code), code_return, openai_api_key=self.openai_api_key)[0] + grade_return = get_score(self.plan, "\n".join(self.prev_working_code), code_return, + openai_api_key=self.openai_api_key)[0] print(f"@@@@ SUBMISSION: model score {grade_return}", REWARD_MODEL_LLM=self.llm_str) f"Your code was properly submitted and you have just received a grade for your model.\nYour score was {grade_return}.\n\n" reflect_prompt = f"This is your code: {code_str}\n\nYour code successfully returned a submission csv. Consider further improving your technique through advanced learning techniques, data augmentation, or hyperparamter tuning to increase the score. Please provide a detailed reflection on how to improve your performance, which lines in the code could be improved upon, and exactly (line by line) how you hope to improve this in the next update. This step is mostly meant to reflect in order to help your future self." @@ -507,7 +531,8 @@ def reflection(self, reflect_prompt, code_str, code_return): @param code_str: (str) code string @return: (str) reflection string """ - refl = query_model(prompt=reflect_prompt, system_prompt=self.system_prompt(commands=False), model_str=f"{self.llm_str}", openai_api_key=self.openai_api_key) + refl = query_model(prompt=reflect_prompt, system_prompt=self.system_prompt(commands=False), + model_str=f"{self.llm_str}", openai_api_key=self.openai_api_key, base_url=self.base_url) return f"During the previous execution, the following code was run: \n\n{code_str}\n\nThis code returned the following: \n{code_return}\nThe following is your reflection from this feedback {refl}\n" def generate_dataset_descr_prompt(self): @@ -518,7 +543,7 @@ def generate_dataset_descr_prompt(self): """ return f"\n- The following dataset code will be added to the beginning of your code always, so this does not need to be rewritten: {self.dataset_code}" - def phase_prompt(self,): + def phase_prompt(self, ): """ Describe system role and general tips for mle-solver @return: (str) system role @@ -569,7 +594,3 @@ def run_code(self): elif self.should_execute_code: return execute_code("\n".join(self.code_lines)) return "Changes have not yet been made to the code." - - - - diff --git a/papersolver.py b/papersolver.py index 18e7a95..9f682de 100755 --- a/papersolver.py +++ b/papersolver.py @@ -243,7 +243,8 @@ def parse_command(self, *args) -> tuple: } class PaperSolver: - def __init__(self, llm_str, notes=None, max_steps=10, insights=None, plan=None, exp_code=None, exp_results=None, lit_review=None, ref_papers=None, topic=None, openai_api_key=None, compile_pdf=True): + def __init__(self, llm_str, notes=None, max_steps=10, insights=None, plan=None, exp_code=None, exp_results=None, lit_review=None, ref_papers=None, topic=None, + openai_api_key=None, compile_pdf=True, base_url=''): if notes is None: self.notes = [] else: self.notes = notes if plan is None: self.plan = "" @@ -271,6 +272,7 @@ def __init__(self, llm_str, notes=None, max_steps=10, insights=None, plan=None, self.prev_paper_ret = str() self.section_related_work = {} self.openai_api_key = openai_api_key + self.base_url=base_url def solve(self): num_attempts = 0 @@ -284,7 +286,7 @@ def solve(self): system_prompt=self.system_prompt(), prompt=f"\nNow please enter a command: ", temp=1.0, - openai_api_key=self.openai_api_key) + openai_api_key=self.openai_api_key, base_url=self.base_url) #print(model_resp) model_resp = self.clean_text(model_resp) cmd_str, paper_lines, prev_paper_ret, score = self.process_command(model_resp) @@ -351,7 +353,7 @@ def gen_initial_report(self): break if not first_attempt: att_str = "This is not your first attempt please try to come up with a simpler search query." - search_query = query_model(model_str=f"{self.llm_str}", prompt=f"Given the following research topic {self.topic} and research plan: \n\n{self.plan}\n\nPlease come up with a search query to find relevant papers on arXiv. Respond only with the search query and nothing else. This should be a a string that will be used to find papers with semantically similar content. {att_str}", system_prompt=f"You are a research paper finder. You must find papers for the section {_section}. Query must be text nothing else.", openai_api_key=self.openai_api_key) + search_query = query_model(model_str=f"{self.llm_str}", prompt=f"Given the following research topic {self.topic} and research plan: \n\n{self.plan}\n\nPlease come up with a search query to find relevant papers on arXiv. Respond only with the search query and nothing else. This should be a a string that will be used to find papers with semantically similar content. {att_str}", system_prompt=f"You are a research paper finder. You must find papers for the section {_section}. Query must be text nothing else.", openai_api_key=self.openai_api_key, base_url=self.base_url) search_query.replace('"', '') papers = arx.find_papers_by_str(query=search_query, N=10) first_attempt = False @@ -374,7 +376,7 @@ def gen_initial_report(self): system_prompt=self.system_prompt(section=_section), prompt=f"{prompt}", temp=0.8, - openai_api_key=self.openai_api_key) + openai_api_key=self.openai_api_key, base_url=self.base_url) model_resp = self.clean_text(model_resp) if _section == "scaffold": # minimal scaffold (some other sections can be combined) @@ -437,7 +439,7 @@ def process_command(self, model_resp, scoring=True): else: paper_lines = copy(args[1]) # if scoring: - score, cmd_str, is_valid = get_score(self.plan, "\n".join(paper_lines), reward_model_llm=self.llm_str) + score, cmd_str, is_valid = get_score(self.plan, "\n".join(paper_lines), reward_model_llm=self.llm_str, base_url=self.base_url) else: score, cmd_str, is_valid = 0.0, "Paper scored successfully", True if is_valid: failed = False @@ -459,7 +461,7 @@ def process_command(self, model_resp, scoring=True): if success: paper_lines = copy(args[0]) # if scoring: - score, cmd_str, is_valid = get_score(self.plan, "\n".join(paper_lines), reward_model_llm=self.llm_str) + score, cmd_str, is_valid = get_score(self.plan, "\n".join(paper_lines), reward_model_llm=self.llm_str, base_url=self.base_url) else: score, cmd_str, is_valid = 0.0, "Paper scored successfully", True if is_valid: failed = False diff --git a/requirements.txt b/requirements.txt index 55ae922..de6adff 100755 --- a/requirements.txt +++ b/requirements.txt @@ -66,7 +66,7 @@ murmurhash==1.0.11 namex==0.0.8 nest-asyncio==1.6.0 networkx==3.2.1 -NEURON==8.2.0 +NEURON nltk==3.9.1 numpy==2.0.2 openai==1.55.1 diff --git a/tools.py b/tools.py index fe4413a..8bcaa3a 100755 --- a/tools.py +++ b/tools.py @@ -202,25 +202,64 @@ def __init__(self): # Construct the default API client. self.sch_engine = arxiv.Client() + def _process_query(self, query: str) -> str: + """Process query string to fit within MAX_QUERY_LENGTH while preserving as much information as possible""" + MAX_QUERY_LENGTH = 300 + + if len(query) <= MAX_QUERY_LENGTH: + return query + + # Split into words + words = query.split() + processed_query = [] + current_length = 0 + + # Add words while staying under the limit + # Account for spaces between words + for word in words: + # +1 for the space that will be added between words + if current_length + len(word) + 1 <= MAX_QUERY_LENGTH: + processed_query.append(word) + current_length += len(word) + 1 + else: + break + + return ' '.join(processed_query) + def find_papers_by_str(self, query, N=20): - search = arxiv.Search( - query="abs:" + query, - max_results=N, - sort_by=arxiv.SortCriterion.Relevance) + processed_query = self._process_query(query) + max_retries = 3 + retry_count = 0 - paper_sums = list() - # `results` is a generator; you can iterate over its elements one by one... - for r in self.sch_engine.results(search): - paperid = r.pdf_url.split("/")[-1] - pubdate = str(r.published).split(" ")[0] - paper_sum = f"Title: {r.title}\n" - paper_sum += f"Summary: {r.summary}\n" - paper_sum += f"Publication Date: {pubdate}\n" - paper_sum += f"Categories: {' '.join(r.categories)}\n" - paper_sum += f"arXiv paper ID: {paperid}\n" - paper_sums.append(paper_sum) - time.sleep(2.0) - return "\n".join(paper_sums) + while retry_count < max_retries: + try: + search = arxiv.Search( + query="abs:" + processed_query, + max_results=N, + sort_by=arxiv.SortCriterion.Relevance) + + paper_sums = list() + # `results` is a generator; you can iterate over its elements one by one... + for r in self.sch_engine.results(search): + paperid = r.pdf_url.split("/")[-1] + pubdate = str(r.published).split(" ")[0] + paper_sum = f"Title: {r.title}\n" + paper_sum += f"Summary: {r.summary}\n" + paper_sum += f"Publication Date: {pubdate}\n" + paper_sum += f"Categories: {' '.join(r.categories)}\n" + paper_sum += f"arXiv paper ID: {paperid}\n" + paper_sums.append(paper_sum) + time.sleep(2.0) + return "\n".join(paper_sums) + + except Exception as e: + retry_count += 1 + if retry_count < max_retries: + # 递增延时 + time.sleep(2 * retry_count) + continue + + return None def retrieve_full_paper_text(self, query): pdf_text = str() @@ -247,6 +286,7 @@ def retrieve_full_paper_text(self, query): time.sleep(2.0) return pdf_text + """ import multiprocessing import sys @@ -298,8 +338,6 @@ def run_code(queue): import traceback import concurrent.futures - - import multiprocessing import io import sys @@ -311,7 +349,7 @@ def run_code(queue): def execute_code(code_str, timeout=60, MAX_LEN=1000): - #print(code_str) + # print(code_str) # prevent plotting errors import matplotlib @@ -323,7 +361,7 @@ def execute_code(code_str, timeout=60, MAX_LEN=1000): return "[CODE EXECUTION ERROR] pubmed Download took way too long. Program terminated" if "exit(" in code_str: return "[CODE EXECUTION ERROR] The exit() command is not allowed you must remove this." - #print(code_str) + # print(code_str) # Capturing the output output_capture = io.StringIO() sys.stdout = output_capture @@ -356,4 +394,3 @@ def run_code(): return output_capture.getvalue()[:MAX_LEN] -