Change UI and config

yourbench_space/app.py

@@ -2,7 +2,6 @@ import os
 import sys
 import time
 import gradio as gr
-import yaml
 from pathlib import Path
 from loguru import logger
 from huggingface_hub import whoami
@@ -13,9 +12,9 @@ from yourbench_space.utils import (
     UPLOAD_DIRECTORY,
     SubprocessManager,
     save_files,
+    STAGES,
 )
 
-# Short project description
 project_description = """
 # YourBench 🚀  
 A Dynamic Benchmark Generation Framework 
@@ -35,44 +34,43 @@ logger.add(sys.stderr, level="INFO")
 command = ["uv", "run", "yourbench", f"--config={CONFIG_PATH}"]
 manager = SubprocessManager(command)
 
-# Create a citation section
 docs_path = Path(__file__).parent / "docs.md"
-
-# Read the file safely
-if docs_path.exists():
-    docs_content = docs_path.read_text()
-else:
-    docs_content = "# Citation\n\nDocumentation file not found."
-
-citation_content = docs_content.split("# Citation")[-1].strip()
+citation_content = (
+    docs_path.read_text().split("# Citation")[-1].strip()
+    if docs_path.exists()
+    else "# Citation\n\nDocumentation file not found."
+)
 
 def generate_and_return(hf_org, hf_prefix):
-    """Handles config generation and validates file existence before enabling download"""
     generate_and_save_config(hf_org, hf_prefix)
-
-    # Wait until the config file is actually created
     for _ in range(5):
         if CONFIG_PATH.exists():
             break
         time.sleep(0.5)
 
-    if CONFIG_PATH.exists():
-        return "✅ Config saved!", gr.update(value=str(CONFIG_PATH), visible=True, interactive=True)
-    else:
-        return "❌ Config generation failed.", gr.update(visible=False, interactive=False)
+    return (
+        (
+            "✅ Config saved!",
+            gr.update(value=str(CONFIG_PATH), visible=True, interactive=True),
+        )
+        if CONFIG_PATH.exists()
+        else (
+            "❌ Config generation failed.",
+            gr.update(visible=False, interactive=False),
+        )
+    )
 
-def prepare_task(oauth_token: gr.OAuthToken | None, model_token: str):
-    """Prepares and starts the subprocess with environment variables."""
+
+def prepare_task(oauth_token: gr.OAuthToken | None, hf_dataset_prefix: str, _=None):
     new_env = os.environ.copy()
     if oauth_token:
         new_env["HF_TOKEN"] = oauth_token.token
-    new_env["MODEL_API_KEY"] = model_token
+    new_env["DATASET_PREFIX"] = hf_dataset_prefix
     manager.start_process(custom_env=new_env)
 
+
 def update_hf_org_dropdown(oauth_token: gr.OAuthToken | None):
-    """Updates the dropdown with the user's Hugging Face organizations"""
     if oauth_token is None:
-        print("Please deploy this on Spaces and log in to view the list of available organizations")
         return gr.Dropdown([], label="Organization")
 
     try:
@@ -83,38 +81,52 @@ def update_hf_org_dropdown(oauth_token: gr.OAuthToken | None):
         return gr.Dropdown(org_names, value=user_name, label="Organization")
 
     except Exception as e:
-        print(f"Error retrieving user info: {e}")
         return gr.Dropdown([], label="Organization")
 
+
 def enable_button(files):
-    """Enables the button if files are uploaded"""
     return gr.update(interactive=bool(files))
 
 
 with gr.Blocks() as app:
     gr.Markdown(project_description)
-
     gr.Markdown("## YourBench Setup")
-    
+
     with gr.Row():
         login_btn = gr.LoginButton()
 
     with gr.Tab("Setup"):
         with gr.Row():
             with gr.Accordion("Hugging Face Settings"):
-                hf_org_dropdown = gr.Dropdown(choices=[], label="Organization", allow_custom_value=True)
+                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")
+                hf_dataset_prefix = gr.Textbox(
+                    label="Dataset Prefix",
+                    value="yourbench",
+                    info="Prefix applied to all datasets",
+                )
 
             with gr.Accordion("Upload documents"):
-                file_input = gr.File(label="Upload text files", file_count="multiple", file_types=[".txt", ".md", ".html"])
+                file_input = gr.File(
+                    label="Upload text files",
+                    file_count="multiple",
+                    file_types=[".txt", ".md", ".html"],
+                )
                 output = gr.Textbox(label="Log")
-                file_input.upload(lambda files: save_files([file.name for file in files]), file_input, output)
+                file_input.upload(
+                    lambda files: save_files([file.name for file in files]),
+                    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)
+        download_button = gr.File(
+            label="Download Config", visible=False, interactive=False
+        )
 
         file_input.change(enable_button, inputs=file_input, outputs=preview_button)
 
@@ -125,18 +137,29 @@ with gr.Blocks() as app:
         )
 
     with gr.Tab("Run Generation"):
-        log_output = gr.Code(label="Log Output", language=None, lines=20, interactive=False)
-        log_timer = gr.Timer(0.05, active=True)
-        log_timer.tick(manager.read_and_get_output, outputs=log_output)
+        with gr.Row():
+            with gr.Accordion("Log Output", open=True):
+                log_output = gr.Code(language=None, lines=20, interactive=False)
+
+            with gr.Accordion("Stages", open=True):
+                stages_table = gr.CheckboxGroup(
+                    choices=STAGES,
+                    value=[],
+                    label="Pipeline Stages Completed",
+                    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(0.05, active=True)
+            status_timer = gr.Timer(1.0, active=True)
             status_timer.tick(manager.is_running, outputs=process_status)
 
         with gr.Row():
             start_button = gr.Button("Start Task")
-            start_button.click(prepare_task, inputs=[hf_org_dropdown])
+            start_button.click(prepare_task, inputs=[login_btn, hf_dataset_prefix])
 
             stop_button = gr.Button("Stop Task")
             stop_button.click(manager.stop_process)
@@ -144,9 +167,7 @@ with gr.Blocks() as app:
             kill_button = gr.Button("Kill Task")
             kill_button.click(manager.kill_process)
 
-    # Citation section at the end
     with gr.Accordion("📜 Citation", open=False):
         gr.Markdown(citation_content)
 
-
-app.launch(allowed_paths=["/app"])
+app.launch(allowed_paths=["/app"])

yourbench_space/config.py

@@ -2,42 +2,46 @@ import yaml
 from loguru import logger
 from yourbench_space.utils import CONFIG_PATH
 
+
 def generate_base_config(hf_org, hf_prefix):
     """Creates the base config dictionary"""
     return {
         "hf_configuration": {
             "token": "$HF_TOKEN",
             "private": True,
-            "hf_organization": hf_org,  
-            "hf_dataset_name": hf_prefix
+            "hf_organization": hf_org,
+            "hf_dataset_name": hf_prefix,
         },
         "local_dataset_dir": "results/",
         "model_list": [
             {
                 "model_name": "meta-llama/Llama-3.3-70B-Instruct",
-                "provider": "huggingface",
-                "base_url": "https://jsq69lxgkhvpnliw.us-east-1.aws.endpoints.huggingface.cloud",
-                "api_key": "$HF_TOKEN",
-                "max_concurrent_requests": 16
+                "provider": "sambanova",
+                "max_concurrent_requests": 32,
+            },
+            {
+                "model_name": "Qwen/Qwen2.5-72B-Instruct",
+                "provider": "nebius",
+                "max_concurrent_requests": 32,
             }
         ],
         "model_roles": {
             "ingestion": ["meta-llama/Llama-3.3-70B-Instruct"],
-            "summarization": ["meta-llama/Llama-3.3-70B-Instruct"],
+            "summarization": ["Qwen/Qwen2.5-72B-Instruct"],
             "single_shot_question_generation": ["meta-llama/Llama-3.3-70B-Instruct"],
             "multi_hop_question_generation": ["meta-llama/Llama-3.3-70B-Instruct"],
-            "answer_generation": ["meta-llama/Llama-3.3-70B-Instruct"],
-            "judge_answers": ["meta-llama/Llama-3.3-70B-Instruct"]
+            "answer_generation": ["Qwen/Qwen2.5-72B-Instruct"],
+            "judge_answers": ["meta-llama/Llama-3.3-70B-Instruct"],
         },
         "pipeline": {
             "ingestion": {
                 "source_documents_dir": "/app/uploaded_files",
                 "output_dir": "/app/ingested",
-                "run": True
+                "run": True,
             },
             "upload_ingest_to_hub": {
                 "source_documents_dir": "/app/ingested",
-                "run": True
+                "run": True,
             },
             "summarization": {"run": True},
             "chunking": {
@@ -46,42 +50,52 @@ def generate_base_config(hf_org, hf_prefix):
                     "l_max_tokens": 128,
                     "tau_threshold": 0.3,
                     "h_min": 2,
-                    "h_max": 4
+                    "h_max": 4,
                 },
-                "run": True
+                "run": True,
             },
             "single_shot_question_generation": {
                 "diversification_seed": "24 year old adult",
-                "run": True
+                "run": True,
             },
             "multi_hop_question_generation": {"run": True},
             "answer_generation": {
                 "question_type": "single_shot",
                 "run": True,
                 "strategies": [
-                    {"name": "zeroshot", "prompt": "ZEROSHOT_QA_USER_PROMPT", "model_name": "meta-llama/Llama-3.3-70B-Instruct"},
-                    {"name": "gold", "prompt": "GOLD_QA_USER_PROMPT", "model_name": "meta-llama/Llama-3.3-70B-Instruct"}
-                ]
+                    {
+                        "name": "zeroshot",
+                        "prompt": "ZEROSHOT_QA_USER_PROMPT",
+                        "model_name": "meta-llama/Llama-3.3-70B-Instruct",
+                    },
+                    {
+                        "name": "gold",
+                        "prompt": "GOLD_QA_USER_PROMPT",
+                        "model_name": "meta-llama/Llama-3.3-70B-Instruct",
+                    },
+                ],
             },
             "judge_answers": {
                 "run": True,
                 "comparing_strategies": [["zeroshot", "gold"]],
                 "chunk_column_index": 0,
-                "random_seed": 42
-            }
-        }
+                "random_seed": 42,
+            },
+        },
     }
 
+
 def save_yaml_file(config):
     """Saves the given config dictionary to a YAML file"""
     with open(CONFIG_PATH, "w") as file:
         yaml.dump(config, file, default_flow_style=False, sort_keys=False)
     return CONFIG_PATH
 
+
 def generate_and_save_config(hf_org, hf_prefix):
     """Generates and saves the YAML configuration file"""
     logger.debug(f"Generating config with org: {hf_org}, prefix: {hf_prefix}")
     config = generate_base_config(hf_org, hf_prefix)
     file_path = save_yaml_file(config)
     logger.success(f"Config saved at: {file_path}")
-    return file_path
+    return file_path

yourbench_space/utils.py

@@ -1,5 +1,6 @@
 import io
 import os
+import re
 import pathlib
 import shutil
 from loguru import logger
@@ -12,6 +13,18 @@ CONFIG_PATH = pathlib.Path("/app/yourbench_config.yml")
 # Ensure the upload directory exists
 UPLOAD_DIRECTORY.mkdir(parents=True, exist_ok=True)
 
+STAGES = [
+    "ingestion",
+    "upload_ingest_to_hub",
+    "summarization",
+    "chunking",
+    "single_shot_question_generation",
+    "multi_hop_question_generation",
+    "answer_generation",
+    "judge_answers",
+]
+
+
 def save_files(files: List[pathlib.Path]) -> str:
     """Save uploaded files to the UPLOAD_DIRECTORY safely"""
     saved_paths = []
@@ -31,7 +44,12 @@ def save_files(files: List[pathlib.Path]) -> str:
         except Exception as e:
             print(f"Error moving file {file}: {e}")
 
-    return f"Files saved to: {', '.join(saved_paths)}" if saved_paths else "No files were saved"
+    return (
+        f"Files saved to: {', '.join(saved_paths)}"
+        if saved_paths
+        else "No files were saved"
+    )
+
 
 class SubprocessManager:
     def __init__(self, command):
@@ -53,24 +71,29 @@ class SubprocessManager:
             text=True,
             bufsize=1,
             start_new_session=True,
-            env=custom_env
+            env=custom_env,
         )
         os.set_blocking(self.process.stdout.fileno(), False)
         logger.info("Started the process")
 
     def read_and_get_output(self):
-        """Read available subprocess output and return the captured output"""
+        """Read subprocess output, capture it, and return log and completed stages."""
         if self.process and self.process.stdout:
             try:
                 while True:
                     line = self.process.stdout.readline()
                     if line:
-                        self.output_stream.write(line)  # Capture in StringIO
+                        self.output_stream.write(line)
                     else:
                         break
             except BlockingIOError:
                 pass
-        return self.output_stream.getvalue()
+
+        current_output = self.output_stream.getvalue()
+        completed_stages = list(
+            set(re.findall(r"Successfully completed stage: (\w+)", current_output))
+        )
+        return current_output, completed_stages
 
     def stop_process(self):
         """Terminate the subprocess."""
@@ -79,9 +102,9 @@ class SubprocessManager:
             return
         logger.info("Sending SIGTERM to the Process")
         self.process.terminate()
-        exit_code = self.process.wait() # Wait for process to terminate
+        exit_code = self.process.wait()  # Wait for process to terminate
         logger.info(f"Process stopped exit code {exit_code}")
-        #return exit_code
+        # return exit_code
 
     def kill_process(self):
         """Forcefully kill the subprocess"""
@@ -90,9 +113,9 @@ class SubprocessManager:
             return
         logger.info("Sending SIGKILL to the Process")
         self.process.kill()
-        exit_code = self.process.wait() # Wait for process to be killed
+        exit_code = self.process.wait()  # Wait for process to be killed
         logger.info(f"Process killed exit code {exit_code}")
-        #return exit_code
+        # return exit_code
 
     def is_running(self):
         """Check if the subprocess is still running"""