Update app.py

app.py

@@ -5,6 +5,7 @@ from reactagent.agents.agent_research import ResearchAgent
 from reactagent.runner import create_parser
 from reactagent import llm
 from reactagent.users.user import User
+import os
 import json
 
 
@@ -25,22 +26,29 @@ example_text = [
 # Load example JSON file
 def load_example_data():
     with open("example/example_data.json", "r") as json_file:
-        return json.load(json_file)
+        example_data = json.load(json_file)
+    
+    for idx in example_data.keys():
+        try:
+            with open(os.path.join("example", example_data[idx]["code_init"]), "r") as f:
+                example_data[idx]["code_init"] = f.read()
+        except FileNotFoundError:
+            print(f"File not found: {example_data[idx]["code_init"]}. Skipping key: {idx}")
+        try:
+            with open(os.path.join("example", example_data[idx]["code_final"]), "r") as f:
+                example_data[idx]["code_final"] = f.read()
+        except FileNotFoundError:
+            print(f"File not found: {example_data[idx]["code_final"]}. Skipping key: {idx}")     
+    return example_data
 
 example_data = load_example_data()
 
-with open("example/ex1_init.py", "r") as f:
-    predefined_code = f.read()
-
-with open("example/ex1_final.py", "r") as f:
-    final_code = f.read()
-
 # Function to handle the selection of an example and populate the respective fields
 def load_example(example_id):
     global index_ex 
     index_ex = str(example_id)
     example = example_data[index_ex]
-    paper_text = 'Title:\t' +  example['title'] + '\nAbstract:\t' + example['abstract']
+    paper_text = 'Title:\t' +  example['title'] + '\n\nAbstract:\t' + example['abstract']
     return paper_text
 
 example_text = [load_example(1), load_example(2)]
@@ -58,9 +66,8 @@ def load_example_and_set_index(paper_text_input):
 ########## Phase 1 ##############
 
 def extract_research_elements(paper_text):
-    global state_extract 
+    global state_extract, index_ex
     state_extract = True
-    global index_ex
     example = example_data[index_ex]
     tasks = example['research_tasks']
     gaps = example['research_gaps']
@@ -73,9 +80,8 @@ def extract_research_elements(paper_text):
 def generate_and_store(tasks, gaps, keywords, recent_works):
     if (not state_extract):
         return "", "", "", ""
-    global state_generate 
+    global state_generate, index_ex
     state_generate = True
-    global index_ex
     hypothesis = example_data[index_ex]['hypothesis']
     experiment_plan = example_data[index_ex]['experiment_plan']
     return hypothesis, experiment_plan, hypothesis, experiment_plan
@@ -83,17 +89,18 @@ def generate_and_store(tasks, gaps, keywords, recent_works):
 ##########  Phase 2 & 3  ##############
 def start_experiment_agent(hypothesis, plan):
     if (not state_extract or not  state_generate):
-        return "", ""
-    global state_agent 
+        return "", "", ""
+    global state_agent, step_index, state_complete
     state_agent = True
-    predefined_message = f"Implement the following hypothesis and experiment plan:\n\nHypothesis:\n{hypothesis}\n\nExperiment Plan:\n{plan}"
-    return predefined_code, predefined_action_log
+    step_index = 0
+    state_complete = False
+    # predefined_message = f"Implement the following hypothesis and experiment plan:\n\nHypothesis:\n{hypothesis}\n\nExperiment Plan:\n{plan}"
+    return example_data[index_ex]['code_init'], predefined_action_log, ""
 
 def submit_feedback(user_feedback, history, previous_response):
     if (not state_extract or not  state_generate or not state_agent):
         return "", "", ""
-    global step_index
-    global state_complete
+    global step_index, state_complete
     step_index += 1
     msg = history
     if step_index < len(process_steps):
@@ -108,7 +115,7 @@ def submit_feedback(user_feedback, history, previous_response):
         response = "Agent Finished."
         
 
-    return msg, response, predefined_code if state_complete else final_code
+    return msg, response, example_data[index_ex]['code_init'] if state_complete else example_data[index_ex]['code_final']
 
 def load_phase_2_inputs(hypothesis, plan):
     return hypothesis, plan, "# Code implementation will be displayed here after Start ExperimentAgent."
@@ -121,7 +128,7 @@ predefined_action_log = """
 Input: {"script_name": "train.py", "start_line_number": "1", "end_line_number": "74"}
 Objective: Understand the training script, including data processing, [...]
 [Observation]: The train.py script imports [...]. Sets random seeds [...]. Defines [...] Placeholder functions [...] exist without implementation. [...]
-[Feedback]: The script structure is clear, but key functions (train_model, predict) need proper implementation for proposed model training and prediction.
+[Feedback]: The script structure is clear, but key functions (train_model, predict) need proper implementation for proposed model training and prediction.\n
 """
 
 
@@ -131,7 +138,7 @@ Train MSE: 0.543,
 Test MSE: 0.688
 Epoch [2/10],
 Train MSE: 0.242,
-Test MSE: 0.493
+Test MSE: 0.493\n
 """
 
 # Initialize the global step_index and history
@@ -198,9 +205,13 @@ def handle_example_click(example_index):
     return load_example(index_ex)  # Simply return the text to display it in the textbox
 
 # Gradio Interface
-with gr.Blocks() as app:
+with gr.Blocks(theme=gr.themes.Default()) as app:
     gr.Markdown("# MLR- Copilot: Machine Learning Research based on LLM Agents")
-    gr.Markdown("MLR-Copilot is a framework where LLMs mimic researchers’ thought processes, designed to enhance the productivity of machine learning research by automating the generation and implementation of research ideas.It begins with a research paper, autonomously generating and validating these ideas, while incorporating human feedback to help reach executable research outcomes.")
+    gr.Markdown("###   ")
+
+    gr.Markdown("MLR-Copilot is a framework where LLMs mimic researchers’ thought processes, designed to enhance the productivity of machine learning research by automating the generation and implementation of research ideas. It begins with a research paper, autonomously generating and validating these ideas, while incorporating human feedback to help reach executable research outcomes.")
+
+
 
 
     # Use state variables to store generated hypothesis and experiment plan
@@ -208,51 +219,35 @@ with gr.Blocks() as app:
     experiment_plan_state = gr.State("")
 
     ########## Phase 1:  Research Idea Generation Tab ##############
-    with gr.Tab("Phase 1: Research Idea Generation"):
+    with gr.Tab("💡Stage 1: Research Idea Generation"):
         gr.Markdown("### Extract Research Elements and Generate Research Ideas")
         
         with gr.Row():
             with gr.Column():
-                paper_text_input = gr.Textbox(value=load_example(1), lines=10, label="Research Paper Text")
-                extract_button = gr.Button("Extract Research Elements")
+                paper_text_input = gr.Textbox(value=load_example(1), lines=10, label="📑 Research Paper Text")
+                extract_button = gr.Button("🔍 Extract Research Elements")
                 with gr.Row():
                     tasks_output = gr.Textbox(placeholder="Research task definition", label="Research Tasks", lines=2, interactive=True)
                     gaps_output = gr.Textbox(placeholder="Research gaps of current works", label="Research Gaps", lines=2, interactive=True)
                     keywords_output = gr.Textbox(placeholder="Paper keywords", label="Keywords", lines=2, interactive=True)
                     recent_works_output = gr.Textbox(placeholder="Recent works extracted from Semantic Scholar", label="Recent Works", lines=2, interactive=True)
             with gr.Column():
-                with gr.Row():  # Move the button to the top right
-                    generate_button = gr.Button("Generate Research Hypothesis & Experiment Plan")
+                with gr.Row():  # Move the button to the top 
+                    generate_button = gr.Button("✍️ Generate Research Hypothesis & Experiment Plan")
                 with gr.Group():
-                    gr.Markdown("### Research Idea")
+                    gr.Markdown("### 🌟 Research Idea")
                     with gr.Row():
                         hypothesis_output = gr.Textbox(label="Generated Hypothesis", lines=20, interactive=False)
                         experiment_plan_output = gr.Textbox(label="Generated Experiment Plan", lines=20, interactive=False)
 
-        # with gr.Row():
-        #     example_1_button = gr.Button("Load Example 1: " + example_data["1"]["title"])
-        #     example_2_button = gr.Button("Load Example 2: " + example_data["2"]["title"])
-        # Example buttons
         gr.Examples(
             examples=example_text,
             inputs=[paper_text_input],
             outputs=[paper_text_input, tasks_output, gaps_output, keywords_output, recent_works_output, hypothesis_output, experiment_plan_output],
             fn=load_example_and_set_index,
             run_on_click = True,
-            label="Click an example to load"
+            label="⬇️ Click an example to load"
         )
-
-       
-        # # Pre-step: load example         
-        # example_1_button.click(
-        #     fn=lambda: load_example(1), x
-        #     outputs=[paper_text_input]
-        # )
-        
-        # example_2_button.click(
-        #     fn=lambda: load_example(2), 
-        #     outputs=[paper_text_input]
-        # )
     
         # Step 1: Extract Research Elements
         extract_button.click(
@@ -270,24 +265,27 @@ with gr.Blocks() as app:
 
 
     ########## Phase 2 & 3:  Experiment implementation and execution ##############
-    with gr.Tab("Phase 2 & Phase 3: Experiment implementation and execution"):
+    with gr.Tab("🧪 Stage 2 & Stage 3: Experiment implementation and execution"):
         gr.Markdown("### Interact with the ExperimentAgent")
 
         with gr.Row():
             with gr.Column():
-                idea_input = gr.Textbox(label="Research Hypothesis", lines=30, interactive=False)
-                plan_input = gr.Textbox(label="Experiment Plan", lines=30, interactive=False)
+                with gr.Group():
+                    gr.Markdown("### 🌟 Generated Research Idea")
+                    with gr.Row():
+                        idea_input = gr.Textbox(label="Generated Research Hypothesis", lines=30, interactive=False)
+                        plan_input = gr.Textbox(label="Generated Experiment Plan", lines=30, interactive=False)
         
             with gr.Column():
-                start_exp_agnet = gr.Button("Start ExperimentAgent", elem_classes=["agent-btn"])
+                start_exp_agnet = gr.Button("⚙️ Start / Reset ExperimentAgent", elem_classes=["agent-btn"])
                 with gr.Group():
                     gr.Markdown("### Implementation + Execution Log")    
-                    log = gr.Textbox(label="Execution Log", lines=20, interactive=False)
-                    code_display = gr.Code(label="Implementation", language="python", interactive=False)
+                    log = gr.Textbox(label="📖 Execution Log", lines=20, interactive=False)
+                    code_display = gr.Code(label="🧑‍💻 Implementation", language="python", interactive=False)
 
             with gr.Column():
-                response = gr.Textbox(label="ExperimentAgent Response", lines=30, interactive=False)
-                feedback = gr.Textbox(placeholder="N/A", label="User Feedback", lines=3, interactive=True)
+                response = gr.Textbox(label="🤖 ExperimentAgent Response", lines=30, interactive=False)
+                feedback = gr.Textbox(placeholder="N/A", label="🧑‍🔬 User Feedback", lines=3, interactive=True)
                 submit_button = gr.Button("Submit", elem_classes=["Submit-btn"])
         
         hypothesis_state.change(
@@ -300,7 +298,7 @@ with gr.Blocks() as app:
         start_exp_agnet.click(
             fn=start_experiment_agent, 
             inputs=[hypothesis_state, experiment_plan_state], 
-            outputs=[code_display, log]
+            outputs=[code_display, log, response]
         )
         
         submit_button.click(
@@ -311,4 +309,4 @@ with gr.Blocks() as app:
 
 if __name__ == "__main__":
     step_index = 0
-    app.launch(share=True)
+    app.launch(share=True)