Fix tests

EleutherAI · lauritowal · Mar 28, 2023 · Mar 9, 2023 · Mar 9, 2023 · Mar 10, 2023
commit fe94c22c83579e093b610c317ede48524403b211
diff --git a/tests/distilgpt2_copa_cfg.yaml b/tests/distilgpt2_copa_cfg.yaml
@@ -2,7 +2,8 @@ layers: []
 model: distilgpt2
 prompts:
  balance: true
- dataset: super_glue copa
+ datasets:
+ - "super_glue copa"
  label_column: null
  max_examples:
  - 2

diff --git a/tests/distilgpt2_dbpedia_cfg.yaml b/tests/distilgpt2_dbpedia_cfg.yaml
@@ -2,7 +2,8 @@ layers: []
 model: distilgpt2
 prompts:
  balance: true
- dataset: dbpedia_14
+ datasets:
+ - "dbpedia_14"
  label_column: null
  max_examples:
  - 3

diff --git a/tests/test_prompt_dataset.py b/tests/test_prompt_dataset.py
@@ -7,7 +7,7 @@
 def test_prompt_dataset_getitem_boolq():
  def test_prompt_dataset_getitem(cfg: ExtractionConfig, split: str):
  prompt_ds = PromptDataset(cfg.prompts, rank=0, world_size=1, split=split)
- ds_name, _, config_name = cfg.prompts.dataset.partition(" ")
+ ds_name, _, config_name = cfg.prompts.datasets[0].partition(" ")
 
  prompter = DatasetTemplates(ds_name, config_name or None)
  assert len(prompt_ds) == cfg.prompts.max_examples[-1]