Save dataset to huggingface

app.py

@@ -2,11 +2,10 @@ import os
 import gradio as gr
 import requests
 import pandas as pd
+import datasets
 from mini_agents import master_agent
 from utils import get_full_file_path
-import subprocess
-
-subprocess.run(["playwright", "install"])
+from smolagents.memory import ActionStep, PlanningStep, TaskStep, SystemPromptStep, FinalAnswerStep
 
 # (Keep Constants as is)
 # --- Constants ---
@@ -14,13 +13,63 @@ DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 
 # --- Basic Agent Definition ---
 # ----- THIS IS WERE YOU CAN BUILD WHAT YOU WANT ------
+columns = [
+    'task_id',
+    'step_class',
+    # Common attributes (from MemoryStep base class)
+    'model_input_messages',
+    'tool_calls',
+    'start_time',
+    'end_time',
+    'step_number',
+    'error',
+    'duration',
+    'model_output_message',
+    'model_output',
+    'observations',
+    'observations_images',
+    'action_output',
+    # PlanningStep attributes
+    'plan',
+    # TaskStep attributes
+    'task',
+    'task_images',
+    # SystemPromptStep attributes
+    'system_prompt',
+    # FinalAnswerStep attributes
+    'final_answer'
+]
+df_agent_steps = pd.DataFrame(columns=columns)
+
 class BasicAgent:
     def __init__(self):
         self.agent = master_agent
         print("Master Agent initialized.")
-    def __call__(self, question: str) -> str:
+    def __call__(self, question: str, task_id: str) -> str:
         print(f"Agent received question (first 50 chars): {question[:50]}...")
         fixed_answer = self.agent.run(question)
+        all_steps = self.agent.memory.get_full_steps()
+        for step in all_steps:
+            if isinstance(step, ActionStep):
+                step_class = "ActionStep"
+            elif isinstance(step, PlanningStep):
+                step_class = "PlanningStep"
+            elif isinstance(step, TaskStep):
+                step_class = "TaskStep"
+            elif isinstance(step, SystemPromptStep):
+                step_class = "SystemPromptStep"
+            elif isinstance(step, FinalAnswerStep):
+                step_class = "FinalAnswerStep"
+            else:
+                step_class = "UnknownStep"
+
+            step_dict = step.dict()
+            df_agent_steps.loc[len(df_agent_steps)] = None
+            df_agent_steps.at[len(df_agent_steps), 'task_id'] = task_id
+            df_agent_steps.at[len(df_agent_steps), 'step_class'] = step_class
+            for key, value in step_dict.items():
+                df_agent_steps.at[len(df_agent_steps), key] = value
+            
         print(f"Agent returning fixed answer: {fixed_answer}")
         return fixed_answer
 
@@ -90,7 +139,7 @@ def run_and_submit_all( profile: gr.OAuthProfile | None, mock_submission: bool =
         try:
             if file_path:
                 question_text = question_text + f"\n\nHere is also the path to the file for the task (file name matches with task ID and is not in plain English): {file_path}"
-            submitted_answer = agent(question_text)
+            submitted_answer = agent(question_text, task_id)
             answers_payload.append({"task_id": task_id, "submitted_answer": submitted_answer})
             results_log.append({"Task ID": task_id, "Question": question_text, "Submitted Answer": submitted_answer})
         except Exception as e:
@@ -106,7 +155,12 @@ def run_and_submit_all( profile: gr.OAuthProfile | None, mock_submission: bool =
     status_update = f"Agent finished. Submitting {len(answers_payload)} answers for user '{username}'..."
     print(status_update)
 
-    # 5. Submit
+    # 5. Save steps data to huggingface dataset
+    print("Commiting steps data to huggingface dataset...")
+    dataset = datasets.Dataset.from_pandas(df_agent_steps)
+    dataset.push_to_hub("huytofu92/agent_steps_huggingface_course_unit4")
+    print("Agent steps data committed to huggingface dataset.")
+    # 6. Submit
     print(f"Submitting {len(answers_payload)} answers to: {submit_url}")
     if mock_submission:
         answer_df = pd.DataFrame(results_log, columns=["Task ID", "Question", "Submitted Answer"])