diff --git a/src/ctx_to_lora/data/processing.py b/src/ctx_to_lora/data/processing.py index 6857751..4a081bf 100644 --- a/src/ctx_to_lora/data/processing.py +++ b/src/ctx_to_lora/data/processing.py @@ -419,13 +419,7 @@ def load_and_process_dataset( is_pretrain: bool, streaming: bool, num_proc: int, - ds_path: str, ): - if path.exists(ds_path) and split == "train": - # load the cached ds - logger.info(f"Loaded processed dataset from {ds_path}") - return datasets.load_from_disk(ds_path) - logger.info(f"Loading dataset {ds_name} with split {split}...") try: ds_kwargs = get_ds_kwargs(ds_name, split) @@ -500,7 +494,6 @@ def load_and_process_dataset( batch_size=100_000, # num_proc=num_proc, ) - ds.save_to_disk(ds_path, num_proc=num_proc) return ds @@ -524,9 +517,10 @@ def get_tokenized_dataset( assert not use_kl_loss, "KL loss is deprecated" if add_repeat_prompt: assert repeat_prob > 0, f"add_repeat_prompt is set but repeat_prob = 0" - logger.debug(f"Loading dataset {ds_name} with split {split}...") + logger.info(f"Loading dataset {ds_name} with split {split}...") need_ctx_ids = not add_ctx_to_chat and bool(ctx_model_max_len) is_pretrain = "pretrain" in ds_name + is_paraphrased = "aug_pretrain" in ds_name load_and_process_kwargs = dict( ds_name=ds_name, @@ -537,35 +531,48 @@ def get_tokenized_dataset( is_pretrain=is_pretrain, streaming=streaming, ) - ds_hash = hashlib.sha256(json.dumps(load_and_process_kwargs).encode()).hexdigest() - ds_path = f"{TRANSFORMED_DATA_DIR}/{ds_hash}" - num_proc = None if streaming and split == "train" else 8 + tokenize_kwargs = dict( + base_model_max_len=base_model_max_len, + tokenizer=tokenizer, + tokenizer_kwargs=tokenizer_kwargs, + ctx_model_max_len=ctx_model_max_len, + ctx_tokenizer=ctx_tokenizer, + ctx_tokenizer_kwargs=ctx_tokenizer_kwargs, + add_ctx_to_chat=add_ctx_to_chat, + use_kl_loss=use_kl_loss, + need_ctx_ids=need_ctx_ids, + is_pretrain=is_pretrain, + is_paraphrased=is_paraphrased, + split=split, + set_format=set_format, + ) + all_kwargs = {**load_and_process_kwargs, **tokenize_kwargs} + kwargs_str = json.dumps(all_kwargs) + logger.debug(f"Tokenizing dataset with kwargs: {kwargs_str}") + ds_hash = hashlib.sha256(kwargs_str.encode()).hexdigest() + ds_path = f"{TRANSFORMED_DATA_DIR}/{ds_hash}" + + if path.exists(ds_path) and split == "train": + # load the cached ds + logger.info(f"Loaded tokenized dataset from {ds_path}") + tokenized_ds = datasets.load_from_disk(ds_path) + return tokenized_ds + + num_proc = None if streaming and split == "train" else 8 ds = load_and_process_dataset( **load_and_process_kwargs, - ds_path=ds_path, num_proc=num_proc, ) - if split == "train": - # force loading from disk when processing for the first time - ds = datasets.load_from_disk(ds_path) - + logger.info(f"Constructing and tokenizing {ds_name} with {split} split...") tokenized_ds = construct_and_tokenize_ctx_qa( - base_model_max_len, - tokenizer, - tokenizer_kwargs, - ctx_model_max_len, - ctx_tokenizer, - ctx_tokenizer_kwargs, - add_ctx_to_chat, - use_kl_loss, - need_ctx_ids, - is_pretrain, - ds, - split, - set_format, - num_proc, + ds=ds, + num_proc=num_proc, + **tokenize_kwargs, ) + if split == "train": + tokenized_ds.save_to_disk(ds_path, num_proc=num_proc) + return tokenized_ds @@ -580,45 +587,34 @@ def construct_and_tokenize_ctx_qa( use_kl_loss, need_ctx_ids, is_pretrain, + is_paraphrased, ds, split, set_format=None, num_proc=None, ): - kwargs = dict( - base_model_max_len=base_model_max_len, - tokenizer=repr(tokenizer), - tokenizer_kwargs=json.dumps(tokenizer_kwargs), - ctx_model_max_len=ctx_model_max_len, - ctx_tokenizer=repr(ctx_tokenizer), - ctx_tokenizer_kwargs=json.dumps(ctx_tokenizer_kwargs), - add_ctx_to_chat=add_ctx_to_chat, - use_kl_loss=use_kl_loss, - need_ctx_ids=need_ctx_ids, - is_pretrain=is_pretrain, - ds=ds._fingerprint, - split=split, - set_format=set_format, - ) - kwargs_str = json.dumps(kwargs) - logger.debug(f"Tokenizing dataset with kwargs: {kwargs_str}") - ds_hash = hashlib.sha256(kwargs_str.encode()).hexdigest() - ds_path = f"{TRANSFORMED_DATA_DIR}/{ds_hash}" - if path.exists(ds_path) and split == "train": - # load the cached ds - logger.info(f"Loaded tokenized dataset from {ds_path}") - ds = datasets.load_from_disk(ds_path) - return ds if is_pretrain: - ds = ds.map( - build_intx_pretrain, - # batched=True, - # batch_size=100_000, - # # num_proc=num_proc, - ) + if is_paraphrased: + ds = ds.map( + build_paraphrase_pretrain, + batched=True, + batch_size=100_000, + remove_columns=["variation"], + ) + else: + ds = ds.map( + build_intx_pretrain, + # batched=True, + # batch_size=100_000, + # # num_proc=num_proc, + ) tokenized_ds = ds.map( tokenize_pretrain, - fn_kwargs={"tokenizer": tokenizer, "tokenizer_kwargs": tokenizer_kwargs}, + fn_kwargs={ + "tokenizer": tokenizer, + "tokenizer_kwargs": tokenizer_kwargs, + "add_eos_token": is_paraphrased, + }, batched=True, batch_size=100_000, # # num_proc=num_proc, @@ -766,7 +762,9 @@ def construct_and_tokenize_ctx_qa( ) if is_pretrain: - tokenized_ds = tokenized_ds.remove_columns(["context", "qas", "text"]) + tokenized_ds = tokenized_ds.remove_columns(["context", "text"]) + if not is_paraphrased: + tokenized_ds = tokenized_ds.remove_columns(["qas"]) else: tokenized_ds = tokenized_ds.remove_columns( ["messages", "context", "prompt", "response"] @@ -777,13 +775,16 @@ def construct_and_tokenize_ctx_qa( # # the columns are unknown when using streaming dataset # tokenized_ds = tokenized_ds._resolve_features() # validate_columns(tokenized_ds) - if split == "train": - tokenized_ds.save_to_disk(ds_path, num_proc=num_proc) + return tokenized_ds -# def build_intx_pretrain(samples: dict[str, any]): -# return {"text": [ctx + qa for ctx, qa in zip(samples["context"], samples["qas"])]} +def build_paraphrase_pretrain(samples: dict[str, list[str]]): + # use "context" to predict "variation" as vice versa + return { + "text": samples["context"] + samples["variation"], + "context": samples["variation"] + samples["context"], + } def build_intx_pretrain(sample: dict[str, str]): @@ -1034,6 +1035,7 @@ def tokenize_pretrain( samples: dict[str, Any], tokenizer: PreTrainedTokenizerBase, tokenizer_kwargs: dict[str, Any] | None = None, + add_eos_token: bool = False, ): """ Tokenize text for pre-training tasks. @@ -1049,7 +1051,7 @@ def tokenize_pretrain( add_bos_val = tokenizer.add_bos_token add_eos_val = tokenizer.add_eos_token tokenizer.add_bos_token = True - tokenizer.add_eos_token = False # no eos for pre-training samples + tokenizer.add_eos_token = add_eos_token tokens = tokenizer( samples["text"], add_special_tokens=True,