Add resulted datasets [wip]

yourbench_space/app.py

@@ -12,6 +12,7 @@ from yourbench_space.utils import (
     UPLOAD_DIRECTORY,
     SubprocessManager,
     save_files,
+    update_dataset,
     STAGES,
 )
 
@@ -60,6 +61,8 @@ def generate_and_return(hf_org, hf_prefix):
         )
     )
 
+final_dataset = None
+
 def update_process_status():
     """Update process status and include exit details if process has terminated"""
     is_running = manager.is_running()
@@ -191,5 +194,23 @@ with gr.Blocks(theme=gr.themes.Default()) as app:
 
                 kill_button = gr.Button("Kill Task")
                 kill_button.click(manager.kill_process)
+            
+            with gr.Row():
+                with gr.Accordion("Ingestion"):
+                    ingestion_df = gr.DataFrame()
+                    
+                with gr.Accordion("Summarization"):
+                    summarization_df = gr.DataFrame()
+                    
+                with gr.Accordion("Single-Hop"):
+                    single_hop = gr.DataFrame()
+
+                with gr.Accordion("Answer Generation"):
+                    answers_df = gr.DataFrame()
+            
+                stages_table.change(
+                    update_dataset, inputs=[stages_table, hf_org_dropdown, hf_dataset_prefix], outputs=[ingestion_df, summarization_df, single_hop, answers_df]
+                )
+
 
 app.launch(allowed_paths=["/app"])

yourbench_space/config.py

@@ -98,3 +98,4 @@ def generate_and_save_config(hf_org, hf_prefix):
     file_path = save_yaml_file(config)
     logger.success(f"Config saved at: {file_path}")
     return file_path
+

yourbench_space/utils.py

@@ -3,8 +3,10 @@ import os
 import re
 import pathlib
 import shutil
-from loguru import logger
 import subprocess
+import pandas as pd
+from datasets import load_dataset, get_dataset_config_names
+from loguru import logger
 from typing import List
 
 UPLOAD_DIRECTORY = pathlib.Path("/app/uploaded_files")
@@ -49,6 +51,24 @@ def save_files(files: List[pathlib.Path]) -> str:
         else "No files were saved"
     )
 
+def update_dataset(stages, hf_org, hf_prefix):
+    """
+    Updates the dataset based on the provided stages and dataset configuration.
+    """
+    # If not the final step, return empty dataframes
+    if stages[-1] != STAGES[-1]:
+        return (pd.DataFrame(), pd.DataFrame(), pd.DataFrame(), pd.DataFrame())
+    
+    # Construct dataset name from config
+    dataset_name = f"{hf_org}/{hf_prefix}"
+    
+    # Load datasets and convert to DataFrame
+    ingestion_df = pd.DataFrame(load_dataset(dataset_name, name="ingested", split="train", streaming=True))
+    summarization_df = pd.DataFrame(load_dataset(dataset_name, name="summarization", split="train", streaming=True))
+    single_hop = pd.DataFrame(load_dataset(dataset_name, name="single_shot_question_generation", split="train", streaming=True))
+    answers_df = pd.DataFrame(load_dataset(dataset_name, name="answer_generation", split="train", streaming=True))
+    
+    return (ingestion_df, summarization_df, single_hop, answers_df)
 
 class SubprocessManager:
     def __init__(self, command):