wassname/Open-Assistant
1
0
Fork 0
mirror of https://github.com/wassname/Open-Assistant.git synced 2026-09-09 11:15:08 +08:00
Code Issues Packages Projects Releases Wiki Activity

[feature] added GSM8k and code refactoring

Browse Source
This commit is contained in:
theblackcat102
2023-01-14 06:24:47 +00:00
parent 3966024871
commit 1546111094
5 changed files with 53 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
model/supervised_finetuning/tests/test_datasets.py
+20 -7
View File
@@ -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()