only show relevant columns as eval goes on

yourbench_space/utils.py

@@ -67,12 +67,13 @@ def update_dataset(stages, hf_org, hf_prefix):
 
     # TODO: add cache dir
     # Will be able to group everything in one pass once the names get homogeneized
+    # TODO: make sure the questions are loaded with a set
     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", streaming=True)
+        ingestion_ds = load_dataset(dataset_name, name="ingested", split="train").select_columns("document_text")
         ingestion_df = pd.DataFrame([next(iter(ingestion_ds)) for _ in range(5)])
     if "summarization" in stages:
-        summarization_ds = load_dataset(dataset_name, name="summarization", split="train", streaming=True)
+        summarization_ds = load_dataset(dataset_name, name="summarization", split="train", streaming=True).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)