Simplify Setup tab

yourbench_space/app.py

@@ -1,16 +1,15 @@
 import os
 import sys
+import time  # Needed for file existence check
 import gradio as gr
+import yaml
 from loguru import logger
-from huggingface_hub import HfApi, whoami
+from huggingface_hub import whoami
 
-from yourbench_space.config import generate_base_config, save_config
+from yourbench_space.config import generate_and_save_config
 from yourbench_space.utils import (
     CONFIG_PATH,
     UPLOAD_DIRECTORY,
-    BASE_API_URLS,
-    AVAILABLE_MODELS,
-    DEFAULT_MODEL,
     SubprocessManager,
     save_files,
 )
@@ -23,21 +22,33 @@ logger.add(sys.stderr, level="INFO")
 command = ["uv", "run", "yourbench", f"--config={CONFIG_PATH}"]
 manager = SubprocessManager(command)
 
+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)  # No need to store the return value
+
+    # 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)
 
 def prepare_task(oauth_token: gr.OAuthToken | None, model_token: str):
+    """Prepares and starts the subprocess with environment variables."""
     new_env = os.environ.copy()
-    # Override env token, when running in gradio space
     if oauth_token:
         new_env["HF_TOKEN"] = oauth_token.token
     new_env["MODEL_API_KEY"] = model_token
     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"
-        )
+        print("Please deploy this on Spaces and log in to view the list of available organizations")
         return gr.Dropdown([], label="Organization")
 
     try:
@@ -46,108 +57,49 @@ def update_hf_org_dropdown(oauth_token: gr.OAuthToken | None):
         user_name = user_info.get("name", "Unknown User")
         org_names.insert(0, user_name)
         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") 
+        return gr.Dropdown([], label="Organization")
 
+def enable_button(files):
+    """Enables the button if files are uploaded"""
+    return gr.update(interactive=bool(files))
 
-config_output = gr.Code(label="Generated Config", language="yaml")
-model_name = gr.Dropdown(
-    label="Model Name",
-    value=DEFAULT_MODEL,
-    choices=AVAILABLE_MODELS,
-    allow_custom_value=True,
-)
-base_url = gr.Textbox(
-    label="Model API Base URL",
-    value=BASE_API_URLS["huggingface"],
-    info="Use a custom API base URL for Hugging Face Inference Endpoints",
-)
 
 with gr.Blocks() as app:
-    gr.Markdown("## YourBench Configuration")
+    gr.Markdown("## YourBench Setup")
+    
     with gr.Row():
         login_btn = gr.LoginButton()
 
-    with gr.Tab("Configuration"):
-        with gr.Accordion("Hugging Face"):
-            hf_org_dropdown = gr.Dropdown(
-                list(),
-                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",
-            )
-            private_dataset = gr.Checkbox(
-                label="Private Dataset",
-                value=True,
-                info="Create private datasets (recommended by default)",
-            )
-
-        with gr.Accordion("Model"):
-            model_name.render()
-
-            provider = gr.Radio(
-                ["huggingface", "openrouter", "openai"],
-                value="huggingface",
-                label="Inference Provider",
-            )
-
-            def set_base_url(provider):
-                return gr.Textbox(
-                    label="Model API Base URL", value=BASE_API_URLS.get(provider, "")
-                )
-
-            provider.change(fn=set_base_url, inputs=provider, outputs=base_url)
-            model_api_key = gr.Textbox(label="Model API Key", type="password")
-            base_url.render()
-            max_concurrent_requests = gr.Radio(
-                [8, 16, 32], value=16, label="Max Concurrent Requests"
-            )
-
-        preview_button = gr.Button("Generate New Config")
-        preview_button.click(
-            generate_base_config,
-            inputs=[
-                hf_org_dropdown,
-                hf_dataset_prefix,
-                model_name,
-                provider,
-                base_url,
-                model_api_key,
-                max_concurrent_requests,
-                private_dataset,
-            ],
-            outputs=config_output,
-        )
+    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)
+                app.load(update_hf_org_dropdown, inputs=None, outputs=hf_org_dropdown)
 
-    with gr.Tab("Raw Configuration"):
-        config_output.render()
-        config_output.change(
-            fn=save_config,
-            inputs=[config_output],
-            outputs=[gr.Textbox(label="Save Status")],
-        )
+                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"])
+                output = gr.Textbox(label="Log")
+                file_input.upload(lambda files: save_files([file.name for file in files]), file_input, output)
 
-    with gr.Tab("Files"):
-        file_input = gr.File(
-            label="Upload text files",
-            file_count="multiple",
-            file_types=[".txt", ".md", ".html"],
+        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],
+            outputs=[log_message, download_button],
         )
-        output = gr.Textbox(label="Log")
-        file_input.upload(save_files, file_input, output)
 
     with gr.Tab("Run Generation"):
-        log_output = gr.Code(
-            label="Log Output", language=None, lines=20, interactive=False
-        )
+        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)
 
@@ -158,7 +110,7 @@ with gr.Blocks() as app:
 
         with gr.Row():
             start_button = gr.Button("Start Task")
-            start_button.click(prepare_task, inputs=[model_api_key])
+            start_button.click(prepare_task, inputs=[hf_org_dropdown])
 
             stop_button = gr.Button("Stop Task")
             stop_button.click(manager.stop_process)
@@ -166,4 +118,4 @@ with gr.Blocks() as app:
             kill_button = gr.Button("Kill Task")
             kill_button.click(manager.kill_process)
 
-app.launch()
+app.launch(allowed_paths=["/app"])

yourbench_space/config.py

@@ -1,49 +1,43 @@
 import yaml
+from loguru import logger
 from yourbench_space.utils import CONFIG_PATH
 
-def generate_base_config(
-    hf_org,
-    hf_dataset_name,
-    model_name,
-    provider,
-    base_url,
-    model_api_key,
-    max_concurrent_requests,
-    private_dataset,
-):
-    config = {
+def generate_base_config(hf_org, hf_prefix):
+    """Creates the base config dictionary"""
+    return {
         "hf_configuration": {
             "token": "$HF_TOKEN",
-            "private": private_dataset,
-            "hf_organization": hf_org,
-            "hf_dataset_name": hf_dataset_name,
+            "private": True,
+            "hf_organization": hf_org,  
+            "hf_dataset_name": hf_prefix
         },
+        "local_dataset_dir": "results/",
         "model_list": [
             {
-                "model_name": model_name,
-                "provider": provider,
-                "base_url": base_url,
-                "api_key": "$MODEL_API_KEY",
-                "max_concurrent_requests": max_concurrent_requests,
+                "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
             }
         ],
         "model_roles": {
-            "ingestion": [model_name],
-            "summarization": [model_name],
-            "single_shot_question_generation": [model_name],
-            "multi_hop_question_generation": [model_name],
-            "answer_generation": [model_name],
-            "judge_answers": [model_name],
+            "ingestion": ["meta-llama/Llama-3.3-70B-Instruct"],
+            "summarization": ["meta-llama/Llama-3.3-70B-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"]
         },
         "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": {
@@ -52,34 +46,42 @@ def generate_base_config(
                     "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": model_name},
-                    {"name": "gold", "prompt": "GOLD_QA_USER_PROMPT", "model_name": model_name},
-                ],
+                    {"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
+            }
+        }
     }
-    return yaml.dump(config, sort_keys=False)
 
-def save_config(yaml_text):
+def save_yaml_file(config):
+    """Saves the given config dictionary to a YAML file"""
     with open(CONFIG_PATH, "w") as file:
-        file.write(yaml_text)
-    return "✅ Config saved!"
+        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

yourbench_space/utils.py

@@ -4,26 +4,34 @@ import pathlib
 import shutil
 from loguru import logger
 import subprocess
+from typing import List
 
 UPLOAD_DIRECTORY = pathlib.Path("/app/uploaded_files")
 CONFIG_PATH = pathlib.Path("/app/yourbench_config.yml")
 
-AVAILABLE_MODELS = [
-    "mistralai/Mistral-Small-24B-Instruct-2501",
-    "meta-llama/Llama-3.3-70B-Instruct",
-]
-DEFAULT_MODEL = AVAILABLE_MODELS[0]
+# Ensure the upload directory exists
+UPLOAD_DIRECTORY.mkdir(parents=True, exist_ok=True)
 
-BASE_API_URLS = {
-    "huggingface": "https://router.huggingface.co/hf-inference/v1",
-    "openrouter":  "https://openrouter.ai/api/v1",
-    "openai": "https://api.openai.com/v1/",
-}
+def save_files(files: List[pathlib.Path]) -> str:
+    """Save uploaded files to the UPLOAD_DIRECTORY safely"""
+    saved_paths = []
 
-def save_files(files: list[str]):
-    saved_paths = [shutil.move(str(pathlib.Path(file)), str(UPLOAD_DIRECTORY / pathlib.Path(file).name)) for file in files]
-    return f"Files saved to: {', '.join(saved_paths)}"
+    for file in files:
+        try:
+            source_path = pathlib.Path(file)
+            destination_path = UPLOAD_DIRECTORY / source_path.name
 
+            if not source_path.exists():
+                print(f"File not found: {source_path}")
+                continue  # Skip missing files
+
+            shutil.move(str(source_path), str(destination_path))
+            saved_paths.append(str(destination_path))
+
+        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"
 
 class SubprocessManager:
     def __init__(self, command):
@@ -51,7 +59,7 @@ class SubprocessManager:
         logger.info("Started the process")
 
     def read_and_get_output(self):
-        """Read available subprocess output and return the captured output."""
+        """Read available subprocess output and return the captured output"""
         if self.process and self.process.stdout:
             try:
                 while True:
@@ -76,7 +84,7 @@ class SubprocessManager:
         #return exit_code
 
     def kill_process(self):
-        """Forcefully kill the subprocess."""
+        """Forcefully kill the subprocess"""
         if not self.is_running():
             logger.info("Process is not running")
             return
@@ -87,5 +95,5 @@ class SubprocessManager:
         #return exit_code
 
     def is_running(self):
-        """Check if the subprocess is still running."""
+        """Check if the subprocess is still running"""
         return self.process and self.process.poll() is None