From 154611109465e3ca1c22efd08ce8bb15e52f519d Mon Sep 17 00:00:00 2001 From: theblackcat102 Date: Sat, 14 Jan 2023 06:24:47 +0000 Subject: [PATCH] [feature] added GSM8k and code refactoring --- .../supervised_finetuning/configs/config.yaml | 12 +++++++++ .../custom_datasets/__init__.py | 10 ++++--- .../custom_datasets/qa_datasets.py | 16 ++++++++--- .../tests/test_datasets.py | 27 ++++++++++++++----- model/supervised_finetuning/utils.py | 3 ++- 5 files changed, 53 insertions(+), 15 deletions(-) diff --git a/model/supervised_finetuning/configs/config.yaml b/model/supervised_finetuning/configs/config.yaml index b0684bda..0440201a 100644 --- a/model/supervised_finetuning/configs/config.yaml +++ b/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: diff --git a/model/supervised_finetuning/custom_datasets/__init__.py b/model/supervised_finetuning/custom_datasets/__init__.py index 6b7906f7..e293af3d 100644 --- a/model/supervised_finetuning/custom_datasets/__init__.py +++ b/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) diff --git a/model/supervised_finetuning/custom_datasets/qa_datasets.py b/model/supervised_finetuning/custom_datasets/qa_datasets.py index d0fcffda..eed9c644 100644 --- a/model/supervised_finetuning/custom_datasets/qa_datasets.py +++ b/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) diff --git a/model/supervised_finetuning/tests/test_datasets.py b/model/supervised_finetuning/tests/test_datasets.py index 721fdfd3..c9363303 100644 --- a/model/supervised_finetuning/tests/test_datasets.py +++ b/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() diff --git a/model/supervised_finetuning/utils.py b/model/supervised_finetuning/utils.py index f598dde1..85fb86db 100644 --- a/model/supervised_finetuning/utils.py +++ b/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