Update config

yourbench_space/app.py

@@ -33,18 +33,23 @@ def prepare_task(oauth_token: gr.OAuthToken | None, model_token: str):
     manager.start_process(custom_env=new_env)
 
 
-def update_hf_org_dropdown(oauth_token: gr.OAuthToken | None) -> str:
+def update_hf_org_dropdown(oauth_token: gr.OAuthToken | None):
     if oauth_token is None:
         print(
             "Please, deploy this on Spaces and log in to view the list of available organizations"
         )
-        return list()
-    user_info = whoami(oauth_token.token)
+        return gr.Dropdown([], label="Organization")
 
-    org_names = [org["name"] for org in user_info["orgs"]]
-    user_name = user_info["name"]
-    org_names.insert(0, user_name)
-    return gr.Dropdown(org_names, value=user_name, label="Organization")
+    try:
+        user_info = whoami(oauth_token.token)
+        org_names = [org["name"] for org in user_info.get("orgs", [])]
+        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") 
 
 
 config_output = gr.Code(label="Generated Config", language="yaml")
@@ -60,74 +65,6 @@ base_url = gr.Textbox(
     info="Use a custom API base URL for Hugging Face Inference Endpoints",
 )
 
-
-def make_models(model_name=None):
-    if model_name is None:
-        model_name = DEFAULT_MODEL
-
-    ingestion_model = gr.Dropdown(
-        label="Model for ingestion",
-        choices=AVAILABLE_MODELS,
-        value=model_name,
-        interactive=False,
-        allow_custom_value=True,
-    )
-    summarization_model = gr.Dropdown(
-        label="Model for summarization",
-        choices=AVAILABLE_MODELS,
-        value=model_name,
-        interactive=False,
-        allow_custom_value=True,
-    )
-    single_shot_question_generation_model = gr.Dropdown(
-        label="Model for single shot question generation",
-        choices=AVAILABLE_MODELS,
-        value=model_name,
-        interactive=False,
-        allow_custom_value=True,
-    )
-    multi_hop_question_generation_model = gr.Dropdown(
-        label="Model for multi hop question generation",
-        choices=AVAILABLE_MODELS,
-        value=model_name,
-        interactive=False,
-        allow_custom_value=True,
-    )
-    answer_generation_model = gr.Dropdown(
-        label="Model for answer generation",
-        choices=AVAILABLE_MODELS,
-        value=model_name,
-        interactive=False,
-        allow_custom_value=True,
-    )
-    judge_answers_model = gr.Dropdown(
-        label="Model for answer judging",
-        choices=AVAILABLE_MODELS,
-        value=model_name,
-        interactive=False,
-        allow_custom_value=True,
-    )
-
-    return [
-        ingestion_model,
-        summarization_model,
-        single_shot_question_generation_model,
-        multi_hop_question_generation_model,
-        answer_generation_model,
-        judge_answers_model,
-    ]
-
-
-(
-    ingestion_model,
-    summarization_model,
-    single_shot_question_generation_model,
-    multi_hop_question_generation_model,
-    answer_generation_model,
-    judge_answers_model,
-) = make_models()
-
-
 with gr.Blocks() as app:
     gr.Markdown("## YourBench Configuration")
     with gr.Row():
@@ -155,19 +92,6 @@ with gr.Blocks() as app:
 
         with gr.Accordion("Model"):
             model_name.render()
-            # TODO handle this better
-            model_name.change(
-                make_models,
-                inputs=[model_name],
-                outputs=[
-                    ingestion_model,
-                    summarization_model,
-                    single_shot_question_generation_model,
-                    multi_hop_question_generation_model,
-                    answer_generation_model,
-                    judge_answers_model,
-                ],
-            )
 
             provider = gr.Radio(
                 ["huggingface", "openrouter", "openai"],
@@ -187,32 +111,18 @@ with gr.Blocks() as app:
                 [8, 16, 32], value=16, label="Max Concurrent Requests"
             )
 
-        with gr.Accordion("Stages"):
-            ingestion_model.render()
-            summarization_model.render()
-            single_shot_question_generation_model.render()
-            multi_hop_question_generation_model.render()
-            answer_generation_model.render()
-            judge_answers_model.render()
-
         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,
-                hf_dataset_prefix,
                 private_dataset,
-                ingestion_model,
-                summarization_model,
-                single_shot_question_generation_model,
-                multi_hop_question_generation_model,
-                answer_generation_model,
-                judge_answers_model,
             ],
             outputs=config_output,
         )

yourbench_space/config.py

@@ -1,31 +1,24 @@
-
 import yaml
 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,
-    hf_dataset_prefix,
     private_dataset,
-    ingestion_model,
-    summarization_model,
-    single_shot_question_generation_model,
-    multi_hop_question_generation_model,
-    answer_generation_model,
-    judge_answers_model,
 ):
     config = {
         "hf_configuration": {
             "token": "$HF_TOKEN",
             "private": private_dataset,
             "hf_organization": hf_org,
+            "hf_dataset_name": hf_dataset_name,
         },
+        "local_dataset_dir": "results/",
         "model_list": [
             {
                 "model_name": model_name,
@@ -36,38 +29,25 @@ def generate_base_config(
             }
         ],
         "model_roles": {
-            role: [model_name]
-            for role in [
-                "ingestion",
-                "summarization",
-                "single_shot_question_generation",
-                "multi_hop_question_generation",
-                "answer_generation",
-                "judge_answers",
-            ]
+            "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],
         },
-        "inference_config": {"max_concurrent_requests": 16},
         "pipeline": {
             "ingestion": {
-                "source_documents_dir": "/app/uploaded_files",
-                "output_dir": "/app/ingested",
+                "source_documents_dir": "/app/example/raw",
+                "output_dir": "/app/example/ingested",
                 "run": True,
             },
             "upload_ingest_to_hub": {
-                "source_documents_dir": "/app/ingested",
-                "hub_dataset_name": f"{hf_dataset_prefix}_ingested_documents",
-                "run": True,
-            },
-            "summarization": {
-                "source_dataset_name": f"{hf_dataset_prefix}_ingested_documents",
-                "output_dataset_name": f"{hf_dataset_prefix}_summaries",
-                "concat_existing_dataset": False,
+                "source_documents_dir": "/app/example/ingested",
                 "run": True,
             },
+            "summarization": {"run": True},
             "chunking": {
-                "source_dataset_name": f"{hf_dataset_prefix}_summaries",
-                "output_dataset_name": f"{hf_dataset_prefix}_chunked_documents",
-                "concat_existing_dataset": False,
                 "chunking_configuration": {
                     "l_min_tokens": 64,
                     "l_max_tokens": 128,
@@ -78,50 +58,28 @@ def generate_base_config(
                 "run": True,
             },
             "single_shot_question_generation": {
-                "source_dataset_name": f"{hf_dataset_prefix}_chunked_documents",
-                "output_dataset_name": f"{hf_dataset_prefix}_single_shot_questions",
                 "diversification_seed": "24 year old adult",
-                "concat_existing_dataset": False,
-                "run": True,
-            },
-            "multi_hop_question_generation": {
-                "source_dataset_name": f"{hf_dataset_prefix}_chunked_documents",
-                "output_dataset_name": f"{hf_dataset_prefix}_multi_hop_questions",
-                "concat_existing_dataset": False,
                 "run": True,
             },
+            "multi_hop_question_generation": {"run": True},
             "answer_generation": {
-                "question_dataset_name": f"{hf_dataset_prefix}_single_shot_questions",
-                "output_dataset_name": f"{hf_dataset_prefix}_answered_questions",
-                "concat_existing_dataset": False,
+                "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": model_name},
+                    {"name": "gold", "prompt": "GOLD_QA_USER_PROMPT", "model_name": model_name},
                 ],
-                "run": True,
             },
             "judge_answers": {
-                "source_judge_dataset_name": f"{hf_dataset_prefix}_answered_questions",
-                "output_judged_dataset_name": f"{hf_dataset_prefix}_judged_comparisons",
-                "concat_existing_dataset": False,
+                "run": True,
                 "comparing_strategies": [["zeroshot", "gold"]],
                 "chunk_column_index": 0,
                 "random_seed": 42,
-                "run": True,
             },
         },
     }
     return yaml.dump(config, sort_keys=False)
 
-
 def save_config(yaml_text):
     with open(CONFIG_PATH, "w") as file:
         file.write(yaml_text)