update oauth to follow the session

yourbench_space/app.py

@@ -108,7 +108,7 @@ def run_evaluation_pipeline(oauth_token: gr.OAuthToken | None, org_name, eval_na
     eval_ds_name = f"{org_name}/{eval_name}"
     # Test dataset existence
     try:
-        load_dataset(eval_ds_name, streaming=True)
+        load_dataset(eval_ds_name, streaming=True, token=oauth_token.token)
     except Exception as e:
         print(f"Error while loading the dataset: {e}")
         return
@@ -232,7 +232,7 @@ 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_name], outputs=[ingestion_df, summarization_df, single_hop, answers_df]
+                        update_dataset, inputs=[stages_table, hf_org_dropdown, hf_dataset_name, login_btn], outputs=[ingestion_df, summarization_df, single_hop, answers_df]
                     )
 
             log_timer = gr.Timer(1.0, active=True)

yourbench_space/utils.py

@@ -52,7 +52,7 @@ def save_files(uuid: str, files: List[pathlib.Path]) -> str:
         else "No files were saved"
     )
 
-def update_dataset(stages, hf_org, hf_prefix):
+def update_dataset(stages, hf_org, hf_prefix, oauth_token: gr.OAuthToken):
     """
     Updates the dataset based on the provided stages and dataset configuration.
     """
@@ -66,16 +66,16 @@ def update_dataset(stages, hf_org, hf_prefix):
 
     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").select_columns("document_text")
+        ingestion_ds = load_dataset(dataset_name, name="ingested", split="train", token=oauth_token.token).select_columns("document_text")
         ingestion_df = pd.DataFrame(ingestion_ds[0]) # only one row
     if "summarization" in stages:
-        summarization_ds = load_dataset(dataset_name, name="summarization", split="train", streaming=True).select_columns(['raw_document_summary', 'document_summary', 'summarization_model'])
+        summarization_ds = load_dataset(dataset_name, name="summarization", split="train", streaming=True, token=oauth_token.token).select_columns(['raw_document_summary', 'document_summary', 'summarization_model'])
         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_ds = load_dataset(dataset_name, name="single_shot_question_generation", split="train", streaming=True, token=oauth_token.token)
         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_ds = load_dataset(dataset_name, name="answer_generation", split="train", streaming=True, token=oauth_token.token)
         answers_df = pd.DataFrame([next(iter(answers_ds)) for _ in range(5)])
     
     return (ingestion_df, summarization_df, single_hop_df, answers_df)