apunkt/doc-to-lora
1
0
Fork 0
mirror of https://github.com/SakanaAI/doc-to-lora.git synced 2026-07-23 17:01:04 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions

iclr cleanup

Browse source
This commit is contained in:
51616 2025-09-30 14:53:21 +00:00
parent 8745db6e11
commit b6679ba755
171 changed files with 219 additions and 248151 deletions
Download patch file Download diff file Expand all files Collapse all files

51
data/create_test_data.py
View file

@ -1,51 +0,0 @@
import argparse
import pandas as pd
def make_test_samples():
return [
{
"context": "Alice was beginning to get very tired of sitting by her sister on the bank.",
"prompts_level_0": [
"Who is sitting by Alice?",
"How is Alice feeling at the start?"
],
"responses_level_0": [
"Her sister is sitting by her.",
"She is very tired."
],
},
{
"context": "<think>debug info</think>Then she saw a White Rabbit with pink eyes run close by her.",
"prompts_level_0": [
"What did Alice see run by?",
"What color were its eyes?"
],
"responses_level_0": [
"She saw a White Rabbit.",
"Its eyes were pink."
],
},
{
"context": "x" * 200,
"prompts_level_0": ["How long is this context?"],
"responses_level_0": ["It is two hundred characters long."],
},
]
def main(output_path: str):
samples = make_test_samples()
df = pd.DataFrame(samples)
df.to_parquet(output_path, index=False)
print(f"Wrote {len(df)} test rows to {output_path}")
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument(
"--output",
"-o",
type=str,
default="test_shard.parquet",
help="Where to write the test Parquet file"
)
args = parser.parse_args()
main(args.output)

20
data/download_facts_grounding.py
View file

@ -1,20 +0,0 @@
# Install dependencies as needed:
# pip install kagglehub[hf-datasets]
import kagglehub
from kagglehub import KaggleDatasetAdapter
# Set the path to the file you'd like to load
file_path = "data/raw_datasets/FACTS/ds.parquet"
# Load the latest version
hf_dataset = kagglehub.load_dataset(
KaggleDatasetAdapter.HUGGING_FACE,
"deepmind/FACTS-grounding-examples",
file_path,
# Provide any additional arguments like
# sql_query, hf_kwargs, or pandas_kwargs. See
# the documenation for more information:
# https://github.com/Kaggle/kagglehub/blob/main/README.md#kaggledatasetadapterhugging_face
)
print("Hugging Face Dataset:", hf_dataset)

9
data/download_generated_fineweb_qa.py
View file

@ -1,9 +0,0 @@
from huggingface_hub import snapshot_download
if __name__ == "__main__":
fw_dir = "./data/raw_datasets/fw_qa_v2/"
snapshot_download(
"SakanaAI/fineweb_qa",
repo_type="dataset",
local_dir=fw_dir,
)

9
data/download_self_gen_qa.py
View file

@ -1,9 +0,0 @@
from huggingface_hub import snapshot_download
if __name__ == "__main__":
fw_dir = "./data/raw_datasets/self_gen/"
snapshot_download(
"SakanaAI/self_gen_qa",
repo_type="dataset",
local_dir=fw_dir,
)

9
data/download_self_gen_qa_logprobs.py
View file

@ -1,9 +0,0 @@
from huggingface_hub import snapshot_download
if __name__ == "__main__":
fw_dir = "./data/raw_datasets/self_gen/"
snapshot_download(
"Rujikorn/self_gen_qa_logprobs",
repo_type="dataset",
local_dir=fw_dir,
)

220535
data/eval/facts/examples.csv
View file

File diff suppressed because one or more lines are too long

11313
data/eval/triviaqa/triviaqa_retrieved.jsonl
View file

File diff suppressed because one or more lines are too long

243
data/generate_ctx_kv.py
View file

@ -1,243 +0,0 @@
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()

95
data/generate_ctx_numbers.py
View file

@ -1,95 +0,0 @@
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}")

104
data/generate_fav.py
View file

@ -1,104 +0,0 @@
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, save_dir: str = None
):
"""
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 text above."
# Generate all unique combinations of k numbers
combinations = get_random_combinations(numbers, n, k)
# random.shuffle(combinations)
for combination in combinations:
entry = {
"context": ", ".join(map(str, combination)),
"prompt": query,
"response": ", ".join(map(str, combination)),
}
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 :]
save_dir = "" if save_dir is None else save_dir
# Save splits to separate files
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")
if __name__ == "__main__":
# Set random seed for reproducibility
random.seed(42)
for k in range(2, 11):
save_dir = f"data/raw_datasets/context_numbers_{k}"
os.makedirs(save_dir, exist_ok=True)
generate_number_dataset(n=12_000, k=k, save_dir=save_dir)
print(f"Dataset generated and saved at {save_dir}")
# # Generate dataset
# for k in range(2, 257):
# save_dir = f"data/raw_datasets/context_numbers_{k}"
# os.makedirs(save_dir, exist_ok=True)
# generate_number_dataset(n=12_000, k=k, save_dir=save_dir)
# print(f"Dataset generated and saved at {save_dir}")
# # for k in range(144, 257, 16):
# # save_dir = f"context_numbers_{k}"
# # os.makedirs(save_dir, exist_ok=True)
# # generate_number_dataset(n=100_000, k=k, save_dir=save_dir)
# # print(f"Dataset generated and saved at {save_dir}")
# for k in [512]:
# save_dir = f"data/raw_datasets/context_numbers_{k}"
# os.makedirs(save_dir, exist_ok=True)
# generate_number_dataset(n=100_000, k=k, save_dir=save_dir)
# print(f"Dataset generated and saved at {save_dir}")

83
data/generate_fav_num_big.py
View file

@ -1,83 +0,0 @@
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, save_dir: str = None
):
"""
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 text above."
# Generate all unique combinations of k numbers
combinations = get_random_combinations(numbers, n, k)
# random.shuffle(combinations)
for combination in combinations:
entry = {
"context": ", ".join(map(str, combination)),
"prompt": query,
"response": ", ".join(map(str, combination)),
}
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 :]
save_dir = "" if save_dir is None else save_dir
# Save splits to separate files
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")
if __name__ == "__main__":
# Set random seed for reproducibility
random.seed(42)
# Generate dataset
for k in [32, 64, 128, 256]:
save_dir = f"context_numbers_{k}_big"
os.makedirs(save_dir, exist_ok=True)
generate_number_dataset(n=240_000, k=k, save_dir=save_dir)
print(f"Dataset generated and saved at {save_dir}")

157
data/generate_fw_edu_augment_vllm.py
View file

@ -1,157 +0,0 @@
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()

2
data/generate_fw_edu_qa_v2.py
View file

@ -157,7 +157,7 @@ if __name__ == "__main__":
parser.add_argument(
"--max_length",
type=int,
default=10_000,
default=2000,
help="Maximum length of the context to consider for generation",
)
parser.add_argument(

317
data/generate_fw_edu_qa_v3.py
View file

@ -1,317 +0,0 @@
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
from ctx_to_lora.data.processing import load_and_process_dataset
STOP_STRINGS = {
"google/gemma-3-12b-it": ["<eos>", "<end_of_turn>"],
"google/gemma-3-4b-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-4b-it"),
help="VLLM model to use for generation",
)
parser.add_argument(
"--shard_pattern",
type=str,
default=None,
help="Pattern to match shard files (e.g., '000_0000*')",
)
parser.add_argument(
"--ds_names", # list of names
type=str,
nargs="+",
default=None,
help="dataset names",
)
parser.add_argument(
"--split",
type=str,
default=None,
help="Dataset split to use (e.g., 'train', 'validation')",
)
# 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()
assert bool(args.shard_pattern) ^ bool(args.ds_names)
if args.ds_names:
assert bool(args.split)
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
if args.shard_pattern is not None:
paths = glob(
f"./data/raw_datasets/fineweb_edu/sample/100BT/{shard_pattern}.parquet"
)
elif args.ds_names:
paths = args.ds_names
else:
raise ValueError("Either shard_pattern or ds_names must be provided.")
split = "train[:100]" if args.debug else args.split or "train"
for path in paths:
if args.shard_pattern is not None:
ds = load_dataset(
"parquet",
data_files=path,
split=split,
)
elif args.ds_names:
ds = load_and_process_dataset(path, args.split, num_proc=16)
ds = ds.rename_column("context", "text")
print(f"Loaded {len(ds)} samples from {path}")
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"
)

235
data/generate_fw_edu_qa_vllm.py
View file

@ -1,235 +0,0 @@
import os
import random
import re
import sys
from glob import glob
import pandas as pd
from datasets import Dataset, load_dataset
from vllm import LLM, SamplingParams
SYSTEM_TEMPLATE = (
"You are a creative and helpful assistant.\n"
"You are tasked with generating questions for reading comprehension tests.\n"
"You will be given a context and you need to generate questions and corresponding answers from the given context.\n"
"The questions should be highly specific to the information provided in the context, not general questions that suits any context.\n"
"**DO NOT** halucinate or make up information."
)
# based on Make Your LLM Fully Utilize the Context (https://arxiv.org/pdf/2404.16811)
PROMPT_TEMPLATE = (
"### Instructions ###\n"
"Generate questions and corresponding answers from the given context. The questions should be highly specific to the "
"information provided in the context, not general questions that suits any context.\n\n\n"
"### Context ###\n"
"{context}"
"\n\n\n"
"### Additional Instructions ###\n"
"Rules to follow when generate the questions:\n"
"1. The questions must be specific to the given context and fully answerable from information present in given context.\n"
"2. Make sure the questions are clear and unambiguous.\n"
"3. Phrases like 'based on the provided context', 'according to the context', etc, are **NOT ALLOWED** to appear in "
"the questions.\n"
"4. The questions should not overlap. They should be diverse, covering many aspects of the context.\n"
"5. 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"
"6. Ignore the text formatting of the context, e.g., bold, italic, underline, etc.\n"
"7. Ignore typos, spacing and grammatical errors in the context.\n\n"
"Rules to follow when generate the answers:\n"
"1. The answers must use the information provided in the context.\n"
"2. Do not just copy words from the context. Answer the question in your own words.\n"
"3. The answers should be detailed and comprehensive. Please include additional specific details from the context when possible.\n\n"
"Response with {n_qa_pairs} question-answer pairs.\n"
"Always use proper grammar and punctuation.\n"
"Try to use different question forms and styles.\n"
"Use simple words and make sure that the answers are clear and comprehensive.\n\n"
"The question-answer pairs should be in the following format:\n"
"Question 1: {{question_1}}\n"
"Answer 1: {{answer_1}}\n"
"Question 2: {{question_2}}\n"
"Answer 2: {{answer_2}}\n"
"..."
)
def get_prompt(context, n_qa_pairs):
prompt = PROMPT_TEMPLATE.format(context=context, n_qa_pairs=n_qa_pairs)
return prompt
def postprocess_qa_pairs(res_txt: str):
"""
Postprocesses the QA pairs from the response text.
Args:
res_txt: The response text.
n_qa_pairs: The number of QA pairs.
Returns:
A tuple of two lists, the first containing the questions and the second containing the answers.
"""
# capture everything after each "Question {number}:" until "Answer"
q_pattern = r"Question \d+:(.*?)(?=Answer|$)" # thanks chatgpt
questions = re.findall(q_pattern, res_txt, flags=re.S)
a_pattern = r"Answer \d+:(.*?)(?=Question|$)" # thanks chatgpt
answers = re.findall(a_pattern, res_txt, flags=re.S)
if len(questions) != len(answers):
print(f"Warning---number of questions and answers do not match")
print(f"Number of questions: {len(questions)}")
print(f"Number of answers: {len(answers)}")
out_q = []
out_a = []
if (len(questions) > 0) and (len(answers) > 0):
for i in range(min(len(questions), len(answers))):
out_q.append(questions[i].strip())
out_a.append(answers[i].strip())
return out_q, out_a
def length_filter(samples):
return [len(text) < 10_000 for text in samples["text"]]
if __name__ == "__main__":
# api_key = os.environ.get("vllm_api_key")
# vllm_port = os.environ.get("vllm_port")
# vllm_ip = os.environ.get("vllm_ip") # "172.16.0.62"
vllm_model = os.environ.get("vllm_model") # "google/gemma-2-27b-it"
print(f"Using model: {vllm_model}")
# print(f"Using API key: {api_key}")
# print(f"Using VLLM IP: {vllm_ip}")
# print(f"Using VLLM port: {vllm_port}")
# client = OpenAI(
# base_url=f"http://{vllm_ip}:{vllm_port}/v1",
# api_key=api_key,
# )
mm_kwargs = {"image": 0} # if "gemma-3" in vllm_model else {}
llm_kwargs = dict(
model=vllm_model,
dtype="bfloat16",
# tokenizer_mode="mistral",
# config_format="mistral",
# load_format="mistral",
enable_prefix_caching=True,
max_model_len=2**14, # 11264,
limit_mm_per_prompt=mm_kwargs,
# enable_chunked_prefill=True,
)
if vllm_model == "mistralai/Mistral-Small-3.1-24B-Instruct-2503":
llm_kwargs["tokenizer_mode"] = "mistral"
llm_kwargs["config_format"] = "mistral"
llm_kwargs["load_format"] = "mistral"
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]
n_qa_pairs = int(sys.argv[2])
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)]
messages = [
[
{"role": "system", "content": SYSTEM_TEMPLATE},
{"role": "user", "content": get_prompt(ctx, n_qa_pairs)},
]
for ctx in ctxs
]
print(f"Generating from {len(messages)} contexts")
completions = llm.chat(
messages,
sampling_params=SamplingParams(
max_tokens=2048,
temperature=(
0.1
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):
questions, answers = postprocess_qa_pairs(completion.outputs[0].text)
for q, a in zip(questions, answers):
samples.append(
{
"context": ctx,
"prompt": q,
"response": a,
}
)
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_3/{shard_name}.parquet")
val_ds.to_parquet(f"data/raw_datasets/fw_qa_3/{shard_name}_val.parquet")
print(f"Saved to data/raw_datasets/fw_qa_3/{shard_name}.parquet")
print(f"Saved to data/raw_datasets/fw_qa_3/{shard_name}_val.parquet")
# debug
# ds = load_dataset(
# "parquet",
# data_files=paths[0],
# split="train[:10]",
# )
# ctxs = [sample["text"] for sample in iter(ds)]
# messages = [
# [
# {"role": "system", "content": SYSTEM_TEMPLATE},
# {"role": "user", "content": get_prompt(ctx, n_qa_pairs)},
# ]
# for ctx in ctxs
# ]
#
# print(f"Generating from {len(messages)} contexts")
# completions = llm.chat(
# messages,
# sampling_params=SamplingParams(
# max_tokens=2048,
# temperature=0.1 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):
# questions, answers = postprocess_qa_pairs(completion.outputs[0].text)
# for q, a in zip(questions, answers):
# samples.append(
# {
# "context": ctx,
# "prompt": q,
# "response": a,
# }
# )
# for sample in samples:
# print(f"{sample['context']=}")
# print(f"{sample['prompt']=}")
# print(f"{sample['response']=}")
# print()

64
data/generate_negative_natural_questions.py
View file

@ -1,64 +0,0 @@
import os
from datasets import load_dataset
def process_sample(sample):
"""Process a single sample and return the processed data or None if invalid."""
ctx = " ".join(
[
token
for is_html, token in zip(
sample["document"]["tokens"]["is_html"],
sample["document"]["tokens"]["token"],
)
if not is_html
]
).strip()
if not ctx:
return {"prompt": None, "context": None, "answer": None}
answers = []
for answer in sample["annotations"]["short_answers"]:
answers += answer["text"]
if not answers:
return {"prompt": None, "context": None, "answer": None}
answer = answers[0].strip()
if not answer:
return {"prompt": None, "context": None, "answer": None}
prompt = sample["question"]["text"].capitalize().strip() + " ?"
return {"prompt": prompt, "context": ctx, "answer": answer.capitalize()}
def shift_contexts(examples):
"""Shift contexts by one position to create negative examples."""
contexts = examples["context"]
shifted_contexts = contexts[1:] + contexts[:1]
examples["context"] = shifted_contexts
return examples
if __name__ == "__main__":
ds = load_dataset("google-research-datasets/natural_questions", split="validation")
ds = ds.shuffle(seed=42)
# for split in ["validation"]: # , "train"
# Process samples in parallel using map
processed_ds = ds.map(process_sample, remove_columns=ds.column_names, num_proc=16)
# Filter out None values
processed_ds = processed_ds.filter(lambda x: x["context"] is not None)
# Shift contexts to create negative examples
processed_ds = processed_ds.map(
shift_contexts, batched=True, batch_size=len(processed_ds)
)
print(processed_ds)
save_dir = "data/raw_datasets/negative_natural_questions"
os.makedirs(save_dir, exist_ok=True)
processed_ds.to_json(f"{save_dir}/validation.jsonl")

37
data/generate_pretrain_from_fw_qa.py
View file

@ -1,37 +0,0 @@
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 = 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")
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, samples, ctx_qa_dict
gc.collect()

143
data/prepare_distillation_dataset.py
View file

@ -1,143 +0,0 @@
import argparse
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from ctx_to_lora.data.processing import load_and_process_dataset
def add_logits(row, model, tok, device, args):
ctx = row["context"]
qs = row["prompts"]
ans = row["responses"]
responses_masks = []
all_logits_values = []
all_logits_positions = []
for question, answer in zip(qs, ans):
messages = [
{
"role": "user",
"content": f"Answer the question based on the following information:\n{ctx}.\n\n{question}",
},
{"role": "assistant", "content": answer},
]
all_tokens = tok.apply_chat_template(
messages,
tokenize=True,
add_special_tokens=False,
truncation=False,
padding=False,
add_generation_prompt=False,
return_dict=True,
)
input = torch.tensor([all_tokens["input_ids"]], device=device)
with torch.no_grad():
logits = model(input).logits
top_logits, top_logits_positions = torch.topk(
logits, args.max_logits_to_store, dim=-1
)
# Convert logits to specified precision
dtype_map = {
"float32": torch.float32,
"bfloat16": torch.bfloat16,
"float16": torch.float16,
}
target_dtype = dtype_map[args.logit_precision]
top_logits = top_logits.to(target_dtype)
tokens_to_mask = tok.apply_chat_template(
[messages[0]], # Just the user message
tokenize=True,
add_special_tokens=False,
truncation=False,
padding=False,
add_generation_prompt=True,
return_dict=True,
)
number_of_tokens_to_mask = len(tokens_to_mask["input_ids"])
# Create mask: False for prompt tokens, True for response tokens
responses_mask = torch.zeros(len(all_tokens["input_ids"]), dtype=torch.bool)
responses_mask[number_of_tokens_to_mask:] = True
responses_masks.append(responses_mask)
all_logits_values.append(top_logits.cpu())
all_logits_positions.append(top_logits_positions.cpu())
row["labels_logits_values"] = all_logits_values
row["labels_logits_positions"] = all_logits_positions
row["responses_mask"] = responses_masks
return row
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--ds_name", required=True, help="name of the dataset")
ap.add_argument("--split", required=True, help="split of the dataset")
ap.add_argument(
"--max_logits_to_store",
type=int,
default=100,
help="max number of logits kept per completion",
)
ap.add_argument(
"--limit_samples",
type=int,
default=None,
help="process only the first n dataset rows",
)
ap.add_argument(
"--logit_precision",
default="bfloat16",
choices=["float32", "bfloat16", "float16"],
help="dtype for saved logits",
)
ap.add_argument(
"--model_name", required=True, help="model name to use for generating logits"
)
ap.add_argument(
"--output_path", required=True, help="path to save the processed dataset"
)
args = ap.parse_args()
# Load model and tokenizer
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = AutoModelForCausalLM.from_pretrained(args.model_name)
model = model.to(device)
model.eval()
tok = AutoTokenizer.from_pretrained(args.model_name)
# Load dataset
ds = load_and_process_dataset(
args.ds_name,
args.split,
add_negative_prompt=False,
add_repeat_prompt=False,
repeat_prob=0,
is_pretrain=False,
streaming=False,
num_proc=8,
)
# Limit samples if specified
if args.limit_samples:
ds = ds.select(range(min(args.limit_samples, len(ds))))
# Process dataset with logits
ds = ds.map(
lambda row: add_logits(row, model, tok, device, args),
batched=False,
desc="Adding logits to dataset",
)
# Save the processed dataset
ds.save_to_disk(args.output_path)
print(f"Dataset saved to {args.output_path}")
if __name__ == "__main__":
main()

114
data/q_generation_prompts.py
View file

@ -1,114 +0,0 @@
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."

124
data/self_generate_qa.py
View file

@ -11,11 +11,8 @@ from vllm import LLM, SamplingParams
from ctx_to_lora.data.definitions import (
CLOSED_QA_INTX_TEMPLATES,
COT_TEMPLATES,
RAW_DATA_DIR,
SELF_GEN_DATA_DIR,
STRUCTURING_PROMPTS,
SUMMARIZATION_PROMPTS,
)
from ctx_to_lora.data.processing import (
filter_none,
@ -67,29 +64,15 @@ def truncate_middle_if_too_long(
return input_ids
# TODO: self-gen for magic ctx num
def get_prompt(context: str, q: str, remove_qa_template: bool) -> str:
prompt = QA_PROMPT_TEMPLATE if not remove_qa_template else PROMPT_TEMPLATE
return prompt.format(context=context, question=q)
def augment_prompt(sample, closed_qa_prob, cot_prob):
z = random.random()
should_add_closed_qa = ("type" not in sample) or (sample["type"] == "question")
if should_add_closed_qa and (z < closed_qa_prob):
sample["prompts"] = add_closed_qa_prompt(sample["prompts"])
elif closed_qa_prob <= z < (closed_qa_prob + cot_prob):
sample["prompts"] = add_cot_prompt(sample["prompts"])
return sample
def add_cot_prompt(prompts: list[str]) -> str:
return [random.choice(COT_TEMPLATES).format(input=q) for q in prompts]
def add_closed_qa_prompt(prompts: list[str]) -> str:
return [random.choice(CLOSED_QA_INTX_TEMPLATES).format(input=q) for q in prompts]
def add_closed_qa_prompt(q: str, closed_qa_prob: float = 0.1) -> str:
if random.random() <= closed_qa_prob:
q = random.choice(CLOSED_QA_INTX_TEMPLATES).format(input=q)
return q
def load_config(config_path: str) -> dict:
@ -99,49 +82,6 @@ def load_config(config_path: str) -> dict:
return config
def add_paraphrasing_prompt(samples, summ_prob: float, struct_prob: float):
"""Augment batch with summarization / structuring prompts given separate probabilities.
The original samples are extended (in-place) with newly created paraphrased variants.
"""
if summ_prob == 0 and struct_prob == 0:
return samples
n_samples = len(samples["ctx_ids"]) if "ctx_ids" in samples else 0
if n_samples == 0:
return samples
n_summ = int(n_samples * summ_prob)
n_struct = int(n_samples * struct_prob)
all_indices = list(range(n_samples))
random.shuffle(all_indices)
summ_indices = all_indices[:n_summ]
struct_indices = all_indices[n_summ : n_summ + n_struct]
def _add(kind_indices, prompts_source, prompt_type):
out = dict()
if not kind_indices:
return out
for k in samples:
if k == "prompts":
sampled_prompts = random.choices(prompts_source, k=len(kind_indices))
out[k] = [[p] for p in sampled_prompts]
elif k == "type":
out[k] = [prompt_type] * len(kind_indices)
else:
out[k] = [samples[k][i] for i in kind_indices]
return out
summ_samples = _add(summ_indices, SUMMARIZATION_PROMPTS, "summarization")
struct_samples = _add(struct_indices, STRUCTURING_PROMPTS, "structuring")
# Extend original samples
for k in samples:
if summ_indices:
samples[k] += summ_samples[k]
if struct_indices:
samples[k] += struct_samples[k]
return samples
def get_dataset_configs(
ds_names: list[str] | None,
config: dict | None,
@ -334,46 +274,18 @@ def self_generate(
keep_in_memory=True,
)
ds = ds.map(
augment_prompt,
fn_kwargs={"closed_qa_prob": closed_qa_prob, "cot_prob": cot_prob},
keep_in_memory=True,
num_proc=16,
)
ds = ds.map(
add_paraphrasing_prompt,
batched=True,
fn_kwargs={
"summ_prob": summ_prob,
"struct_prob": struct_prob,
},
batch_size=50_000,
keep_in_memory=True,
)
ctxs = [sample["context"] for sample in ds]
questions = [
[add_closed_qa_prompt(q, closed_qa_prob) for q in sample["prompts"] if q]
for sample in ds
]
questions = [q_list for q_list in ds["prompts"] if len(q_list) > 0]
# for sample in ds:
# if ("type" not in sample) or (
# "type" in sample and sample["type"] == "question"
# ):
# # for generated qa, apply to only "question" type
# questions.append(
# [
# add_closed_qa_prompt(q, closed_qa_prob)
# for q in sample["prompts"]
# if q
# ]
# )
# else:
# questions.append([q for q in sample["prompts"] if q])
print(f"Loaded {len(ctxs)} contexts and {len(questions)} questions")
k = 16
fpath = f"{SELF_GEN_DATA_DIR}/{args.vllm_model}_temp_{temp}_closed_qa_prob_{closed_qa_prob}_cot_prob_{cot_prob}_summ_prob_{summ_prob}_struct_prob_{struct_prob}/{ds_name}/{split}/ds{shard_name}"
fpath = f"{SELF_GEN_DATA_DIR}/{args.vllm_model}_temp_{temp}_closed_qa_prob_{closed_qa_prob}/{ds_name}/{split}/ds{shard_name}"
chunk_size = 1_000
for chunk_idx, start in enumerate(range(0, len(ctxs), chunk_size)):
@ -670,24 +582,6 @@ def parse_args() -> argparse.Namespace:
default=0.0,
help="Probability of using closed QA prompt template (default: 0.0)",
)
parser.add_argument(
"--cot_prob",
type=float,
default=0.0,
help="Probability of using chain-of-thought prompt template (default: 0.0)",
)
parser.add_argument(
"--summ_prob",
type=float,
default=0.0,
help="Probability of adding a summarization variant (default: 0.0)",
)
parser.add_argument(
"--struct_prob",
type=float,
default=0.0,
help="Probability of adding a structuring variant (default: 0.0)",
)
parser.add_argument(
"--do_truncate",
action="store_true",

126
data/store_top_logits.py
View file

@ -1,126 +0,0 @@
import argparse
import io
import zipfile
from pathlib import Path
from tqdm import tqdm
import torch
from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForCausalLM, PreTrainedTokenizer
# e.g., python data/store_top_logits.py --data "temp/test_data.parquet" --model "Qwen/Qwen3-0.6B"
def main():
ap = argparse.ArgumentParser()
ap.add_argument('--data', required=True,
help='path to a parquet shard of the dataset')
ap.add_argument('--model', required=True,
help='huggingface model name or checkpoint')
ap.add_argument('--max_logits_to_store', type=int, default=100,
help='max number of logits kept per completion')
ap.add_argument('--limit_samples', type=int, default=None,
help='process only the first n dataset rows')
ap.add_argument('--logit_precision', default='bfloat16',
choices=['float32', 'bfloat16', 'float16'],
help='dtype for saved logits')
ap.add_argument('--index_precision', default='int32',
choices=['int64', 'int32', 'int16'],
help='dtype for saved indices')
args = ap.parse_args()
float_map = {
'float32': torch.float32,
'bfloat16': torch.bfloat16,
'float16': torch.float16,
}
int_map = {
'int64': torch.int64,
'int32': torch.int32,
'int16': torch.int16,
}
l_dtype = float_map[args.logit_precision]
i_dtype = int_map[args.index_precision]
ds = load_dataset('parquet', data_files=args.data)['train']
if args.limit_samples and args.limit_samples > 0:
ds = ds.select(range(min(args.limit_samples, len(ds))))
tok: PreTrainedTokenizer = AutoTokenizer.from_pretrained(
args.model, use_fast=True)
model = AutoModelForCausalLM.from_pretrained(args.model)
model.eval()
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)
base = Path(args.data).with_suffix('')
model_dir = args.model.replace('/', '__')
out_dir = base / model_dir
out_dir.mkdir(parents=True, exist_ok=True)
if args.limit_samples:
zip_name = f'limited_logits_{args.limit_samples}.zip'
else:
zip_name = 'all_logits.zip'
zip_path = out_dir / zip_name
samp_id = 0
with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zf:
for row in tqdm(ds):
ctx = row['context']
qs = row['prompts_level_0']
ans = row['responses_level_0']
q_list = []
a_list = []
vals_list = []
sel_pos_list = []
for q, a in zip(qs, ans):
messages = [
{'role': 'user', 'content': ctx + '\n' + q},
{'role': 'assistant', 'content': a},
]
tpl = tok.apply_chat_template(
messages,
tokenize=True,
add_special_token=False,
truncation=False,
add_generation_prompt=False,
return_dict=True,
)
tokens_to_mask = tok.apply_chat_template(
messages[:-1],
tokenize=True,
add_special_token=False,
truncation=False,
add_generation_prompt=True,
return_dict=True,
)
number_of_tokens_to_mask = len(tokens_to_mask['input_ids'])
assert tokens_to_mask['input_ids'] == tpl['input_ids'][:number_of_tokens_to_mask]
inp = torch.tensor([tpl['input_ids']], device=device)
with torch.no_grad():
logits = model(inp).logits
logits_for_answer = logits[0, number_of_tokens_to_mask - 1:, :]
assert logits_for_answer.shape[0] > 0
assert logits_for_answer.shape[0] + number_of_tokens_to_mask -1 == inp.size(1)
# just in case, logits to store is super large
k = min(args.max_logits_to_store, logits_for_answer.shape[-1])
vals, idxs = torch.topk(logits_for_answer, k, dim=-1)
q_list.append(q)
a_list.append(a)
vals_list.append(vals.cpu())
sel_pos_list.append(idxs.cpu())
record = {
'context': ctx,
'question': q_list,
'answer': a_list,
'top_logits': vals_list,
'top_positions_positions': sel_pos_list,
}
buf = io.BytesIO()
torch.save(record, buf)
samp_id += 1
zf.writestr(f'{samp_id:07d}.pt', buf.getvalue())
if __name__ == '__main__':
main()