toy ctx nums and multi-lora training (#11)

* multi-lora trainable toy number repeat dataset

* per rank bias init

* remove head_bias +simplify merge + skip perplexities metric

* ctx_numbers train example

* self-gen ctx numbers example

data/generate_ctx_kv.py

@@ -0,0 +1,243 @@
+import argparse
+import json
+import os
+import random
+import uuid
+from collections.abc import Iterable
+
+import wonderwords
+
+# -----------------------------
+# Simple, editable config knobs
+# -----------------------------
+# You can change these defaults or pass CLI flags: --key-type / --value-type
+DEFAULT_KEY_TYPE = "uuid"  # choices: "uuid" | "digits" | "words"
+DEFAULT_VALUE_TYPE = "uuid"  # choices: "uuid" | "digits" | "words"
+PAIR_SEP = "\n"  # between pairs
+KV_SEP = " : "  # between key and value
+TOKENS_PER_PAIR = 10  # rough heuristic used for bucketing by context length
+BASE_SAMPLES_PER_BIN = 128_000
+RNG_SEED = 42
+
+nouns = wonderwords.random_word._get_words_from_text_file("nounlist.txt")
+adjs = wonderwords.random_word._get_words_from_text_file("adjectivelist.txt")
+# verbs = wonderwords.random_word._get_words_from_text_file("verblist.txt")
+words = [f"{adj}-{noun}" for adj in adjs for noun in nouns]
+words = sorted(list(set(words)))
+
+
+def save_jsonl(data: list[dict], filepath: str) -> None:
+    """Save data to a JSONL file."""
+    parent_dir = os.path.dirname(filepath)
+    if parent_dir:
+        os.makedirs(parent_dir, exist_ok=True)
+    with open(filepath, "w") as f:
+        for entry in data:
+            json.dump(entry, f)
+            f.write("\n")
+
+
+def _gen_uuid8() -> str:
+    return uuid.uuid4().hex[:8]
+
+
+def _gen_digits6() -> str:
+    return f"{random.randint(0, 999_999):06d}"
+
+
+def _make_generator(kind: str):
+    kind = kind.lower()
+    if kind == "uuid":
+        return _gen_uuid8
+    if kind == "digits":
+        return _gen_digits6
+    if kind == "words":
+        return _gen_word
+    raise ValueError(f"Unknown kind '{kind}', expected 'uuid', 'digits', or 'words'")
+
+
+def _gen_digits4() -> str:
+    return f"{random.randint(0, 9_999):04d}"
+
+
+def _gen_word() -> str:
+    return random.choice(words)
+
+
+def _make_value_generator(kind: str):
+    """Like _make_generator, but when kind is 'digits' use 4 digits for values."""
+    kind = kind.lower()
+    if kind == "uuid":
+        return _gen_uuid8
+    if kind == "digits":
+        return _gen_digits4
+    if kind == "words":
+        return _gen_word
+    raise ValueError(f"Unknown kind '{kind}', expected 'uuid', 'digits', or 'words'")
+
+
+def _unique(seq: Iterable[str]) -> list[str]:
+    s = set()
+    out = []
+    for x in seq:
+        if x not in s:
+            s.add(x)
+            out.append(x)
+    return out
+
+
+def render_context(pairs: list[tuple[str, str]]) -> str:
+    """Render key-value pairs as a single context string.
+
+    Example: "k1:v1 | k2:v2 | k3:v3"
+    """
+    return PAIR_SEP.join([f"{k}{KV_SEP}{v}" for k, v in pairs])
+
+
+def generate_kv_dataset(n: int, k: int, key_type: str, value_type: str):
+    """Generate n examples. Each example has k key-value pairs in the context.
+
+    - context: "k1:v1 | k2:v2 | ..."
+    - prompt:  "What is the value of `{k}`?"
+    - response: "v" (exactly the value corresponding to the chosen k)
+    """
+    key_gen = _make_generator(key_type)
+    val_gen = _make_value_generator(value_type)
+
+    dataset = []
+
+    for _ in range(n):
+        # Generate unique keys to avoid ambiguity
+        keys: list[str] = []
+        while len(keys) < k:
+            keys = _unique([key_gen() for _ in range(k * 2)])  # oversample, then dedupe
+        keys = keys[:k]
+        values = [val_gen() for _ in range(k)]
+
+        pairs = list(zip(keys, values))
+
+        # Choose a random key from the context to ask about
+        q_key = random.choice(keys)
+        q_val = values[keys.index(q_key)]
+
+        entry = {
+            "context": f"key{KV_SEP}value\n" + render_context(pairs),
+            "prompt": f'What is the value of key "{q_key}"? Reply with only the value.',
+            "response": q_val,
+        }
+        dataset.append(entry)
+
+    # Split into train/val/test matching the same structure as generate_ctx_numbers.py
+    total_size = len(dataset)
+    train_size = int(0.98 * total_size)
+    val_size = int(0.01 * total_size)
+
+    train_data = dataset[:train_size]
+    val_data = dataset[train_size : train_size + val_size]
+    test_data = dataset[train_size + val_size :]
+    return train_data, val_data, test_data
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Generate key-value context QA data with the same structure as generate_ctx_numbers.py",
+    )
+    parser.add_argument(
+        "--key-type",
+        choices=["uuid", "digits", "words"],
+        default=DEFAULT_KEY_TYPE,
+        help="Type of keys to generate: 8-char uuid, 6-digit numbers, or words",
+    )
+    parser.add_argument(
+        "--value-type",
+        choices=["uuid", "digits", "words"],
+        default=DEFAULT_VALUE_TYPE,
+        help="Type of values to generate: 8-char uuid, 4-digit numbers, or words",
+    )
+    parser.add_argument("--seed", type=int, default=RNG_SEED, help="Random seed")
+    parser.add_argument(
+        "--base-samples-per-bin",
+        type=int,
+        default=BASE_SAMPLES_PER_BIN,
+        help="Baseline number of samples per token bin (scaled by bin width)",
+    )
+    parser.add_argument(
+        "--out-prefix",
+        type=str,
+        default="data/raw_datasets/ctx_kv",
+        help="Output directory prefix (bin range will be appended)",
+    )
+    parser.add_argument(
+        "--tokens-per-pair",
+        type=int,
+        default=TOKENS_PER_PAIR,
+        help="Heuristic tokens per key:value pair for bucketing",
+    )
+    parser.add_argument(
+        "--only-first-n-bins",
+        type=int,
+        default=None,
+        help="For quick tests: only generate the first N token bins",
+    )
+    parser.add_argument(
+        "--dry-run",
+        action="store_true",
+        help="Print a small sample and exit without writing files",
+    )
+
+    args = parser.parse_args()
+
+    random.seed(args.seed)
+
+    # Same token bins as generate_ctx_numbers.py
+    tok_bins = [(64, 128), (128, 256), (256, 512)] + [
+        (512 + 256 * i, 512 + 256 * (i + 1)) for i in range(14)
+    ]
+
+    if args.only_first_n_bins is not None:
+        tok_bins = tok_bins[: args.only_first_n_bins]
+
+    # Map token bins to pair-length bins using a simple heuristic
+    len_bins = [
+        (lo // args.tokens_per_pair, hi // args.tokens_per_pair)
+        for (lo, hi) in tok_bins
+    ]
+
+    # Optional dry-run: build a tiny sample to preview format and exit
+    if args.dry_run:
+        train, val, test = generate_kv_dataset(
+            n=3,
+            k=max(1, len_bins[0][0] or 1),
+            key_type=args.key_type,
+            value_type=args.value_type,
+        )
+        print("Sample entry:")
+        print(json.dumps(train[0], indent=2))
+        return
+
+    for len_bin, tok_bin in zip(len_bins, tok_bins):
+        bin_size = max(1, len_bin[1] - len_bin[0])  # avoid div-by-zero for tiny bins
+        save_dir = f"{args.out_prefix}_{tok_bin[0]}_{tok_bin[1]}"
+        train_data, val_data, test_data = [], [], []
+
+        # Iterate over different context lengths (number of pairs)
+        for k in range(max(1, len_bin[0]), max(1, len_bin[1])):
+            # scale like the numbers script
+            per_k = max(1, args.base_samples_per_bin // bin_size)
+            train, val, test = generate_kv_dataset(
+                n=per_k, k=k, key_type=args.key_type, value_type=args.value_type
+            )
+            train_data += train
+            val_data += val
+            test_data += test
+
+        os.makedirs(save_dir, exist_ok=True)
+        save_jsonl(train_data, f"{save_dir}/train.jsonl")
+        save_jsonl(val_data, f"{save_dir}/val.jsonl")
+        save_jsonl(test_data, f"{save_dir}/test.jsonl")
+
+        print(f"Dataset generated and saved at {save_dir}")
+
+
+if __name__ == "__main__":
+    main()

data/generate_ctx_magic_number.py

@@ -0,0 +1,195 @@
+import argparse
+import json
+import math
+import os
+import random
+
+# -----------------------------
+# Config knobs (edit or use CLI)
+# -----------------------------
+TOKENS_PER_BLOCK = 40  # rough heuristic tokens per noise block
+BASE_SAMPLES_PER_BIN = 12_800
+RNG_SEED = 42
+NOISE_BLOCK = "The grass is green. The sky is blue. The sun is yellow. Here we go. There and back again."
+SPECIAL_TPL = "The special magic number is {magic_number}."
+SEP = "\n"  # between blocks
+
+
+def save_jsonl(data: list[dict], filepath: str) -> None:
+    parent_dir = os.path.dirname(filepath)
+    if parent_dir:
+        os.makedirs(parent_dir, exist_ok=True)
+    with open(filepath, "w") as f:
+        for entry in data:
+            json.dump(entry, f)
+            f.write("\n")
+
+
+essential_digits4 = lambda: f"{random.randint(0, 9_999):04d}"
+
+
+def _choose_position(total_blocks: int, depth_bin: int) -> int:
+    """Choose an insertion index for the special sentence within [0, total_blocks-1]
+    such that its relative depth falls within the depth bin [i/10, (i+1)/10).
+    """
+    if total_blocks <= 0:
+        return 0
+    # Use floor for start and ceil for end to cover boundaries evenly
+    start = math.floor(total_blocks * (depth_bin / 10))
+    end = math.ceil(total_blocks * ((depth_bin + 1) / 10)) - 1
+    # clamp
+    start = max(0, min(start, total_blocks - 1))
+    end = max(start, min(end, total_blocks - 1))
+    return random.randint(start, end)
+
+
+def _build_example(total_blocks: int, depth_bin: int) -> dict:
+    """Build one example with a special line inserted among noise blocks.
+
+    total_blocks: total number of blocks in the final context (including the special one)
+    depth_bin: integer in [0, 9]
+    """
+    total_blocks = max(1, total_blocks)
+
+    # Prepare blocks
+    magic = essential_digits4()
+    special_line = SPECIAL_TPL.format(magic_number=magic)
+
+    # We'll have (total_blocks - 1) noise blocks and 1 special line
+    noise_count = max(0, total_blocks - 1)
+    blocks: list[str] = [NOISE_BLOCK for _ in range(noise_count)]
+
+    insert_at = _choose_position(total_blocks, depth_bin)
+    # Insert special line at the desired position within the final sequence
+    # If noise_count == 0, we just return special
+    if noise_count == 0:
+        final_blocks = [special_line]
+    else:
+        # Compose by interleaving noise and inserting special at index
+        # Build a list of length `total_blocks` and fill
+        final_blocks = []
+        noise_idx = 0
+        for idx in range(total_blocks):
+            if idx == insert_at:
+                final_blocks.append(special_line)
+            else:
+                final_blocks.append(blocks[noise_idx])
+                noise_idx += 1
+
+    context = SEP.join(final_blocks)
+    prompt = "What is the special magic number? Reply with only the number."
+    response = magic
+    return {"context": context, "prompt": prompt, "response": response}
+
+
+def generate_magic_dataset(n: int, k: int) -> tuple[list[dict], list[dict], list[dict]]:
+    """Generate n samples for a given block length k, evenly distributed across 10 depth bins."""
+    # Evenly divide n across 10 depth bins
+    base = n // 10
+    rem = n % 10
+    counts = [base + (1 if i < rem else 0) for i in range(10)]
+
+    dataset: list[dict] = []
+    for depth_bin, c in enumerate(counts):
+        for _ in range(c):
+            dataset.append(_build_example(total_blocks=k, depth_bin=depth_bin))
+
+    # Randomize before splitting to ensure val/test are random samples
+    random.shuffle(dataset)
+
+    # 98/1/1 split
+    total = len(dataset)
+    train_sz = int(0.98 * total)
+    val_sz = int(0.01 * total)
+    train = dataset[:train_sz]
+    val = dataset[train_sz : train_sz + val_sz]
+    test = dataset[train_sz + val_sz :]
+    return train, val, test
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Generate noise-wrapped special magic number dataset (similar structure to generate_ctx_kv.py)",
+    )
+    parser.add_argument("--seed", type=int, default=RNG_SEED, help="Random seed")
+    parser.add_argument(
+        "--base-samples-per-bin",
+        type=int,
+        default=BASE_SAMPLES_PER_BIN,
+        help="Baseline number of samples per token bin (scaled by bin width)",
+    )
+    parser.add_argument(
+        "--out-prefix",
+        type=str,
+        default="data/raw_datasets/ctx_magic_number",
+        help="Output directory prefix (bin range will be appended)",
+    )
+    parser.add_argument(
+        "--tokens-per-block",
+        "--tokens-per-pair",
+        dest="tokens_per_block",
+        type=int,
+        default=TOKENS_PER_BLOCK,
+        help="Heuristic tokens per noise block for bucketing",
+    )
+    parser.add_argument(
+        "--only-first-n-bins",
+        type=int,
+        default=None,
+        help="For quick tests: only generate the first N token bins",
+    )
+    parser.add_argument(
+        "--dry-run",
+        action="store_true",
+        help="Print a small sample and exit without writing files",
+    )
+
+    args = parser.parse_args()
+
+    random.seed(args.seed)
+
+    # Same token bins pattern as generate_ctx_kv.py (current version)
+    tok_bins = [(64, 128), (128, 256), (256, 512)] + [
+        (512 + 256 * i, 512 + 256 * (i + 1)) for i in range(14)
+    ]
+
+    if args.only_first_n_bins is not None:
+        tok_bins = tok_bins[: args.only_first_n_bins]
+
+    # Map token bins to block-length bins using heuristic
+    len_bins = [
+        (lo // args.tokens_per_block, hi // args.tokens_per_block)
+        for (lo, hi) in tok_bins
+    ]
+
+    if args.dry_run:
+        k = max(1, len_bins[0][0] or 1)
+        train, _, _ = generate_magic_dataset(n=10, k=k)
+        print("Sample entry:")
+        print(json.dumps(train[0], indent=2))
+        return
+
+    for len_bin, tok_bin in zip(len_bins, tok_bins):
+        bin_size = max(1, len_bin[1] - len_bin[0])
+        save_dir = f"{args.out_prefix}_{tok_bin[0]}_{tok_bin[1]}"
+        train_data: list[dict] = []
+        val_data: list[dict] = []
+        test_data: list[dict] = []
+
+        for k in range(max(1, len_bin[0]), max(1, len_bin[1])):
+            per_k = max(1, args.base_samples_per_bin // bin_size)
+            train, val, test = generate_magic_dataset(n=per_k, k=k)
+            train_data += train
+            val_data += val
+            test_data += test
+
+        os.makedirs(save_dir, exist_ok=True)
+        save_jsonl(train_data, f"{save_dir}/train.jsonl")
+        save_jsonl(val_data, f"{save_dir}/val.jsonl")
+        save_jsonl(test_data, f"{save_dir}/test.jsonl")
+
+        print(f"Dataset generated and saved at {save_dir}")
+
+
+if __name__ == "__main__":
+    main()

data/generate_ctx_numbers.py

@@ -0,0 +1,95 @@
+import json
+import os
+import random
+
+
+def save_jsonl(data: list[dict], filepath: str) -> None:
+    """Save data to a JSONL file."""
+    parent_dir = os.path.dirname(filepath)
+    if parent_dir:  # Only create directories if there's a parent path
+        os.makedirs(parent_dir, exist_ok=True)
+    with open(filepath, "w") as f:
+        for entry in data:
+            json.dump(entry, f)
+            f.write("\n")
+
+
+def get_random_combinations(numbers: list[int], n: int, k: int) -> list[list[int]]:
+    """Get n random combinations of k numbers from the list."""
+    # using itertools.combinations hangs with large numbers
+    # so we explicitly generate the n indices
+    indices = [random.sample(range(len(numbers)), k=k) for _ in range(n)]
+    return [[numbers[i] for i in ind] for ind in indices]
+
+
+def generate_number_dataset(max_num: int = 1000, n: int = 12000, k: int = 3):
+    """
+    Generate a dataset of numbers with corresponding query and answer,
+    split into train/val/test sets.
+
+    Args:
+        max_num: Maximum number in the range (exclusive)
+        k: Number of elements in each combination
+    """
+    # Generate list of all numbers and shuffle them
+    numbers = list(range(max_num))
+    random.shuffle(numbers)
+
+    # Create dataset entries
+    dataset = []
+    query = f"Repeat the information above exactly. Do not output anything else."
+
+    # Generate all unique combinations of k numbers
+    combinations = get_random_combinations(numbers, n, k)
+    # random.shuffle(combinations)
+
+    for combination in combinations:
+        num_list = [f"{i + 1}. {num}" for i, num in enumerate(combination)]
+        entry = {
+            "context": "\n".join(num_list),
+            "prompt": query,
+            "response": "\n".join(num_list),
+        }
+        dataset.append(entry)
+
+    # Calculate split sizes
+    total_size = len(dataset)
+    train_size = int(0.98 * total_size)
+    val_size = int(0.01 * total_size)
+
+    # Split dataset
+    train_data = dataset[:train_size]
+    val_data = dataset[train_size : train_size + val_size]
+    test_data = dataset[train_size + val_size :]
+    return train_data, val_data, test_data
+
+
+if __name__ == "__main__":
+    # Set random seed for reproducibility
+    random.seed(42)
+    tok_bins = [(64, 128), (128, 256), (256, 512)] + [
+        (512 + 256 * i, 512 + 256 * (i + 1)) for i in range(14)
+    ]
+    # roughly 9 tokens per number
+    tok_per_num = [9 if bin[0] >= 1024 else 8 for bin in tok_bins]
+    len_bins = [
+        (bin[0] // tok, bin[1] // tok) for bin, tok in zip(tok_bins, tok_per_num)
+    ]
+
+    for len_bin, tok_bin in zip(len_bins, tok_bins):
+        bin_size = len_bin[1] - len_bin[0]
+        save_dir = f"data/raw_datasets/ctx_numbers_{tok_bin[0]}_{tok_bin[1]}"
+        train_data, val_data, test_data = [], [], []
+        for k in range(*len_bin):
+            train, val, test = generate_number_dataset(n=128_000 // bin_size, k=k)
+            train_data += train
+            val_data += val
+            test_data += test
+
+        # Save splits to separate files
+        os.makedirs(save_dir, exist_ok=True)
+        save_jsonl(train_data, f"{save_dir}/train.jsonl")
+        save_jsonl(val_data, f"{save_dir}/val.jsonl")
+        save_jsonl(test_data, f"{save_dir}/test.jsonl")
+
+        print(f"Dataset generated and saved at {save_dir}")

data/self_generate_qa.py

@@ -41,31 +41,21 @@ SYSTEM_TEMPLATE = (
     "### END OF SYSTEM INSTRUCTION ###"
 )
 
+CTX_Q_SEP = "\n\n---\n\n# User\n"
+
 SELF_QA_INTX = (
-    "---\n\n"
+    # "---\n\n"
     "# System Instruction\n"
     "- The information provided is up-to-date information.\n"
     "- When the provided information is not relevant to the question, ***ignore*** it and answer the question based on your knowledge.\n"
-    "- If the provided information is related to the question, incorporate it in your response.\n"
-    "\n---\n\n"
-    "# User\n"
+    "- If the provided information is related to the question, incorporate it in your response."
+    + CTX_Q_SEP
 )
 
-PRE_CTX = "# Information\n"
+PROMPT_TEMPLATE = "{context}" + CTX_Q_SEP + "{question}"
 
-PROMPT_TEMPLATE = (
-    PRE_CTX
-    + "{context}\n\n---\n\n"
-    +
-    # "# System Instruction\n"
-    # "- The information provided is up-to-date information.\n"
-    # "- When the provided information is not relevant to the question, ***ignore*** it and answer the question based on your knowledge.\n"
-    # "- If the provided information is related to the question, incorporate it in your response.\n"
-    # "\n---\n\n"
-    # "# User\n{question}"
-    SELF_QA_INTX
-    + "{question}"
-)
+PRE_CTX = ""
+PROMPT_TEMPLATE_QA = "# Information\n{context}\n\n---\n\n" + SELF_QA_INTX + "{question}"
 
 
 """
@@ -114,8 +104,9 @@ def truncate_middle_if_too_long(
     return input_ids
 
 
-def get_prompt(context: str, q: str) -> str:
-    return PROMPT_TEMPLATE.format(context=context, question=q)
+def get_prompt(context: str, q: str, use_qa_template: bool) -> str:
+    template = PROMPT_TEMPLATE_QA if use_qa_template else PROMPT_TEMPLATE
+    return template.format(context=context, question=q)
 
 
 def add_closed_qa_prompt(q: str, closed_qa_prob: float = 0.1) -> str:
@@ -187,7 +178,11 @@ def get_dataset_configs(
 
 
 def create_messages(
-    ctxs: list[str], questions: list[list[str]], vllm_model: str, system_template: str
+    ctxs: list[str],
+    questions: list[list[str]],
+    vllm_model: str,
+    system_template: str,
+    use_qa_template: bool,
 ) -> list[list[dict]]:
     """Create chat messages for the model."""
     # if "gemma" in vllm_model:
@@ -196,7 +191,9 @@ def create_messages(
         [
             {
                 "role": "user",
-                "content": system_template + "\n\n\n" + get_prompt(ctx, q),
+                "content": system_template
+                + "\n\n\n"
+                + get_prompt(ctx, q, use_qa_template),
             }
         ]
         for ctx, q_list in zip(ctxs, questions)
@@ -221,6 +218,8 @@ def self_generate(
     system_template: str,
     parquet_file: str | None = None,
     do_truncate: bool = False,
+    use_qa_template: bool = True,
+    max_new_tokens: int = 1024,
 ) -> None:
     """Process a single dataset and generate QA pairs."""
 
@@ -254,6 +253,7 @@ def self_generate(
     print(f"Processing dataset: {ds_name}, split: {split}")
     print(f"Using temperature: {temp}")
     print(f"Using closed QA prompt probability: {closed_qa_prob}")
+    print(f"Use QA template: {use_qa_template}")
 
     if parquet_file:
         print(f"Loading dataset from parquet file: {parquet_file}")
@@ -282,8 +282,8 @@ def self_generate(
 
     tk = get_tokenizer(args.vllm_model, train=True)
 
-    self_qa_intx_tokens = tk(SELF_QA_INTX, add_special_tokens=False)["input_ids"]
-    n_self_qa_intx_tokens = len(self_qa_intx_tokens)
+    ctx_q_sep_tokens = tk(CTX_Q_SEP, add_special_tokens=False)["input_ids"]
+    n_qa_sep_tokens = len(ctx_q_sep_tokens)
     pre_ctx_tokens = tk(PRE_CTX, add_special_tokens=False)["input_ids"]
     n_pre_ctx_tokens = len(pre_ctx_tokens)
     sys_tokens = tk(system_template.split("\n")[0], add_special_tokens=False)[
@@ -316,7 +316,11 @@ def self_generate(
         chunk_ctxs = ctxs[start : start + chunk_size]
         chunk_questions = questions[start : start + chunk_size]
         chunk_messages = create_messages(
-            chunk_ctxs, chunk_questions, args.vllm_model, SELF_GEN_SYSTEM_MSG
+            chunk_ctxs,
+            chunk_questions,
+            args.vllm_model,
+            SELF_GEN_SYSTEM_MSG,
+            use_qa_template,
         )
 
         if do_truncate:
@@ -330,7 +334,7 @@ def self_generate(
                 truncate_middle_if_too_long(
                     ids,
                     max_length=MODEL_CTX_LEN[args.vllm_model],
-                    max_new_tokens=1024,
+                    max_new_tokens=max_new_tokens,
                 )
                 for ids in tokenized_contents["input_ids"]
             ]
@@ -350,8 +354,8 @@ def self_generate(
             llm,
             temp,
             tk,
-            self_qa_intx_tokens,
-            n_self_qa_intx_tokens,
+            ctx_q_sep_tokens,
+            n_qa_sep_tokens,
             sys_tokens,
             n_sys_tokens,
             chunk_ctxs,
@@ -359,6 +363,7 @@ def self_generate(
             chunk_questions,
             chunk_messages,
             k,
+            max_new_tokens,
         )
 
 
@@ -368,8 +373,8 @@ def execute_qa_generation(
     llm,
     temp,
     tk,
-    self_qa_intx_tokens,
-    n_self_qa_intx_tokens,
+    ctx_q_sep_tokens,
+    n_qa_sep_tokens,
     sys_tokens,
     n_sys_tokens,
     ctxs,
@@ -377,11 +382,12 @@ def execute_qa_generation(
     questions,
     messages,
     k,
+    max_tokens,
 ):
     completions = llm.chat(
         messages,
         sampling_params=SamplingParams(
-            max_tokens=1024,
+            max_tokens=max_tokens,
             logprobs=k,
             temperature=temp,
             seed=42,
@@ -389,6 +395,8 @@ def execute_qa_generation(
             skip_special_tokens=False,
             include_stop_str_in_output=True,
         ),
+        chat_template=tk.chat_template,
+        chat_template_content_format="string",
     )
 
     self_gen_data = {
@@ -451,13 +459,13 @@ def execute_qa_generation(
 
             q_start = None
             for ii in range(
-                len(prompt_ids) - n_self_qa_intx_tokens,
+                len(prompt_ids) - n_qa_sep_tokens,
                 -1,
                 -1,
             ):
-                if prompt_ids[ii : ii + n_self_qa_intx_tokens] == self_qa_intx_tokens:
+                if prompt_ids[ii : ii + n_qa_sep_tokens] == ctx_q_sep_tokens:
                     # found the start of the user input
-                    q_start = ii + n_self_qa_intx_tokens
+                    q_start = ii + n_qa_sep_tokens
                     break
 
             # bos + question + eos + start model turn + response + eos
@@ -611,6 +619,17 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Truncate contexts to fit model context length",
     )
+    parser.add_argument(
+        "--max_new_tokens",
+        type=int,
+        default=1024,
+        help="Maximum number of new tokens to generate (default: 1024)",
+    )
+    parser.add_argument(
+        "--remove_qa_template",
+        action="store_true",
+        help="Remove the QA template from the prompt (default: False)",
+    )
     return parser.parse_args()
 
 
@@ -650,7 +669,17 @@ if __name__ == "__main__":
         # Process each dataset
         for ds_name, split in dataset_configs:
             print(f"Processing dataset: {ds_name}, split: {split}")
-            self_generate(ds_name, split, args, llm, SYSTEM_TEMPLATE, args.do_truncate)
+            self_generate(
+                ds_name,
+                split,
+                args,
+                llm,
+                system_template=SYSTEM_TEMPLATE,
+                parquet_file=None,
+                do_truncate=args.do_truncate,
+                use_qa_template=not args.remove_qa_template,
+                max_new_tokens=args.max_new_tokens,
+            )
     else:
         assert args.glob_pattern, (
             "glob_pattern must be provided if no ds_names or config"
@@ -666,4 +695,6 @@ if __name__ == "__main__":
                 llm=llm,
                 system_template=SYSTEM_TEMPLATE,
                 do_truncate=args.do_truncate,
+                use_qa_template=not args.remove_qa_template,
+                max_new_tokens=args.max_new_tokens,
             )