[feature] added GSM8k and code refactoring

model/supervised_finetuning/configs/config.yaml

@@ -18,6 +18,18 @@ defaults:
   datasets:
     - webgpt
     - prompt_dialogue
+    - squad_v2
+    - adversarial_qa
+    - trivia_qa_nocontext
+    - xsum
+    - cnn_dailymail
+    - prompt_dialogue
+    - multi_news
+    - scitldr
+    - soda
+    - joke
+    -
+    - joke
   cache_dir: .cache
   loss_fn: CrossEntropyLoss
   eval_size:

model/supervised_finetuning/custom_datasets/__init__.py

@@ -4,8 +4,8 @@ from custom_datasets.summarization import SummarizationDataset
 from sklearn.model_selection import train_test_split
 from torch.utils.data import Subset
 
-QA_DATASETS = ["squad_v2", "adversarial_qa", "trivia_qa_context", "trivia_qa_nocontext"]
-SUMMARIZATION_DATASETS = ["xsum", "cnn_dailymail", "samsum", "multi_news", "scitldr", "billsum", "reddit"]
+QA_DATASETS = ["squad_v2", "adversarial_qa", "trivia_qa_context", "trivia_qa_nocontext", "gsm8k"]
+SUMMARIZATION_DATASETS = ["xsum", "cnn_dailymail", "samsum", "multi_news", "scitldr", "billsum"]
 
 
 def train_val_dataset(dataset, val_split=0.2):
@@ -20,11 +20,13 @@ def get_one_dataset(conf, dataset_name):
 
     if dataset_name in QA_DATASETS:
         train = QADataset(dataset_name, conf.cache_dir, "train")
-        eval = QADataset(dataset_name, conf.cache_dir, "validation")
+        val_name = "validation" if dataset_name not in ["gsm8k"] else "test"
+        eval = QADataset(dataset_name, conf.cache_dir, val_name)
 
     elif dataset_name in SUMMARIZATION_DATASETS:
         train = SummarizationDataset(dataset_name, conf.cache_dir, "train")
-        eval = SummarizationDataset(dataset_name, conf.cache_dir, "validation" if dataset_name != "billsum" else "test")
+        val_name = "validation" if dataset_name not in ["billsum"] else "test"
+        eval = SummarizationDataset(dataset_name, conf.cache_dir, val_name)
     elif dataset_name == "webgpt":
         dataset = WebGPT()
         train, eval = train_val_dataset(dataset, val_split=0.2)

model/supervised_finetuning/custom_datasets/qa_datasets.py

@@ -39,6 +39,10 @@ def index_adversarial_qa(example):
     return example["title"] + ". " + example["context"] + " " + example["question"], example["answers"]["text"][0]
 
 
+def index_gsm8k(example):
+    return example["question"], example["answer"]
+
+
 class QADataset(Dataset):
     def __init__(self, dataset, cache_dir, split):
         if dataset == "squad_v2":
@@ -46,13 +50,19 @@ class QADataset(Dataset):
             self.dataset = load_dataset("squad_v2", cache_dir=cache_dir, split=split)
         elif dataset == "trivia_qa_nocontext":
             self.index_fn = index_trivia_qa_nocontext
-            self.dataset = load_dataset("trivia_qa", "rc.nocontext", split=split)
+            self.dataset = load_dataset("trivia_qa", "rc.nocontext", split=split, cache_dir=cache_dir)
         elif dataset == "trivia_qa_context":
             self.index_fn = index_trivia_qa_context
-            self.dataset = load_dataset("trivia_qa", "rc", split=split)
+            self.dataset = load_dataset("trivia_qa", "rc", split=split, cache_dir=cache_dir)
         elif dataset == "adversarial_qa":
             self.index_fn = index_adversarial_qa
-            self.dataset = load_dataset("adversarial_qa", "adversarialQA", split=split)
+            self.dataset = load_dataset("adversarial_qa", "adversarialQA", split=split, cache_dir=cache_dir)
+        elif dataset == "gsm8k":
+            self.index_fn = index_gsm8k
+            self.dataset = load_dataset("gsm8k", "main", split=split, cache_dir=cache_dir)
+        elif dataset == "adversarial_qa":
+            self.index_fn = index_adversarial_qa
+            self.dataset = load_dataset("adversarial_qa", "adversarialQA", split=split, cache_dir=cache_dir)
         else:
             raise ValueError("Unknown dataset : " + dataset)
 

model/supervised_finetuning/tests/test_datasets.py

@@ -1,12 +1,12 @@
 from argparse import Namespace
 
-from custom_datasets import get_one_dataset
+from custom_datasets import QA_DATASETS, SUMMARIZATION_DATASETS, get_one_dataset
 from custom_datasets.dialogue_collator import DialogueDataCollator
 
 
 def test_all_datasets():
-    qa_base = ["squad_v2", "adversarial_qa", "trivia_qa_context", "trivia_qa_nocontext"]
-    summarize_base = ["scitldr", "xsum", "cnn_dailymail", "samsum", "multi_news", "billsum"]
+    qa_base = QA_DATASETS
+    summarize_base = SUMMARIZATION_DATASETS
     others = ["prompt_dialogue", "webgpt", "soda", "joke"]
 
     config = Namespace(cache_dir=".cache")
@@ -21,21 +21,34 @@ def test_all_datasets():
 
 
 def test_collate_fn():
-    from torch.utils.data import DataLoader
+    from torch.utils.data import ConcatDataset, DataLoader
     from utils import get_tokenizer
 
     config = Namespace(cache_dir=".cache", model_name="Salesforce/codegen-2B-multi")
     tokenizer = get_tokenizer(config)
     collate_fn = DialogueDataCollator(tokenizer, max_length=512)
-    train, eval = get_one_dataset(config, "multi_news")
-    dataloader = DataLoader(train, collate_fn=collate_fn, batch_size=128)
+    qa_base = QA_DATASETS
+    summarize_base = SUMMARIZATION_DATASETS
+    others = ["prompt_dialogue", "webgpt", "soda", "joke", "gsm8k"]
+
+    trains, evals = [], []
+    for dataset_name in others + qa_base + summarize_base:
+        print(dataset_name)
+        train, eval = get_one_dataset(config, dataset_name)
+        trains.append(train)
+        evals.append(eval)
+
+    dataloader = DataLoader(ConcatDataset(trains), collate_fn=collate_fn, batch_size=128)
     for batch in dataloader:
         # print(batch.keys())
         # print(tokenizer.decode(batch['input_ids'][0]))
         # print('-----')
         # print(tokenizer.decode(batch['targets'][0][batch['label_masks'][0]]))
         assert batch["targets"].shape[1] <= 512
+    dataloader = DataLoader(ConcatDataset(evals), collate_fn=collate_fn, batch_size=128)
+    for batch in dataloader:
+        assert batch["targets"].shape[1] <= 512
 
 
 if __name__ == "__main__":
-    test_all_datasets()
+    test_collate_fn()

model/supervised_finetuning/utils.py

@@ -6,8 +6,9 @@ import nltk
 import numpy as np
 import transformers
 import yaml
-from custom_datasets import QA_DATASETS, QA_SPECIAL_TOKENS, SUMMARIZATION_DATASETS, get_one_dataset
+from custom_datasets import QA_DATASETS, SUMMARIZATION_DATASETS, get_one_dataset
 from custom_datasets.dialogue_collator import DialogueDataCollator
+from custom_datasets.qa_datasets import QA_SPECIAL_TOKENS
 from losses import CrossEntropyLoss, PolyLoss
 from models import freeze_top_n_layers, get_specific_model
 from sklearn.model_selection import train_test_split