From 624234109943d026b32e9031cfc5f23e0cf55900 Mon Sep 17 00:00:00 2001
From: 51616 <rujikorn.ch@gmail.com>
Date: Fri, 6 Jun 2025 09:48:42 +0000
Subject: [PATCH] data scripts + readme

---
 README.md                            |  41 ++++---
 configs/qa_no_fw.yaml                |  59 ++++++++++
 data/download_fineweb_edu.py         |  10 ++
 data/generate_fw_edu_augment_vllm.py | 157 +++++++++++++++++++++++++++
 data/generate_pretrain_from_fw_qa.py |   5 +-
 5 files changed, 253 insertions(+), 19 deletions(-)
 create mode 100644 configs/qa_no_fw.yaml
 create mode 100644 data/download_fineweb_edu.py
 create mode 100644 data/generate_fw_edu_augment_vllm.py

diff --git a/README.md b/README.md
index bcaeeac..9b2cb39 100644
--- a/README.md
+++ b/README.md
@@ -91,26 +91,35 @@ conda activate vllm;
 vllm_model=mistralai/Mistral-Small-3.1-24B-Instruct-2503 uv run python generate_fw_edu_qa_vllm.py "0*_*" 3
 # [WIP] there are other datasets where we used openAI to gen QA pairs
 
-# self-generate SFT data
-# gemma-2-2b-it
-vllm_model=google/gemma-2-2b-it uv run python data/self_generate_qa.py "0*_*"
+# 0. download fineweb_edu to `data/raw_datasets/fineweb_edu
+uv run data/download_fineweb_edu.py
+
+# ===
+# 1 and 2 can be run in parallel (depends on step 0)
+
+# 1. generate QA data (3 QAs for each context)
+vllm_model=mistralai/Mistral-Small-3.1-24B-Instruct-2503 uv run python generate_fw_edu_qa_vllm.py "00*_*" 3
+
+# 2. Augmenting pre-training data w/ paraphrasing
+uv run python data/generate_fw_edu_augment_vllm.py
+
+# ===
+# 3 and 4 can be run in parallel (depends on step 0 and 1)
+
+# 3. Pre-training data (no GPU needed) 
+# augment each context by concat'ing QAs generated in step 1 at the end
+uv run python data/generate_pretrain_from_fw_qa.py
+
+# 4. self-generate SFT data
 # each context has 3 chat instances as training data
 # the responses are not the same as the ones in the pretraining though
+# Example commands using gemma-2-2b-it
+# self-gen data for fw_qa
+uv run python data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --glob_pattern 'data/raw_datasets/fw_qa_3/00*_*'
+# self-gen data for other ds listed in qa_no_fw.yaml
+uv run python data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --config configs/qa_no_fw.yaml
 
 
-# pre-training data
-# better to include QAs at the end of pretrain docs
-# # https://arxiv.org/pdf/2406.14491
-uv run python data/generate_pretrain_from_fw_qa.py
-# we could do this online but
-# will have duplicate contexts (each context is used to generate 3 QA pairs)
-# better if we do this offline bc dedup the contexts
-
-
-# self-gen sft
-# 3_mini config
-uv run python data/self_generate_qa.py \
---vllm_model=google/gemma-2-2b-it --config=configs/self_gen_3_mini.yaml
 ```
 
 <!-- 
diff --git a/configs/qa_no_fw.yaml b/configs/qa_no_fw.yaml
new file mode 100644
index 0000000..e59f536
--- /dev/null
+++ b/configs/qa_no_fw.yaml
@@ -0,0 +1,59 @@
+output_dir: "" # just a placeholder
+bf16: true
+model_name_or_path: google/gemma-2-2b-it
+label_names: ["labels"]
+# eval_on_start: True
+# eval_strategy: "steps"
+# eval_steps: 500
+# save_strategy: "no"
+# # save_steps: 500
+# logging_strategy: "steps"
+# logging_steps: 100
+# use_liger_kernel: true
+# remove_unused_columns: false
+
+# needed to avoid OOM by compute the metrics batch by batch
+# w/o this the trainer stores logits of all sample in memory...
+# batch_eval_metrics: true
+
+per_device_train_batch_size: 8
+per_device_eval_batch_size: 8
+max_val_samples_per_ds: 1000
+# optim: schedule_free_adamw
+
+learning_rate: 0.00004
+# lr_scheduler_type: "constant_with_warmup"
+neftune_noise_alpha: 5
+weight_decay: 0.01
+#
+warmup_steps: 100
+
+dataloader_prefetch_factor: 8
+dataloader_num_workers: 8
+# LoRA
+lora_r: 8
+lora_dropout: 0.0
+target_modules:
+  - down_proj
+# data
+train_ds_names:
+  - ctx_qa # 300k
+  - pwc # 240k
+  - hotpot_qa # 90k
+  - squad # 90k
+  - drop # 77k
+  - narrativeqa # 40k
+  - quoref # 11k
+  - ropes # 11k
+  - synthetic_convqa # 40k
+
+val_ds_names:
+  - fw_qa_3
+  - fw_qa_xl
+  - ctx_qa
+  - pwc
+  - hotpot_qa
+  - squad
+
+load_best_model_at_end: true
+metric_for_best_model: eval_pwc_loss
diff --git a/data/download_fineweb_edu.py b/data/download_fineweb_edu.py
new file mode 100644
index 0000000..e0d740d
--- /dev/null
+++ b/data/download_fineweb_edu.py
@@ -0,0 +1,10 @@
+from huggingface_hub import snapshot_download
+
+if __name__ == "__main__":
+    fw_dir = "./data/raw_datasets/fineweb_edu/"
+    snapshot_download(
+        "HuggingFaceFW/fineweb-edu",
+        repo_type="dataset",
+        local_dir=fw_dir,
+        allow_patterns="sample/100BT/*",
+    )
diff --git a/data/generate_fw_edu_augment_vllm.py b/data/generate_fw_edu_augment_vllm.py
new file mode 100644
index 0000000..5f05825
--- /dev/null
+++ b/data/generate_fw_edu_augment_vllm.py
@@ -0,0 +1,157 @@
+import os
+import random
+import sys
+from glob import glob
+
+import pandas as pd
+from datasets import Dataset, load_dataset
+from vllm import LLM, SamplingParams
+
+from ctx_to_lora.utils import clear_gpu
+
+MODEL_CTX_LEN = {
+    "google/gemma-2-27b-it": 8192,
+    "google/gemma-2-2b-it": 8192,
+    "google/gemma-2-9b-it": 8192,
+    # not actually 16k but it's faster
+    "mistralai/Mistral-Small-3.1-24B-Instruct-2503": 2**14,  # 11264,
+}
+
+# based on New News: system-2 finetuning (https://arxiv.org/pdf/2505.01812)
+SYSTEM_TEMPLATE = (
+    "You are an careful and creative paraphraser. Paraphrase the given content without leaving out any important information. "
+    "You will be given a context and you must generate a new variation of the provided context.\n"
+    "You only output the paraphrase itself.\n"
+    "**DO NOT** halucinate or make up information."
+)
+
+# based on On the generalization of language models from in-context learning and finetuning: a controlled study
+# https://arxiv.org/pdf/2505.00661
+PROMPT_TEMPLATE = (
+    "### Context ###\n"
+    "{context}"
+    "\n\n\n"
+    "### Instructions ###\n"
+    "- Please generate rephrasings that can be inferred from each sentence. Even a simple sentence has some alternative phrasings that are logically equivalent.\n"
+    "- Reversing a sentence is also valid but make sure that the reversed sentence is still logically equivalent.\n"
+    "- You can shuffle the order of sentences as long as the content is coherent and the meaning is preserved.\n"
+    "- You can combine multiple sentences into one, or split a long sentence into multiple shorter ones.\n"
+    "- You can combine reversal and paraphasing, e.g., reversing a sentence and then paraphrasing it.\n"
+    "- You can also add new sentences that can be logically inferred from the context. "
+    "Please only use logic and language to draw your conclusions, regardless of whether the entities in question actually exist.\n"
+    "- Try to keep the length to be similar to the original context\n"
+    "- Please accurately maintain facts, such as names and numbers, mentioned in the context."
+    "\n\nOnly output a new variation of the context. Don't include any other words."
+)
+
+
+def get_prompt(context):
+    prompt = PROMPT_TEMPLATE.format(context=context)
+    return prompt
+
+
+def length_filter(samples):
+    return [len(text) < 10_000 for text in samples["text"]]
+
+
+if __name__ == "__main__":
+    vllm_model = os.environ.get("vllm_model")
+    print(f"Using model: {vllm_model}")
+    llm_kwargs = dict(
+        model=vllm_model,
+        dtype="bfloat16",
+        enable_prefix_caching=True,
+        enable_chunked_prefill=True,
+        max_model_len=MODEL_CTX_LEN[vllm_model],
+    )
+
+    if vllm_model == "mistralai/Mistral-Small-3.1-24B-Instruct-2503":
+        mm_kwargs = {"image": 0}
+        llm_kwargs["tokenizer_mode"] = "mistral"
+        llm_kwargs["config_format"] = "mistral"
+        llm_kwargs["load_format"] = "mistral"
+        llm_kwargs["limit_mm_per_prompt"] = mm_kwargs
+        SYSTEM_TEMPLATE = (
+            "You are Mistral Small 3.1, a Large Language Model (LLM) created by Mistral AI, a French startup headquartered in Paris.\n"
+            "You power an AI assistant called Le Chat.\n"
+            "Your knowledge base was last updated on 2023-10-01.\n\n\n"
+        ) + SYSTEM_TEMPLATE
+    llm = LLM(**llm_kwargs)
+    tokenizer = llm.get_tokenizer()
+    shard_pattern = sys.argv[1]
+
+    paths = glob(
+        f"./data/raw_datasets/fineweb_edu/sample/100BT/{shard_pattern}.parquet"
+    )
+
+    for path in paths:
+        ds = load_dataset(
+            "parquet",
+            data_files=path,
+            split="train",
+            # streaming=True,
+        )
+        ds = ds.filter(length_filter, batched=True)
+
+        ctxs = [sample["text"] for sample in iter(ds)]
+        if "gemma" in vllm_model:
+            messages = [
+                [
+                    {
+                        "role": "user",
+                        "content": SYSTEM_TEMPLATE + "\n\n\n" + get_prompt(ctx),
+                    },
+                ]
+                for ctx in ctxs
+            ]
+        else:
+            messages = [
+                [
+                    {"role": "system", "content": SYSTEM_TEMPLATE},
+                    {"role": "user", "content": get_prompt(ctx)},
+                ]
+                for ctx in ctxs
+            ]
+
+        print(f"Generating from {len(messages)} contexts")
+        completions = llm.chat(
+            messages,
+            sampling_params=SamplingParams(
+                max_tokens=2048,
+                temperature=(
+                    0.2
+                    if vllm_model == "mistralai/Mistral-Small-3.1-24B-Instruct-2503"
+                    else 0.7
+                ),
+                # frequency_penalty=0.1,
+            ),
+        )
+        samples = []
+        for ctx, completion in zip(ctxs, completions):
+            samples.append(
+                {
+                    "context": ctx,
+                    "variation": completion.outputs[0].text.strip(),
+                }
+            )
+        for sample in samples:
+            print(f"{sample['context']=}")
+            print(f"{sample['variation']=}")
+            print("=" * 60)
+
+        print(f"Generated {len(samples)} samples")
+        random.shuffle(samples)
+        df = pd.DataFrame(samples)
+        ds = Dataset.from_pandas(df)
+        val_ds = ds.take(10)
+        ds = ds.skip(10)
+
+        shard_name = path.split("/")[-1].split(".")[0]
+        ds.to_parquet(f"data/raw_datasets/fw_qa_aug_pretrain/{shard_name}.parquet")
+        val_ds.to_parquet(
+            f"data/raw_datasets/fw_qa_aug_pretrain/{shard_name}_val.parquet"
+        )
+        print(f"Saved to data/raw_datasets/fw_qa_aug_pretrain/{shard_name}.parquet")
+        print(f"Saved to data/raw_datasets/fw_qa_aug_pretrain/{shard_name}_val.parquet")
+        del ds, val_ds, df, samples, completions, messages, ctxs
+        clear_gpu()
diff --git a/data/generate_pretrain_from_fw_qa.py b/data/generate_pretrain_from_fw_qa.py
index 007a3b2..9531535 100644
--- a/data/generate_pretrain_from_fw_qa.py
+++ b/data/generate_pretrain_from_fw_qa.py
@@ -9,7 +9,7 @@ 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")
+    files = sorted(glob(f"{root_data_dir}/*.parquet"))
     for file in files:
         ctx_qa_dict = defaultdict(str)
         ds = load_dataset("parquet", data_files=file, split="train")
@@ -33,6 +33,5 @@ if __name__ == "__main__":
         print(f"Saving dataset to {save_path}")
         ds.to_parquet(save_path)
         print("=" * 80)
-        del ds
-        del samples
+        del ds, samples, ctx_qa_dict
         gc.collect()