sakana.ai の AI エージェント: AI Scientist のコード解析メモ

idea_reflection_prompt = """Round {current_round}/{num_reflections}.
In your thoughts, first carefully consider the quality, novelty, and feasibility of the idea you just created.
Include any other factors that you think are important in evaluating the idea.
Ensure the idea is clear and concise, and the JSON is the correct format.
Do not make things overly complicated.
In the next attempt, try and refine and improve your idea.
Stick to the spirit of the original idea unless there are glaring issues.

Respond in the same format as before:
THOUGHT:
<THOUGHT>

NEW IDEA JSON:

```json
<JSON>
\```

If there is nothing to improve, simply repeat the previous JSON EXACTLY after the thought and include "I am done" at the end of the thoughts but before the JSON.
ONLY INCLUDE "I am done" IF YOU ARE MAKING NO MORE CHANGES."""

# Iteratively improve task.
if num_reflections > 1:
    for j in range(num_reflections - 1):
        print(f"Iteration {j + 2}/{num_reflections}")
        text, msg_history = get_response_from_llm(
            idea_reflection_prompt.format(
                current_round=j + 2, num_reflections=num_reflections
            ),
            client=client,
            model=model,
            system_message=idea_system_prompt,
            msg_history=msg_history,
        )
        ## PARSE OUTPUT
        json_output = extract_json_between_markers(text)
        assert (
            json_output is not None
        ), "Failed to extract JSON from LLM output"
        print(json_output)

        if "I am done" in text:
            print(f"Idea generation converged after {j + 2} iterations.")
            break

for section in [
    "Introduction",
    "Background",
    "Method",
    "Experimental Setup",
    "Results",
    "Conclusion",
]:
    section_prompt = f"""Please fill in the {section} of the writeup. Some tips are provided below:
{per_section_tips[section]}

Be sure to use \cite or \citet where relevant, referring to the works provided in the file.
Do not cite anything that is not already in `references.bib`. Do not add any new entries to this.

Keep the experimental results (figures and tables) only in the Results section, and make sure that any captions are filled in.
In this pass, do not reference anything in later sections of the paper.

Before every paragraph, please include a brief description of what you plan to write in that paragraph in a comment.

Be sure to first name the file and use *SEARCH/REPLACE* blocks to perform these edits.
"""
    coder_out = coder.run(section_prompt)

    "Introduction": """
- Longer version of the Abstract, i.e. of the entire paper
- What are we trying to do and why is it relevant?
- Why is this hard?
- How do we solve it (i.e. our contribution!)
- How do we verify that we solved it (e.g. Experiments and results)
- New trend: specifically list your contributions as bullet points
- Extra space? Future work!
""",

for _ in range(num_cite_rounds):
    with open(osp.join(folder_name, "latex", "template.tex"), "r") as f:
        draft = f.read()
    prompt, done = get_citation_aider_prompt(
        cite_client, cite_model, draft, _, num_cite_rounds
    )
    if done:
        break
    if prompt is not None:
        # extract bibtex string
        bibtex_string = prompt.split('"""')[1]
        # insert this into draft before the "\end{filecontents}" line
        search_str = r"\end{filecontents}"
        draft = draft.replace(search_str, f"{bibtex_string}{search_str}")
        with open(osp.join(folder_name, "latex", "template.tex"), "w") as f:
            f.write(draft)
        coder_out = coder.run(prompt)

for section in [
    "Abstract",
    "Related Work",
    "Introduction",
    "Background",
    "Method",
    "Experimental Setup",
    "Results",
    "Conclusion",
]:
    coder_out = coder.run(
        second_refinement_prompt.format(
            section=section, tips=per_section_tips[section]
        )
        .replace(r"{{", "{")
        .replace(r"}}", "}")
    )

def get_citation_aider_prompt(
    client, model, draft, current_round, total_rounds
) -> Tuple[Optional[str], bool]:
    msg_history = []
    try:
        text, msg_history = get_response_from_llm(
            citation_first_prompt.format(
                draft=draft, current_round=current_round, total_rounds=total_rounds
            ),
            client=client,
            model=model,
            system_message=citation_system_msg.format(total_rounds=total_rounds),
            msg_history=msg_history,
        )
        if "No more citations needed" in text:
            print("No more citations needed.")
            return None, True

        ## PARSE OUTPUT
        json_output = extract_json_between_markers(text)
        assert json_output is not None, "Failed to extract JSON from LLM output"
        query = json_output["Query"]
        papers = search_for_papers(query)
    except Exception as e:
        print(f"Error: {e}")
        return None, False

    # ... (以下、検索結果の処理と最適な論文の選択)

citation_first_prompt = '''Round {current_round}/{total_rounds}:

You have written this LaTeX draft so far:

"""
{draft}
"""

Identify the most important citation that you still need to add, and the query to find the paper.

Respond in the following format:

THOUGHT:
<THOUGHT>

RESPONSE:
```json
<JSON>
\```

In <THOUGHT>, first briefly reason over the paper and identify where citations should be added.
If there are more citations needed, add "No more citations needed" to your thoughts.
Do not add "No more citations needed" if you are adding citations this round.

In <JSON>, respond in JSON format with the following fields:
- "Description": A precise description of the required edit, along with the proposed text and location where it should be made.
- "Query": The search query to find the paper (e.g. attention is all you need).

Ensure the description is sufficient to make the change without further context. Someone else will make the change.
The query will work best if you are able to recall the exact name of the paper you are looking for, or the authors.
This JSON will be automatically parsed, so ensure the format is precise.

citation_second_prompt = """Search has recovered the following articles:

{papers}

Respond in the following format:

THOUGHT:
<THOUGHT>

RESPONSE:
```json
<JSON>
\```

In <THOUGHT>, first briefly reason over the search results and identify which citation best fits your paper and the location is to be added at.
If none are appropriate, add "Do not add any" to your thoughts.

In <JSON>, respond in JSON format with the following fields:
- "Selected": A list of the indices of the selected papers to be cited, e.g. "[0, 1]". Can be "[]" if no papers are selected. This must be a string.
- "Description": Update the previous description of the required edit if needed. Ensure that any cites precisely match the name in the bibtex!!!

Do not select papers that are already in the `references.bib` file at the top of the draft, or if the same citation exists under a different name.
This JSON will be automatically parsed, so ensure the format is precise."""

dependencies = [
    "openai>=1.40.6",
    "anthropic[bedrock]>=0.34.0",
    "aider-chat>=0.50.1",
    "backoff>=2.2.1",
    "matplotlib>=3.9.2",
    "pypdf>=4.3.1",
    "pymupdf4llm>=0.0.10",
    "torch>=2.4.0",
    "numpy>=1.26.4",
    "transformers>=4.44.0",
    "datasets>=2.21.0",
    "tiktoken>=0.7.0",
    "wandb>=0.17.7",
    "tqdm>=4.66.5",
]

export OPENAI_API_KEY="YOUR KEY HERE"

python launch_scientist.py --model "gpt-4o-2024-05-13" --experiment nanoGPT_lite --num-ideas 2

import sys

from launch_scientist import main

if __name__ == "__main__":
    # コマンドライン引数を設定
    sys.argv = [
        "launch_scientist.py",
        "--model",
        "gpt-4o-2024-05-13",
        "--experiment",
        "nanoGPT_lite",
        "--num-ideas",
        "2",
    ]

    # main関数を実行
    main()