From e1a3926e25aa10745f46e9e144ca50188633d36d Mon Sep 17 00:00:00 2001
From: ShinnosukeUesakaSakana <shinnosukeuesaka@sakana.ai>
Date: Thu, 28 Aug 2025 21:03:53 +0900
Subject: [PATCH] Improved training data generation (#14)

* add catridge prompts

* fix bugs

* add instruction

* add generic prompt + improved prompts and template

* default max len

---------

Co-authored-by: 51616 <rujikorn.ch@gmail.com>
---
 README.md                     |   2 +-
 data/generate_fw_edu_qa_v3.py | 287 ++++++++++++++++++++++++++++++++++
 data/q_generation_prompts.py  | 114 ++++++++++++++
 3 files changed, 402 insertions(+), 1 deletion(-)
 create mode 100644 data/generate_fw_edu_qa_v3.py
 create mode 100644 data/q_generation_prompts.py

diff --git a/README.md b/README.md
index 61ab0cc..452c7f9 100644
--- a/README.md
+++ b/README.md
@@ -175,6 +175,7 @@ run uv run data/generate_fw_edu_qa_v2.py --shard_pattern "000_00000" --n_qa_pair
 run uv run data/generate_fw_edu_qa_v2_repeat.py --shard_pattern "min_0_to_2000/000*level_0*" --n_qa_pairs=5 --vllm_model=google/gemma-3-12b-it;
 run uv run data/generate_fw_edu_qa_v2_repeat.py --shard_pattern "min_0_to_2000/000*level_1*" --n_qa_pairs=5 --vllm_model=google/gemma-3-12b-it;
 run uv run data/generate_fw_edu_qa_v2_repeat.py --shard_pattern "min_0_to_2000/000*level_2*" --n_qa_pairs=5 --vllm_model=google/gemma-3-12b-it
+run uv run data/generate_fw_edu_qa_v3.py --n_qa_pairs 3 --shard_pattern '*' --debug --question_weight 1  --use_case_weight 1 --creative_weight 1 --generic_weight 1
 ```
 
 2. Self-generated response QA data (depends on step 0 and 1)
@@ -190,7 +191,6 @@ uv run data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --glob_pattern
 uv run data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --config configs/qa_short_ctx_self_gen_no_fw_qa.yaml
 ```
 
-
 <!-- 
 ### GSM8k LoRA
 ```bash
diff --git a/data/generate_fw_edu_qa_v3.py b/data/generate_fw_edu_qa_v3.py
new file mode 100644
index 0000000..e9e7717
--- /dev/null
+++ b/data/generate_fw_edu_qa_v3.py
@@ -0,0 +1,287 @@
+import argparse
+import os
+import random
+import re
+from glob import glob
+
+import pandas as pd
+from datasets import Dataset, load_dataset
+from q_generation_prompts import (  # get_structuring_seed_prompt,; get_summarization_seed_prompt,
+    PROMPT_TEMPLATE,
+    SYSTEM_PROMPT,
+    get_creative_seed_prompt,
+    get_generic_seed_prompt,
+    get_question_seed_prompt,
+    get_use_case_seed_prompt,
+)
+from vllm import LLM, SamplingParams
+
+STOP_STRINGS = {
+    "google/gemma-3-12b-it": ["<eos>", "<end_of_turn>"],
+}
+
+
+def build_messages(
+    context: str,
+    question_weight: float,
+    use_case_weight: float,
+    creative_weight: float,
+    generic_weight: float,
+) -> list[dict]:
+    # Create list of instruction functions with their weights
+    instruction_options = [
+        ("question", get_question_seed_prompt, question_weight),
+        ("use_case", get_use_case_seed_prompt, use_case_weight),
+        ("creative", get_creative_seed_prompt, creative_weight),
+        ("generic", get_generic_seed_prompt, generic_weight),
+    ]
+
+    q_types, functions, weights = zip(*instruction_options)
+
+    chosen_idx = random.choices(range(len(instruction_options)), weights=weights)[0]
+
+    q_type = q_types[chosen_idx]
+    function = functions[chosen_idx]
+
+    custom_instruction = function()
+    messages = [
+        {"role": "system", "content": SYSTEM_PROMPT},
+        {
+            "role": "user",
+            "content": PROMPT_TEMPLATE.format(
+                context=context, instruction=custom_instruction
+            ),
+        },
+    ]
+    return messages, q_type
+
+
+def check_should_skip(txt: str, vllm_model: str) -> bool:
+    """Check if the response should be skipped based on stop strings."""
+    for stop in STOP_STRINGS[vllm_model]:
+        if stop in txt[-len(stop) :]:
+            return (txt.split(stop)[0], False)  # Found a valid stop string
+    return (txt, True)  # No valid stop string found, skip this response
+
+
+def postprocess_questions(res_txt: str):
+    """
+    Postprocesses the questions from the response text.
+
+    Args:
+        res_txt: The response text.
+
+    Returns:
+        A list containing the questions.
+    """
+    # capture everything after each "Question {number}:" until the next question or end
+    res_txt = res_txt.split("</think>")[-1]
+
+    q_pattern = r"Message:(.*?)(?=Message \d+|$)"  # thanks chatgpt
+    questions = re.findall(q_pattern, res_txt, flags=re.S)
+
+    out_q = []
+    n_skips = 0
+    if len(questions) > 0:
+        for i, question in enumerate(questions):
+            question = question.strip()
+            if not question:
+                print(f"Skipping empty question at index {i}")
+                continue
+            # Check if this is the last question and handle stop strings
+            if i == len(questions) - 1:
+                question, skip = check_should_skip(question, vllm_model)
+                if skip:
+                    print(f"Skipping due to missing stop string")
+                    n_skips += 1
+                    continue
+            out_q.append(question.strip())
+    print(f"Skipped {n_skips} responses due to missing stop strings")
+
+    return out_q
+
+
+def length_filter(sample, min_len, max_len):
+    return min_len <= len(sample["text"]) <= max_len
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Generate QA pairs from FineWeb Edu dataset"
+    )
+    parser.add_argument(
+        "--vllm_model",
+        type=str,
+        default=os.environ.get("vllm_model", "google/gemma-3-12b-it"),
+        help="VLLM model to use for generation",
+    )
+    parser.add_argument(
+        "--shard_pattern",
+        type=str,
+        required=True,
+        help="Pattern to match shard files (e.g., '000_0000*')",
+    )
+    # parser.add_argument(
+    #     "--n_qa_pairs",
+    #     type=int,
+    #     required=True,
+    #     help="Number of question-answer pairs to generate per context",
+    # )
+    parser.add_argument(
+        "--min_length",
+        type=int,
+        default=0,
+        help="Minimum length of the context to consider for generation",
+    )
+    parser.add_argument(
+        "--max_length",
+        type=int,
+        default=20_000,
+        help="Maximum length of the context to consider for generation",
+    )
+    parser.add_argument(
+        "--max_model_length",
+        type=int,
+        default=2**13,
+        help="Maximum length of the model input (context + prompt + response) in tokens",
+    )
+    parser.add_argument(
+        "--debug",
+        action="store_true",
+        help="Debug mode - process only first 100 samples",
+    )
+    parser.add_argument(
+        "--question_weight",
+        type=float,
+        default=0,
+        help="Weight for question seed prompts",
+    )
+    parser.add_argument(
+        "--use_case_weight",
+        type=float,
+        default=0,
+        help="Weight for use case seed prompts",
+    )
+    parser.add_argument(
+        "--creative_weight",
+        type=float,
+        default=0,
+        help="Weight for creative seed prompts",
+    )
+    parser.add_argument(
+        "--generic_weight",
+        type=float,
+        default=0,
+        help="Weight for generic seed prompts",
+    )
+
+    args = parser.parse_args()
+
+    weights = [
+        args.question_weight,
+        args.use_case_weight,
+        args.creative_weight,
+        args.generic_weight,
+    ]
+    total_weight = sum(weights)
+    assert all(w >= 0 for w in weights), "All weights must be non-negative"
+    assert total_weight > 0, "At least one weight must be greater than zero"
+
+    vllm_model = args.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=args.max_model_length,
+        limit_mm_per_prompt={"image": 0},
+    )
+
+    llm = LLM(**llm_kwargs)
+    tokenizer = llm.get_tokenizer()
+    shard_pattern = args.shard_pattern
+    # n_qa_pairs = args.n_qa_pairs
+
+    paths = glob(
+        f"./data/raw_datasets/fineweb_edu/sample/100BT/{shard_pattern}.parquet"
+    )
+
+    split = "train[:100]" if args.debug else "train"
+    for path in paths:
+        ds = load_dataset(
+            "parquet",
+            data_files=path,
+            split=split,
+        )
+        ds = ds.filter(
+            length_filter,
+            fn_kwargs={"min_len": args.min_length, "max_len": args.max_length},
+            num_proc=8,
+        )
+
+        messages = []
+        q_types = []
+        for ctx in ds["text"]:
+            message, q_type = build_messages(
+                ctx,
+                args.question_weight,
+                args.use_case_weight,
+                args.creative_weight,
+                args.generic_weight,
+            )
+            messages.append(message)
+            q_types.append(q_type)
+
+        print(f"Generating from {len(messages)} contexts")
+        completions = llm.chat(
+            messages,
+            sampling_params=SamplingParams(
+                max_tokens=256,
+                temperature=0.0,
+                # needed for checking if stop tokens are present
+                skip_special_tokens=False,
+                include_stop_str_in_output=True,
+            ),
+        )
+        samples = []
+        for ctx, completion, q_type in zip(ds["text"], completions, q_types):
+            questions = postprocess_questions(completion.outputs[0].text)
+            samples.append(
+                {
+                    "context": ctx,
+                    "prompts": questions,
+                    "type": q_type,
+                }
+            )
+            if args.debug:
+                print(f"{ctx=}")
+                print(f"{completion.outputs[0].text=}")
+                print(f"{q_type=}")
+                for q in questions:
+                    print(f"{q=}")
+                    print()
+                print("=" * 80)
+
+        print(f"Generated {len(samples)} 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]
+        shard_name += "_level_0"
+        if args.debug:
+            shard_name += "_debug"
+        ds.to_parquet(
+            f"data/raw_datasets/fw_qa_v3/min_{args.min_length}_to_{args.max_length}/{shard_name}.parquet"
+        )
+        val_ds.to_parquet(
+            f"data/raw_datasets/fw_qa_v3/min_{args.min_length}_to_{args.max_length}/{shard_name}_val.parquet"
+        )
+        print(
+            f"Saved to data/raw_datasets/fw_qa_v3/min_{args.min_length}_to_{args.max_length}/{shard_name}.parquet"
+        )
+        print(
+            f"Saved to data/raw_datasets/fw_qa_v3/min_{args.min_length}_to_{args.max_length}/{shard_name}_val.parquet"
+        )
diff --git a/data/q_generation_prompts.py b/data/q_generation_prompts.py
new file mode 100644
index 0000000..810ad6e
--- /dev/null
+++ b/data/q_generation_prompts.py
@@ -0,0 +1,114 @@
+import random
+
+SYSTEM_PROMPT = (
+    "You are a creative and helpful assistant. "
+    "You will be given a context and you need to generate questions from the given context. "
+    "**DO NOT** hallucinate or make up information."
+)
+
+PROMPT_TEMPLATE = (
+    "### Context ###\n{context}\n\n"
+    "### Instruction ###\n{instruction}\n\n"
+    "### Rules ###\n"
+    "Phrases like 'based on the provided context', 'according to the context', etc., must not to appear in your response.\n\n"
+    # "2. The questions should not overlap. They should be diverse, covering many aspects of the context.\n"
+    # "3. Do not give away too much information in the questions. For example, ask 'Who is X?' instead of 'Who is X that did Y?' when Y is clear from the context.\n"
+    # "4. Ignore the text formatting of the context, e.g., bold, italic, underline, etc.\n"
+    # "5. Ignore typos, spacing, and grammatical errors in the context.\n"
+    # "6. Always use proper grammar and punctuation.\n"
+    # "7. Try to use different question forms and styles.\n"
+    "### Output Format ###\n"
+    "The question/instruction/task/request should be in the following format:\n\n"
+    "Message: {{question}}\n\n"
+    "Use this output template regardless of the type of the output."
+)
+
+# taken from https://github.com/HazyResearch/cartridges/blob/2ac563d79c2f3367a9e780a7bb0b3cf3039a8d50/cartridges/data/resources.py#L195
+# def get_structuring_seed_prompt() -> str:
+#     DATA_FORMATS = [
+#         "JSON",
+#         "YAML",
+#         "TOML",
+#         "INI",
+#         "XML",
+#         "plain text",
+#     ]
+
+#     data_format = random.choice(DATA_FORMATS)
+
+#     EXAMPLES = [
+#         dedent(f"""
+#             Can you structure the information in the context in the following format: {data_format}? Be sure to include precise information like any dates, times, names, and numerical values.
+#         """).strip(),
+#     ]
+
+#     example = random.choice(EXAMPLES)
+
+#     return dedent(f"""
+#         Please generate a single chat message instructing an LLM to structure the information in {data_format}. The message can follow the following template, filling in details from the context:
+
+#         '{example}'
+#     """).strip()
+
+
+# def get_summarization_seed_prompt() -> str:
+#     prompts = [
+#         dedent("""
+#             Please generate a single chat message instructing an LLM to summarize part of the context.
+#             Make sure the instruction is very explicit about the section of the context that you want to summarize.
+#             Include details (ids, names, titles, dates, etc.) that make it clear what you are asking about.
+#         """).strip(),
+#         dedent("""
+#             Please generate a single chat message instructing an LLM to summarize a section.
+#             Make sure the instruction is explicit about the section that should be summarized and the document it is from.
+#         """).strip(),
+#     ]
+#     prompt = random.choice(prompts)
+#     return prompt
+
+
+def get_question_seed_prompt() -> str:
+    prompts = [
+        (
+            "Generate a question for an LLM that will test its knowledge of the information in the context above. "
+            "In your question be sure to include details (ids, names, titles, dates, etc.) that make it clear what you are asking about. "
+            "Output only a single question. Do NOT include any other text or explanation other than the question."
+        ),
+        (
+            "Generate a message for an LLM that will test its knowledge of the information in the context above. "
+            "Be sure to include details (ids, names, titles, dates, etc.) in the question so that it can be answered without access to the context (i.e. closed-book setting). "
+            "Output only a single question. Do NOT include any other text or explanation other than the question."
+        ),
+        (
+            "You are helping to quiz a user about the information in the context. "
+            "Please generate a question about the subsection of the context above. "
+            "Be sure to include details (ids, names, titles, dates, etc.) in the question to make it clear what you are asking about. "
+            "Answer only with the question, do not include any other text."
+        ),
+    ]
+    prompt = random.choice(prompts)
+    return prompt
+
+
+def get_use_case_seed_prompt() -> str:
+    prompt = (
+        "Your primary goal is to think about practical, real-world tasks or applications that someone could achieve using the knowledge contained within the provided context. "
+        "Consider how a user might want to apply this information in a real-world scenario, not just recall it. Put yourself into someone else's shoes and think how you might use the information. "
+        "After considering potential use cases, your task will be to generate an instruction or task that reflects one of these downstream applications. "
+        "This instruction or task should be something a user, who has access to this context, might ask when trying to accomplish their specific goal. "
+        "Output only a single instruction or task. Do NOT include any other text or explanation."
+    )
+    return prompt
+
+
+def get_creative_seed_prompt() -> str:
+    prompt = (
+        "You are having a creative open-ended conversation inspired by the information in the context. "
+        "Please generate an open question for your conversation partner to start off the discussion. "
+        "Answer only with the question, do not include any other text."
+    )
+    return prompt
+
+
+def get_generic_seed_prompt() -> str:
+    return "Please generate a single chat message to begin a conversation about the information in the context. Make an open-ended request or provide a task."