uvos/QRotaryTraining
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add chat datamodules

Browse Source
This commit is contained in:
uvos
2024-07-20 21:46:34 +02:00
parent 0b39ba0843
commit 2f35689355
2 changed files with 148 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files

123
datamodules.py
View File

@ -7,22 +7,23 @@ import transformers
import os
from dataclasses import dataclass
from torch.nn.utils.rnn import pad_sequence
from tqdm import tqdm
from arguments import DataArguments
from arguments import DataArguments, DatasetType
IGNORE_INDEX = -100
def group_texts(examples, block_size: int):
def group_texts(examples, source_max_len: int):
# Concatenate all texts.
concatenated_examples = {k: list(itertools.chain(*examples[k])) for k in examples.keys()}
total_length = len(concatenated_examples[list(examples.keys())[0]])
# We drop the small remainder, we could add padding if the model supported it instead of this drop, you can
# customize this part to your needs.
if total_length >= block_size:
total_length = (total_length // block_size) * block_size
if total_length >= source_max_len:
total_length = (total_length // source_max_len) * source_max_len
# Split by chunks of max_len.
result = {k: [t[i: i + block_size] for i in range(0, total_length, block_size)] for k, t in concatenated_examples.items()}
result = {k: [t[i: i + source_max_len] for i in range(0, total_length, source_max_len)] for k, t in concatenated_examples.items()}
result["labels"] = result["input_ids"].copy()
return result
@ -199,14 +200,15 @@ def create_data_module_hub(tokenizer: transformers.PreTrainedTokenizer, data_arg
)
def create_data_module(tokenizer: transformers.PreTrainedTokenizer, data_args: DataArguments, do_train: bool, do_eval: bool, do_predict: bool) -> typing.Dict:
def create_data_module_txt(tokenizer: transformers.PreTrainedTokenizer,
data_args: DataArguments, do_train: bool, do_eval: bool, do_predict: bool) -> typing.Dict:
try:
dataset = datasets.load_dataset('text', data_files={'train': [data_args.dataset]})
except FileNotFoundError as ex:
raise ValueError(f"Error loading dataset from {data_args.dataset}, {ex}")
if data_args.block_size > tokenizer.model_max_length:
raise ValueError(f"Block size of {data_args.block_size} is larger than the maximum size supported by the model: {tokenizer.model_max_length}")
if data_args.source_max_len > tokenizer.model_max_length:
raise ValueError(f"Max source length of {data_args.source_max_len} is larger than the maximum size supported by the model: {tokenizer.model_max_length}")
def add_newline_fn(example):
example['text'] = example['text'] + '\n'
@ -219,9 +221,7 @@ def create_data_module(tokenizer: transformers.PreTrainedTokenizer, data_args: D
eval_dataset = dataset['eval']
else:
print('Splitting train dataset in train and validation according to `eval_dataset_size`')
dataset = dataset['train'].train_test_split(
test_size=data_args.eval_dataset_size, shuffle=True, seed=42
)
dataset = dataset['train'].train_test_split(test_size=data_args.eval_dataset_size, shuffle=False)
eval_dataset = dataset['test']
if 'train' in dataset:
@ -233,14 +233,14 @@ def create_data_module(tokenizer: transformers.PreTrainedTokenizer, data_args: D
lambda example: tokenizer(example['text']),
batched=True,
remove_columns='text',
num_proc=32,
num_proc=os.cpu_count(),
load_from_cache_file=True)
train_dataset_tokenized = train_dataset_tokenized.map(
lambda example: group_texts(example, data_args.block_size),
lambda example: group_texts(example, data_args.source_max_len),
batched=True,
num_proc=max(1, min(os.cpu_count(), int(len(train_dataset_tokenized['input_ids']) / (data_args.block_size * 10)))),
num_proc=max(1, min(os.cpu_count(), int(len(train_dataset_tokenized['input_ids']) / (data_args.source_max_len * 10)))),
load_from_cache_file=True,
desc=f"Grouping texts in chunks of {data_args.block_size}")
desc=f"Grouping texts in chunks of {data_args.source_max_len}")
eval_dataset_tokenized = None
if eval_dataset is not None:
@ -248,18 +248,18 @@ def create_data_module(tokenizer: transformers.PreTrainedTokenizer, data_args: D
lambda example: tokenizer(example['text']),
batched=True,
remove_columns='text',
num_proc=32)
num_proc=os.cpu_count())
eval_dataset_tokenized = eval_dataset_tokenized.map(
lambda example: group_texts(example, data_args.block_size),
lambda example: group_texts(example, data_args.source_max_len),
batched=True,
num_proc=max(1, min(os.cpu_count(), int(len(eval_dataset_tokenized['input_ids']) / (data_args.block_size * 10)))),
num_proc=max(1, min(os.cpu_count(), int(len(eval_dataset_tokenized['input_ids']) / (data_args.source_max_len * 10)))),
load_from_cache_file=True,
desc=f"Grouping texts in chunks of {data_args.block_size}")
desc=f"Grouping texts in chunks of {data_args.source_max_len}")
for ids in train_dataset_tokenized['input_ids']:
assert len(ids) == data_args.block_size
assert len(ids) == data_args.source_max_len
for ids in eval_dataset_tokenized['input_ids']:
assert len(ids) == data_args.block_size
assert len(ids) == data_args.source_max_len
return dict(
train_dataset=train_dataset_tokenized if do_train else None,
@ -267,3 +267,84 @@ def create_data_module(tokenizer: transformers.PreTrainedTokenizer, data_args: D
predict_dataset=eval_dataset_tokenized if do_predict else None,
data_collator=transformers.default_data_collator
)
def create_data_module_chat(tokenizer, data_args, do_train, do_eval, do_predict):
try:
dataset = datasets.Dataset.from_json(path_or_paths=data_args.dataset)
except FileNotFoundError as ex:
raise ValueError(f"Error loading dataset from {data_args.dataset}, {ex}")
if data_args.dataset_chat_template is not None:
tokenizer.chat_template = data_args.dataset_chat_template
target_len = data_args.source_max_len * 0.5
grouped_chats = list()
last_len = 0
for row in tqdm(dataset, desc="Grouping chat messages"):
content_length = len(tokenizer(row['content'])['input_ids'])
if last_len + content_length <= target_len and len(grouped_chats) > 0:
grouped_chats[-1]['chat'].append(row)
last_len += content_length
else:
last_len = 0
grouped_chats.append({'chat': [row]})
dataset = datasets.Dataset.from_list(grouped_chats)
dataset = dataset.map(lambda x: {"text": tokenizer.apply_chat_template(x["chat"], tokenize=False, add_generation_prompt=False)})
dataset.remove_columns('chat')
eval_dataset = None
if do_eval or do_predict:
print('Splitting train dataset in train and validation according to `eval_dataset_size`')
dataset_split = dataset.train_test_split(test_size=data_args.eval_dataset_size, shuffle=True)
train_dataset = dataset_split["train"]
eval_dataset = dataset_split["test"]
data_collator = DataCollatorForCausalLMText(
tokenizer=tokenizer,
max_len=data_args.source_max_len,
)
return dict(
train_dataset=train_dataset if do_train else None,
eval_dataset=eval_dataset,
predict_dataset=eval_dataset,
data_collator=data_collator
)
def get_data_loaders(tokenizer, data_args: DataArguments, batch_size: int, eval_batch_size: int,
do_train: bool, do_eval: bool, do_predict: bool = False):
data_type = DatasetType.from_string(data_args.dataset_type)
if data_type == DatasetType.S2S:
print("Loading dataset in s2s mode")
data_module = create_data_module_s2s(tokenizer, data_args, do_train, do_eval, do_predict)
elif data_type == DatasetType.HUB:
print("Loading dataset from hub, expecting alpaca style")
data_module = create_data_module_hub(tokenizer, data_args, do_train, do_eval, do_predict)
elif data_type == DatasetType.TEXT:
print("Loading dataset in txt mode")
data_module = create_data_module_txt(tokenizer, data_args, do_train, do_eval, do_predict)
elif data_type == DatasetType.CHAT:
print("Loading dataset in chat mode")
data_module = create_data_module_chat(tokenizer, data_args, do_train, do_eval, do_predict)
else:
raise RuntimeError("Unkown dataset type")
train_dataloader = None
eval_dataloader = None
if do_train:
train_dataloader = torch.utils.data.DataLoader(
data_module['train_dataset'],
shuffle=True,
collate_fn=data_module['data_collator'],
batch_size=batch_size
)
if do_eval:
eval_dataloader = torch.utils.data.DataLoader(
data_module['eval_dataset'],
shuffle=True,
collate_fn=data_module['data_collator'],
batch_size=eval_batch_size
)
return train_dataloader, eval_dataloader