minifix

yourbench_space/utils.py

@@ -80,7 +80,7 @@ def update_dataset(stages: list, hf_org: str, hf_prefix: str, oauth_token: gr.OA
         ingestion_df = pd.DataFrame([next(iter(ingestion_ds)) for _ in range(1)]) # only one row
     if "summarization" in stages:
         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)])
+        summarization_df = pd.DataFrame([next(iter(summarization_ds)) for _ in range(1)])
     if "single_shot_question_generation" in stages:
         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)])