added evaluation + leaderboard generation + reorg of the viz

yourbench_space/app.py

@@ -1,10 +1,13 @@
+import asyncio
 import os
 import sys
 import time
 import gradio as gr
-from pathlib import Path
-from loguru import logger
+
+from datasets import load_dataset
 from huggingface_hub import whoami
+from loguru import logger
+from pathlib import Path
 
 from yourbench_space.config import generate_and_save_config
 from yourbench_space.utils import (
@@ -15,6 +18,8 @@ from yourbench_space.utils import (
     update_dataset,
     STAGES,
 )
+from yourbench_space.evaluation import create_eval_file, run_evaluations
+from yourbench_space.leaderboard_space.env import HF_TOKEN
 
 project_description = """
 # YourBench 🚀  
@@ -74,11 +79,11 @@ def update_process_status():
     
     return gr.update(value=True, label="Process Status: Running")
 
-def prepare_task(oauth_token: gr.OAuthToken | None, hf_dataset_prefix: str, _=None):
+def prepare_task(oauth_token: gr.OAuthToken | None, hf_dataset_name: str, _=None):
     new_env = os.environ.copy()
     if oauth_token:
         new_env["HF_TOKEN"] = oauth_token.token
-    new_env["DATASET_PREFIX"] = hf_dataset_prefix
+    new_env["DATASET_PREFIX"] = hf_dataset_name
     manager.start_process(custom_env=new_env)
 
 
@@ -104,31 +109,55 @@ def switch_to_run_generation_tab():
 def enable_button(files):
     return gr.update(interactive=bool(files))
 
+def run_evaluation_pipeline(oauth_token: gr.OAuthToken | None, org_name, eval_name):
+    # Test dataset existence
+    eval_ds_name = f"{org_name}/{eval_name}"
+    # Test dataset existence
+    try:
+        load_dataset(eval_ds_name, streaming=True)
+    except Exception as e:
+        print(f"Error while loading the dataset: {e}")
+        return
+    # Run evaluations
+    create_eval_file(eval_ds_name)
+    status = asyncio.run(run_evaluations(eval_ds_name=eval_ds_name, org=org_name))
+    # Create space
+    from huggingface_hub import HfApi
+    repo_id = f"{org_name}/leaderboard_yourbench_{eval_ds_name.replace('/', '_')}"
+    api = HfApi()
+
+    try:
+        api.create_repo(repo_id=repo_id, repo_type="space", space_sdk="gradio")
+        api.upload_folder(repo_id=repo_id, repo_type="space", folder_path="src/")
+        api.add_space_secret(repo_id=repo_id, key="HF_TOKEN", value=HF_TOKEN)
+        api.add_space_variable(repo_id=repo_id, key="TASK", value=eval_ds_name)
+        api.add_space_variable(repo_id=repo_id, key="ORG_NAME", value=org_name)
+    except Exception as e:
+        status = "Evaluation" + status + "\nLeaderboard creation:" + e
+    return status
+
 
 with gr.Blocks(theme=gr.themes.Default()) as app:
     gr.Markdown(project_description)
 
-    with gr.Row():
-        gr.Markdown("## ⚙️ YourBench Setup", elem_id="setup-title")
-        with gr.Column(scale=0, min_width=150):
-            login_btn = gr.LoginButton()
-
     with gr.Tabs() as tabs:
         with gr.Tab("Setup", id=0):
             with gr.Row():
-                with gr.Accordion("Hugging Face Settings"):
-                    hf_org_dropdown = gr.Dropdown(
-                        choices=[], label="Organization", allow_custom_value=True
-                    )
-                    app.load(
-                        update_hf_org_dropdown, inputs=None, outputs=hf_org_dropdown
-                    )
-
-                    hf_dataset_prefix = gr.Textbox(
-                        label="Dataset Prefix",
-                        value="yourbench",
-                        info="Prefix applied to all datasets",
-                    )
+                with gr.Column():
+                    login_btn = gr.LoginButton()
+                    with gr.Accordion("Hugging Face Settings"):
+                        hf_org_dropdown = gr.Dropdown(
+                            choices=[], label="Organization", allow_custom_value=True
+                        )
+                        app.load(
+                            update_hf_org_dropdown, inputs=None, outputs=hf_org_dropdown
+                        )
+
+                        hf_dataset_name = gr.Textbox(
+                            label="Dataset name",
+                            value="yourbench",
+                            info="Name of your new evaluation dataset",
+                        )
 
                 with gr.Accordion("Upload documents"):
                     file_input = gr.File(
@@ -142,18 +171,18 @@ with gr.Blocks(theme=gr.themes.Default()) as app:
                         file_input,
                         output,
                     )
-
-            preview_button = gr.Button("Generate New Config", interactive=False)
-            log_message = gr.Textbox(label="Log Message", visible=True)
-            download_button = gr.File(
-                label="Download Config", visible=False, interactive=False
-            )
+            with gr.Row():
+                preview_button = gr.Button("Generate New Config", interactive=False)
+                log_message = gr.Textbox(label="Log Message", visible=True)
+                download_button = gr.File(
+                    label="Download Config", visible=False, interactive=False
+                )
 
             file_input.change(enable_button, inputs=file_input, outputs=preview_button)
 
             preview_button.click(
                 generate_and_return,
-                inputs=[hf_org_dropdown, hf_dataset_prefix],
+                inputs=[hf_org_dropdown, hf_dataset_name],
                 outputs=[log_message, download_button],
             )
             preview_button.click(
@@ -164,9 +193,26 @@ with gr.Blocks(theme=gr.themes.Default()) as app:
 
         with gr.Tab("Run Generation", id=1):
             with gr.Row():
-                with gr.Accordion("Log Output", open=True):
-                    log_output = gr.Code(language=None, lines=20, interactive=False)
+                start_button = gr.Button("Start Task")
+                start_button.click(prepare_task, inputs=[login_btn, hf_dataset_name])
+
+                stop_button = gr.Button("Stop Task")
+                stop_button.click(manager.stop_process)
+
+                kill_button = gr.Button("Kill Task")
+                kill_button.click(manager.kill_process)
 
+            with gr.Column():
+                with gr.Row():
+                    with gr.Accordion("Log Output", open=True):
+                        log_output = gr.Code(language=None, lines=20, interactive=False)
+
+                with gr.Row():
+                    process_status = gr.Checkbox(label="Process Status", interactive=False)
+                    status_timer = gr.Timer(1.0, active=True)
+                    status_timer.tick(update_process_status, outputs=process_status)
+
+            with gr.Column():
                 with gr.Accordion("Stages", open=True):
                     stages_table = gr.CheckboxGroup(
                         choices=STAGES,
@@ -175,27 +221,6 @@ with gr.Blocks(theme=gr.themes.Default()) as app:
                         interactive=False,
                     )
 
-            log_timer = gr.Timer(1.0, active=True)
-            log_timer.tick(
-                manager.read_and_get_output, outputs=[log_output, stages_table]
-            )
-
-            with gr.Row():
-                process_status = gr.Checkbox(label="Process Status", interactive=False)
-                status_timer = gr.Timer(1.0, active=True)
-                status_timer.tick(update_process_status, outputs=process_status)
-
-            with gr.Row():
-                start_button = gr.Button("Start Task")
-                start_button.click(prepare_task, inputs=[login_btn, hf_dataset_prefix])
-
-                stop_button = gr.Button("Stop Task")
-                stop_button.click(manager.stop_process)
-
-                kill_button = gr.Button("Kill Task")
-                kill_button.click(manager.kill_process)
-            
-            with gr.Row():
                 with gr.Accordion("Ingestion"):
                     ingestion_df = gr.DataFrame()
                     
@@ -209,8 +234,19 @@ with gr.Blocks(theme=gr.themes.Default()) as app:
                     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]
+                    update_dataset, inputs=[stages_table, hf_org_dropdown, hf_dataset_name], outputs=[ingestion_df, summarization_df, single_hop, answers_df]
                 )
 
+            log_timer = gr.Timer(1.0, active=True)
+            log_timer.tick(
+                manager.read_and_get_output, outputs=[log_output, stages_table]
+            )
+        with gr.Tab("Evaluate", id=2):
+            with gr.Row():
+                btn_launch_evals = gr.Button("Launch evaluations")
+                status = gr.Textbox(label="Status")
+
+            btn_launch_evals.click(run_evaluation_pipeline, [hf_org_dropdown, hf_dataset_name], status)
+
 
 app.launch(allowed_paths=["/app"])

yourbench_space/evaluation.py

@@ -0,0 +1,89 @@
+import asyncio, os
+from src.env import INIT_MODELS
+
+ON_SPACES=os.environ.get("system") == "spaces"
+OUTPUT_DIR = "/data" if ON_SPACES else "."
+
+def create_eval_file(eval_ds_name):
+    # TODO: replace by Nathan's call
+    content = """
+from aenum import extend_enum
+
+from lighteval.metrics.metrics import Metrics
+from lighteval.metrics.utils.metric_utils import (
+    CorpusLevelMetricGrouping,
+    MetricCategory,
+    MetricUseCase,
+)
+from lighteval.tasks.lighteval_task import LightevalTaskConfig
+from lighteval.tasks.extended.hle.main import JudgeLLMHLE
+from lighteval.tasks.requests import Doc
+
+
+def prompt_function(line, task_name: str = None):
+    if line["image"] not in [None, ""]:
+        return
+
+    return Doc(
+        task_name=task_name,
+        query="Question: " + line["question"] + "\\nAnswer:",
+        choices=[line["answer"]],
+        gold_index=0,
+        specific={"question": line["question"]},
+    )
+""" + f"""
+
+hle = LightevalTaskConfig(
+    name="{eval_ds_name.replace('/', '_')}",
+    suite=["custom"],
+    prompt_function=prompt_function,
+    hf_repo="{eval_ds_name}",
+    hf_subset="default",
+    hf_avail_splits=["test"],
+    evaluation_splits=["test"],
+    few_shots_split=None,
+    few_shots_select=None,
+    generation_size=8192,
+    metric=[Metrics.exact_match],
+    stop_sequence=[],
+    trust_dataset=True,
+    version=0,
+)
+
+
+TASKS_TABLE = [hle]
+""" 
+    
+    with open(f"{OUTPUT_DIR}/custom_task.py", "w") as f:
+        f.write(content)
+
+async def run_process(args: list) -> dict:
+    process = await asyncio.create_subprocess_exec(
+        *args,
+        stdout=asyncio.subprocess.PIPE,
+        stderr=asyncio.subprocess.PIPE
+    )
+    await asyncio.wait_for(process.wait(), timeout=180)
+    stdout = await process.stdout.read()
+    stderr = await process.stderr.read()
+    return {
+        'pid': process.pid,
+        'stdout': stdout.decode(),
+        'stderr': stderr.decode()
+    }
+
+async def run_evaluations(eval_ds_name: str, org: str) -> list:
+    tasks = []
+    for model_name, provider in INIT_MODELS:
+        args = [
+            "lighteval", 
+            "endpoint", "inference-providers", f"model={model_name},provider={provider}",
+            f"custom|{eval_ds_name.replace('/', '_')}|0|0", "--custom-tasks", f"{OUTPUT_DIR}/custom_task.py", "--max-samples", "10", 
+            "--output-dir", f"{OUTPUT_DIR}", "--save-details", "--results-org", org, "--push-to-hub"
+        ]
+        tasks.append(run_process(args))
+    # Will capture the task if failed
+    processes = await asyncio.gather(*tasks, return_exceptions=True)
+    if all(not isinstance(result, Exception) for result in processes):
+        return "✅"
+    return "At least one model failed"

yourbench_space/leaderboard_space/app.py

@@ -0,0 +1,31 @@
+import gradio as gr
+
+from utils import run_pipeline, update_examples
+from env import TASK
+
+with gr.Blocks(
+    title="YourBench Leaderboard",
+    theme=gr.themes.Soft(),
+    css="button { margin: 0 10px; padding: 5px 15px; }",
+) as demo:
+    # DISPLAY TABLE AND ANALYSIS
+    title = gr.Markdown(f"YourBench auto-Leaderboard for {TASK}")
+    leaderboard = gr.DataFrame(label="Results", interactive=False)
+    samples_ix = gr.Number(
+            label="Example Index", 
+            value=0,
+            step=1,
+            info="Navigate through different examples"
+        )
+    with gr.Tab("Hardest samples"):
+        hard_samples = gr.HTML()
+    with gr.Tab("Easiest samples"):
+        easy_samples = gr.HTML()
+    with gr.Tab("All samples"):
+        all_samples = gr.HTML()
+
+    samples_ix.change(update_examples, samples_ix, [easy_samples, hard_samples, all_samples])
+
+    demo.load(run_pipeline, [samples_ix], [leaderboard, easy_samples, hard_samples, all_samples])
+
+demo.launch()

yourbench_space/leaderboard_space/env.py

@@ -0,0 +1,16 @@
+import os
+INIT_MODELS = [
+    # 70B
+    ("Qwen/Qwen2.5-72B-Instruct", "sambanova"),
+    ("meta-llama/Llama-3.3-70B-Instruct", "together"),
+    ("deepseek-ai/DeepSeek-R1-Distill-Llama-70B", "sambanova"),
+    # 20 to 30B
+    ("Qwen/QwQ-32B", "sambanova"),
+    ("mistralai/Mistral-Small-24B-Instruct-2501", "together"),
+    #("allenai/OLMo-2-0325-32B-Instruct", "hf-inference")
+] 
+MODELS = [m[0] for m in INIT_MODELS]
+TASK = os.getenv("TASK")
+# With storage
+HF_TOKEN=os.getenv("HF_TOKEN")
+ORG_NAME = os.getenv("ORG_NAME")

yourbench_space/leaderboard_space/requirements.txt

@@ -0,0 +1,2 @@
+datasets
+huggingface_hub

yourbench_space/leaderboard_space/utils.py

@@ -0,0 +1,198 @@
+from datasets import load_dataset, Dataset
+from functools import lru_cache
+from typing import Tuple
+import gradio as gr
+import json
+
+from env import MODELS, TASK, ORG_NAME
+
+def aggregate_results() -> list:
+    """From the path of outputs and model list, extracts the current scores and stores them in a list of dicts with model, score, time as keys 
+    """
+    all_results = []
+    for org_model in MODELS:
+        try:
+            path = f"{ORG_NAME}/details_{org_model.replace('/', '__')}_private"
+            ds = load_dataset(path, "results", split="latest")
+            config = json.loads(ds["config_general"][0])
+            results = json.loads(ds["results"][0])
+                
+            # Model data
+            org, model = org_model.split("/")
+
+            cur_result = {
+                "Org": org, 
+                "Model": model, 
+                "Duration (s)": config["end_time"] - config["start_time"]
+            }
+                    
+            # Extract the task from the JSON data
+            for k_metric, v_dict in results.items():
+                if k_metric != "all":
+                    for k, v in v_dict.items():
+                        cur_result[f"{k}({k_metric})"] = v
+            all_results.append(cur_result)
+        except Exception as e:
+            print(f"Error processing {model} {ORG_NAME}: {e}")
+    return all_results
+
+def extract_dataviz() -> Tuple[list, list]:
+    """From the path of outputs and model list, extracts from the details the worst samples, best samples
+    """
+    all_samples = {}
+    for org_model in MODELS:
+        try:
+            path = f"{ORG_NAME}/details_{org_model.replace('/', '__')}_private"
+            ds = load_dataset(path, f"custom_{TASK.replace('/', '_')}_0", split="latest")
+
+            for ix, row in enumerate(ds):
+                prompt = row["full_prompt"]
+                gold = row["gold"]
+                score = list(row["metrics"].values())[0]
+                prediction = row["predictions"][0]
+
+
+                # We store flattened samples in a dict
+                # ix -> ix, prompt, gold, model_score for each model, model_prediction for each model
+                # then 2 lists: model_scores and models, to aggreg more easily
+                if ix not in all_samples:
+                    all_samples[ix] = {
+                        "ix": ix,
+                        "prompt": prompt,
+                        "gold": gold[0] if isinstance(gold, list) else gold,
+                        # A bit redundant, but put in their own boxes for simplicity of access later
+                        "model_scores": [],
+                        "models": []
+                    }
+                if org_model not in all_samples[ix]["models"]:
+                    all_samples[ix][f"{org_model}_score"] = row["metrics"]
+                    all_samples[ix][f"{org_model}_prediction"] = prediction
+                    all_samples[ix]["model_scores"].append(score)
+                    all_samples[ix]["models"].append(org_model)
+
+        except Exception as e:
+            print(f"Error processing {org_model}: {e}")
+
+    full_samples = sorted(list(all_samples.values()), key= lambda r: r['ix'])
+    hard_samples = sorted([sample for sample in all_samples.values() if sum(sample["model_scores"]) == 0], key= lambda r: r['ix'])
+    easy_samples = sorted([sample for sample in all_samples.values() if sum(sample["model_scores"]) == len(sample["model_scores"])], key= lambda r: r['ix'])
+    return easy_samples, hard_samples, full_samples
+
+def samples_to_box_display(samples: list, example_index: int = 0):
+    """Adapted from Nathan's code in https://huggingface.co/spaces/SaylorTwift/OpenEvalsModelDetails/
+    """
+    if len(samples) == 0:
+        return "No samples in this category!"
+    outputs = []
+    sample = samples[example_index]
+    for model in sample["models"]:
+        try:
+            outputs.append({
+                'Model': model,
+                'Prediction': sample[f'{model}_prediction'],
+                'Prompt': sample['prompt'],
+                'Metrics': sample[f'{model}_score'],
+                'Gold': sample['gold']
+            })
+        except (KeyError, IndexError):
+            continue
+    
+    if not outputs:
+        return "No results found for the selected combination."
+    
+    # Create HTML output with all models
+    html_output = "<div style='max-width: 800px; margin: 0 auto;'>\n\n"
+    
+    # Show gold answer at the top with distinct styling
+    if outputs:
+        html_output += "<div style='background: #e6f3e6; padding: 20px; border-radius: 10px; margin-bottom: 20px;'>\n"
+        html_output += "<h3 style='margin-top: 0;'>Ground Truth</h3>\n"
+        html_output += "<div style='overflow-x: auto; max-width: 100%;'>\n"
+        html_output += f"<pre style='white-space: pre-wrap; word-wrap: break-word; margin: 0;'><code>{outputs[0]['Gold']}</code></pre>\n"
+        html_output += "</div>\n"
+        html_output += "</div>\n"
+    
+    for output in outputs:
+        html_output += "<div style='background: #f5f5f5; padding: 20px; margin-bottom: 20px; border-radius: 10px;'>\n"
+        html_output += f"<h2 style='margin-top: 0;'>{output['Model']}</h2>\n"
+        
+        # Format metrics as a clean table
+        html_output += "<details open style='margin-bottom: 15px;'>\n"
+        html_output += "<summary><h3 style='display: inline; margin: 0;'>Metrics</h3></summary>\n"
+        metrics = output['Metrics']
+        if isinstance(metrics, str):
+            metrics = eval(metrics)
+        html_output += "<div style='overflow-x: auto;'>\n"
+        html_output += "<table style='width: 100%; margin: 10px 0; border-collapse: collapse;'>\n"
+        for key, value in metrics.items():
+            if isinstance(value, float):
+                value = f"{value:.3f}"
+            html_output += f"<tr><td style='padding: 5px; border-bottom: 1px solid #ddd;'><strong>{key}</strong></td><td style='padding: 5px; border-bottom: 1px solid #ddd;'>{value}</td></tr>\n"
+        html_output += "</table>\n"
+        html_output += "</div>\n"
+        html_output += "</details>\n\n"
+        
+        # Handle prompt formatting with better styling
+        html_output += "<details style='margin-bottom: 15px;'>\n"
+        html_output += "<summary><h3 style='display: inline; margin: 0;'>Prompt</h3></summary>\n"
+        html_output += "<div style='background: #ffffff; padding: 15px; border-radius: 5px; margin-top: 10px;'>\n"
+        
+        prompt_text = output['Prompt']
+        if isinstance(prompt_text, list):
+            for i, msg in enumerate(prompt_text):
+                if isinstance(msg, dict) and 'content' in msg:
+                    role = msg.get('role', 'message').title()
+                    html_output += "<div style='margin-bottom: 10px;'>\n"
+                    html_output += f"<strong>{role}:</strong>\n"
+                    html_output += "<div style='overflow-x: auto;'>\n"
+                    html_output += f"<pre style='white-space: pre-wrap; word-wrap: break-word; margin: 5px 0;'><code>{msg['content']}</code></pre>\n"
+                    html_output += "</div>\n"
+                    html_output += "</div>\n"
+                else:
+                    html_output += "<div style='margin-bottom: 10px;'>\n"
+                    html_output += "<div style='overflow-x: auto;'>\n"
+                    html_output += f"<pre style='white-space: pre-wrap; word-wrap: break-word; margin: 5px 0;'><code>{json.dumps(msg, indent=2)}</code></pre>\n"
+                    html_output += "</div>\n"
+                    html_output += "</div>\n"
+        else:
+            html_output += "<div style='overflow-x: auto;'>\n"
+            if isinstance(prompt_text, dict) and 'content' in prompt_text:
+                html_output += f"<pre style='white-space: pre-wrap; word-wrap: break-word; margin: 5px 0;'><code>{prompt_text['content']}</code></pre>\n"
+            else:
+                html_output += f"<pre style='white-space: pre-wrap; word-wrap: break-word; margin: 5px 0;'><code>{prompt_text}</code></pre>\n"
+            html_output += "</div>\n"
+        
+        html_output += "</div>\n"
+        html_output += "</details>\n\n"
+        
+        # Style prediction output - now in a collapsible section
+        html_output += "<details open style='margin-bottom: 15px;'>\n"
+        html_output += "<summary><h3 style='display: inline; margin: 0;'>Prediction</h3>"
+        # Add word count in a muted style
+        word_count = len(output['Prediction'].split())
+        html_output += f"<span style='color: #666; font-size: 0.8em; margin-left: 10px;'>({word_count} words)</span>"
+        html_output += "</summary>\n"
+        html_output += "<div style='background: #ffffff; padding: 15px; border-radius: 5px; margin-top: 10px;'>\n"
+        html_output += "<div style='overflow-x: auto;'>\n"
+        html_output += f"<pre style='white-space: pre-wrap; word-wrap: break-word; margin: 0;'><code>{output['Prediction']}</code></pre>\n"
+        html_output += "</div>\n"
+        html_output += "</div>\n"
+        html_output += "</details>\n"
+        html_output += "</div>\n\n"
+    
+    html_output += "</div>"
+    return html_output
+
+def run_pipeline(samples_ix: int = 0):
+    results = aggregate_results()
+    best_samples, worst_samples, all_samples = extract_dataviz()
+    return gr.Dataframe(Dataset.from_list(results).to_pandas(), visible=True), \
+        gr.HTML(samples_to_box_display(best_samples, samples_ix), label="Easiest samples (always found)", visible=True), \
+        gr.HTML(samples_to_box_display(worst_samples, samples_ix), label="Hardest samples (always failed)", visible=True), \
+        gr.HTML(samples_to_box_display(all_samples, samples_ix), label="All samples", visible=True)
+
+def update_examples(samples_ix: int = 0):
+    best_samples, worst_samples, all_samples = extract_dataviz()
+    return samples_to_box_display(best_samples, samples_ix), \
+        samples_to_box_display(worst_samples, samples_ix), \
+        samples_to_box_display(all_samples, samples_ix)

yourbench_space/utils.py

@@ -22,6 +22,8 @@ STAGES = [
     "chunking",
     "single_shot_question_generation",
     "answer_generation",
+    #"evaluate_models",
+    #"create_leaderboard"
     # "judge_answers", # to uncomment when fixed 
 ]
 
@@ -55,20 +57,31 @@ 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())
-    
+    ingestion_df = pd.DataFrame()
+    summarization_df = pd.DataFrame()
+    single_hop_df = pd.DataFrame()
+    answers_df = pd.DataFrame()
+
     # Construct dataset name from config
     dataset_name = f"{hf_org}/{hf_prefix}"
+
+    # TODO: add cache dir
+    # Will be able to group everything in one pass once the names get homogeneized
+    if "ingestion" in stages:
+        # TODO: why is the key "ingested" and not "ingestion"? (does not match the other splits)
+        ingestion_ds = load_dataset(dataset_name, name="ingested", split="train", streaming=True)
+        ingestion_df = pd.DataFrame([next(iter(ingestion_ds)) for _ in range(5)])
+    if "summarization" in stages:
+        summarization_ds = load_dataset(dataset_name, name="summarization", split="train", streaming=True)
+        summarization_df = pd.DataFrame([next(iter(summarization_ds)) for _ in range(5)])
+    if "single_shot_question_generation" in stages:
+        single_hop_ds = load_dataset(dataset_name, name="single_shot_question_generation", split="train", streaming=True)
+        single_hop_df = pd.DataFrame([next(iter(single_hop_ds)) for _ in range(5)])
+    if "answer_generation" in stages:
+        answers_ds = load_dataset(dataset_name, name="answer_generation", split="train", streaming=True)
+        answers_df = pd.DataFrame([next(iter(answers_ds)) for _ in range(5)])
     
-    # 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)
+    return (ingestion_df, summarization_df, single_hop_df, answers_df)
 
 class SubprocessManager:
     def __init__(self, command):