diff --git a/data/generate_pretrain_from_fw_qa.py b/data/generate_pretrain_from_fw_qa.py new file mode 100644 index 0000000..007a3b2 --- /dev/null +++ b/data/generate_pretrain_from_fw_qa.py @@ -0,0 +1,38 @@ +import gc +from collections import defaultdict +from glob import glob + +from datasets import Dataset, load_dataset +from tqdm import tqdm + +QA_TEMPLATE = "\n\nQuestion: {question}\nAnswer: {answer}" + +if __name__ == "__main__": + root_data_dir = "./data/raw_datasets/fw_qa_3" + files = glob(f"{root_data_dir}/*.parquet") + for file in files: + ctx_qa_dict = defaultdict(str) + ds = load_dataset("parquet", data_files=file, split="train") + print(f"Loading dataset from {file}") + print(f"Original size: {len(ds)}") + for i, sample in tqdm(enumerate(ds)): + ctx = sample["context"] + question = sample["prompt"] + answer = sample["response"] + ctx_qa_dict[ctx] += QA_TEMPLATE.format(question=question, answer=answer) + print(f"Unique contexts: {len(ctx_qa_dict)}") + sampled_data = ctx_qa_dict[ctx] + print(f"Sampled context-qa pairs: {ctx}{''.join(sampled_data)}") + # convert ctx_qa_dict to a list of dictionaries + samples = [ + {"context": ctx, "qas": qa_pairs} for ctx, qa_pairs in ctx_qa_dict.items() + ] + # save to a new dataset + ds = Dataset.from_list(samples) + save_path = f"./data/raw_datasets/fw_qa_intx_pretrain/{file.split('/')[-1]}" + print(f"Saving dataset to {save_path}") + ds.to_parquet(save_path) + print("=" * 80) + del ds + del samples + gc.collect() diff --git a/src/ctx_to_lora/data/processing.py b/src/ctx_to_lora/data/processing.py index bd2c046..fb0f368 100644 --- a/src/ctx_to_lora/data/processing.py +++ b/src/ctx_to_lora/data/processing.py @@ -543,9 +543,6 @@ def construct_and_tokenize_ctx_qa( set_format=None, num_proc=None, ): - is_eval = split != "train" - # TODO: update hash function to include max_len - # TODO: max_len could be included in the kwargs kwargs = dict( tokenizer=repr(tokenizer), tokenizer_kwargs=json.dumps(tokenizer_kwargs), @@ -555,7 +552,7 @@ def construct_and_tokenize_ctx_qa( use_kl_loss=use_kl_loss, need_ctx_ids=need_ctx_ids, ds=ds._fingerprint, - # TODO: add split (after moving to continued pre-training pipeline) + split=split, set_format=set_format, ) kwargs_str = json.dumps(kwargs)