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

self gen data for all ds both train + val

Browse source
This commit is contained in:
51616 2025-06-03 14:30:38 +00:00
parent 57b63d3b8b
commit 5aa83b96e2
5 changed files with 403 additions and 111 deletions
Download patch file Download diff file Expand all files Collapse all files

66
configs/self_gen_3_mini.yaml Normal file
View file

@ -0,0 +1,66 @@
output_dir: "" # just a placeholder
bf16: true
model_name_or_path: meta-llama/Llama-3.2-1B-Instruct
label_names: ["labels"]
# eval_on_start: True
# eval_strategy: "steps"
# eval_steps: 500
# save_strategy: "no"
# # save_steps: 500
# logging_strategy: "steps"
# logging_steps: 100
# use_liger_kernel: true
# remove_unused_columns: false
# needed to avoid OOM by compute the metrics batch by batch
# w/o this the trainer stores logits of all sample in memory...
# batch_eval_metrics: true
per_device_train_batch_size: 8
per_device_eval_batch_size: 8
max_val_samples_per_ds: 1000
# optim: schedule_free_adamw
learning_rate: 0.00004
# lr_scheduler_type: "constant_with_warmup"
neftune_noise_alpha: 5
weight_decay: 0.01
#
warmup_steps: 100
dataloader_prefetch_factor: 8
dataloader_num_workers: 8
# LoRA
lora_r: 8
lora_dropout: 0.0
target_modules:
- down_proj
# data
train_ds_names:
- self_gen/fw_qa_3_mini # 100k
- self_gen/ctx_qa # 300k
- self_gen/pwc # 240k
- self_gen/hotpot_qa # 90k
- self_gen/squad # 90k
- self_gen/drop # 77k
- self_gen/narrativeqa # 40k
- self_gen/quoref # 11k
- self_gen/ropes # 11k
- self_gen/synthetic_convqa # 40k
val_ds_names:
- self_gen/fw_qa_3
- self_gen/fw_qa_xl
- self_gen/ctx_qa
- self_gen/pwc
- self_gen/hotpot_qa
- self_gen/squad
- fw_qa_3
- fw_qa_xl
- ctx_qa
- pwc
- hotpot_qa
- squad
load_best_model_at_end: true
metric_for_best_model: eval_pwc_loss

47
configs/self_gen_pwc_hotpot_qa.yaml Normal file
View file

@ -0,0 +1,47 @@
output_dir: "" # just a placeholder
bf16: true
model_name_or_path: google/gemma-2-2b-it
label_names: ["labels"]
# eval_on_start: True
# eval_strategy: "steps"
# eval_steps: 500
# save_strategy: "no"
# # save_steps: 500
# logging_strategy: "steps"
# logging_steps: 100
# use_liger_kernel: true
# remove_unused_columns: false
# needed to avoid OOM by compute the metrics batch by batch
# w/o this the trainer stores logits of all sample in memory...
# batch_eval_metrics: true
per_device_train_batch_size: 8
per_device_eval_batch_size: 8
max_val_samples_per_ds: 1000
# optim: schedule_free_adamw
learning_rate: 0.00004
# lr_scheduler_type: "constant_with_warmup"
neftune_noise_alpha: 1
weight_decay: 0.01
warmup_steps: 100
dataloader_prefetch_factor: 8
dataloader_num_workers: 8
# LoRA
lora_r: 8
lora_dropout: 0.02
target_modules:
- down_proj
# data
train_ds_names:
- self_gen/pwc
- self_gen/hotpot_qa
val_ds_names:
- self_gen/pwc
- self_gen/hotpot_qa

340
data/self_generate_qa.py
View file

@ -1,16 +1,21 @@
import argparse
import os
import random
import sys
from glob import glob
import pandas as pd
import yaml
from datasets import Dataset, load_dataset
from vllm import LLM, SamplingParams
from ctx_to_lora.data.definitions import SELF_GEN_DATA_DIR
from ctx_to_lora.data.processing import load_and_process_dataset
from ctx_to_lora.utils import clear_gpu
SYSTEM_TEMPLATE = (
"### System Instructions ###\n"
"You are a creative and helpful assistant.\n"
"**DO NOT** halucinate or make up information."
"**DO NOT** hallucinate or make up information."
)
PROMPT_TEMPLATE = "### Context ###\n{context}\n\n\n### Question ###\n{question}"
@ -22,15 +27,214 @@ MODEL_CTX_LEN = {
}
def get_prompt(context, q):
prompt = PROMPT_TEMPLATE.format(context=context, question=q)
return prompt
def get_prompt(context: str, q: str) -> str:
return PROMPT_TEMPLATE.format(context=context, question=q)
def load_config(config_path: str) -> dict:
"""Load dataset names from YAML config file."""
with open(config_path) as f:
config = yaml.safe_load(f)
return config
def get_dataset_configs(
ds_names: list[str] | None,
config: dict | None,
split: str | None,
) -> list[tuple[str, str]]:
assert not (ds_names and config), "Cannot provide both ds_names and config"
if ds_names:
assert split, "When using ds_names, --split must be provided"
# Validate ds_names format
for ds_name in ds_names:
if not isinstance(ds_name, str):
raise ValueError(f"Invalid dataset name: {ds_name}")
return [(ds_name, split) for ds_name in ds_names]
if config:
dataset_configs = []
# Process train datasets
train_ds_names = config.get("train_ds_names", [])
self_gen_train_ds_names = [
(ds_name, "train")
for ds_name in train_ds_names
if ds_name.startswith("self_gen/")
]
if not self_gen_train_ds_names:
print("No self_gen datasets found in train_ds_names")
dataset_configs.extend(self_gen_train_ds_names)
# Process validation datasets
val_ds_names = config.get("val_ds_names", [])
self_gen_val_ds_names = [
(ds_name, "validation")
for ds_name in val_ds_names
if ds_name.startswith("self_gen/")
]
if not self_gen_val_ds_names:
print("No self_gen datasets found in val_ds_names")
dataset_configs.extend(self_gen_val_ds_names)
return dataset_configs
def create_messages(
ctxs: list[str], questions: list[str], vllm_model: str, system_template: str
) -> list[list[dict]]:
"""Create chat messages for the model."""
if "gemma" in vllm_model:
# gemma models do not support system messages
return [
[
{
"role": "user",
"content": system_template + "\n\n\n" + get_prompt(ctx, q),
}
]
for ctx, q in zip(ctxs, questions)
]
else:
return [
[
{"role": "system", "content": system_template},
{"role": "user", "content": get_prompt(ctx, q)},
]
for ctx, q in zip(ctxs, questions)
]
def self_generate(
ds_name: str,
split: str,
args: argparse.Namespace,
llm: LLM,
system_template: str,
parquet_file: str | None = None,
) -> None:
"""Process a single dataset and generate QA pairs."""
shard_name = ""
if parquet_file:
print(f"Loading dataset from parquet file: {parquet_file}")
ds = load_dataset(path="parquet", data_files=[parquet_file], split="train")
split = "train"
ds_name = os.path.basename(os.path.dirname(parquet_file))
shard_name = "_" + os.path.basename(parquet_file).replace(".parquet", "")
else:
ds_name = ds_name.replace("self_gen/", "")
print(f"Loading dataset: {ds_name} with split: {split}")
kwargs = dict(
ds_name=ds_name,
split=split,
add_negative_prompt=False,
add_repeat_prompt=False,
repeat_prob=0,
)
ds = load_and_process_dataset(**kwargs, streaming=False, num_proc=8)
if args.debug:
ds = ds.take(10)
ctxs = [sample["context"] for sample in ds]
questions = [sample["prompt"] for sample in ds]
messages = create_messages(ctxs, questions, args.vllm_model, system_template)
print(f"Generating from {len(messages)} contexts")
completions = llm.chat(
messages,
sampling_params=SamplingParams(
max_tokens=2048,
temperature=0.1
if args.vllm_model == "mistralai/Mistral-Small-3.1-24B-Instruct-2503"
else 1.0,
),
)
samples = [
{"context": ctx, "prompt": q, "response": completion.outputs[0].text}
for ctx, q, completion in zip(ctxs, questions, completions)
]
if args.debug:
for sample in samples:
print(f"context={sample['context']}")
print(f"prompt={sample['prompt']}")
print(f"response={sample['response']}")
print("=" * 80)
print(f"Generated {len(samples)} samples")
random.shuffle(samples)
# Save results
df = pd.DataFrame(samples)
ds_out = Dataset.from_pandas(df)
fpath = f"{SELF_GEN_DATA_DIR}/{args.vllm_model}/{ds_name}/{split}/ds{shard_name}"
os.makedirs(os.path.dirname(fpath), exist_ok=True)
fpath = f"{fpath}.parquet"
ds_out.to_parquet(fpath)
print(f"Saved to {fpath}")
# Cleanup
del samples, df, ds_out, completions, messages, ctxs, questions
clear_gpu()
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="Generate QA pairs using VLLM")
parser.add_argument(
"--vllm_model",
type=str,
required=True,
help="VLLM model name (e.g., google/gemma-2-2b-it)",
)
parser.add_argument(
"--debug",
action="store_true",
help="Enable debug mode (process only 10 samples)",
)
# Either config file OR ds_names + split
group = parser.add_mutually_exclusive_group(required=True)
group.add_argument(
"--config",
type=str,
help="Path to YAML config file with train_ds_names/val_ds_names",
)
group.add_argument(
"--ds_names",
type=str,
nargs="+",
help="List of dataset names/shard patterns",
)
group.add_argument(
"--glob_pattern",
type=str,
help="Glob pattern to match dataset names (e.g., 'data/raw_datasets/fw_qa_3/*')",
)
parser.add_argument(
"--split",
type=str,
help="Dataset split to use when using --ds_names (required with --ds_names)",
)
return parser.parse_args()
if __name__ == "__main__":
vllm_model = os.environ.get("vllm_model") # "google/gemma-2-27b-it"
debug = bool(os.environ.get("debug", False))
args = parse_args()
# Validate arguments
if args.ds_names and not args.split:
raise ValueError("--split is required when using --ds_names")
vllm_model = args.vllm_model
print(f"Using model: {vllm_model}")
# Setup model-specific configurations
llm_kwargs = dict(
model=vllm_model,
dtype="bfloat16",
@ -39,98 +243,52 @@ if __name__ == "__main__":
max_model_len=MODEL_CTX_LEN.get(vllm_model),
)
system_template = SYSTEM_TEMPLATE
if vllm_model == "mistralai/Mistral-Small-3.1-24B-Instruct-2503":
mm_kwargs = {"image": 0} # if "gemma-3" in vllm_model else {}
llm_kwargs["tokenizer_mode"] = "mistral"
llm_kwargs["config_format"] = "mistral"
llm_kwargs["load_format"] = "mistral"
SYSTEM_TEMPLATE = (
mm_kwargs = {"image": 0}
llm_kwargs.update(
{
"tokenizer_mode": "mistral",
"config_format": "mistral",
"load_format": "mistral",
"max_model_len": 2**14,
"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_kwargs["max_model_len"] = 2**14
llm_kwargs["limit_mm_per_prompt"] = mm_kwargs
print(f"{llm_kwargs=}")
llm = LLM(**llm_kwargs)
tokenizer = llm.get_tokenizer()
shard_pattern = sys.argv[1]
paths = glob(f"./data/raw_datasets/fw_qa_3/{shard_pattern}.parquet")
for path in paths:
if debug:
ds = load_dataset(
"parquet",
data_files=path,
split="train[:10]",
)
else:
ds = load_dataset(
"parquet",
data_files=path,
split="train",
# streaming=True,
)
ctxs = []
questions = []
for sample in iter(ds):
ctxs.append(sample["context"])
questions.append(sample["prompt"])
if "gemma" in vllm_model:
# gemma models do not support system messages
messages = [
[
{
"role": "user",
"content": SYSTEM_TEMPLATE + "\n\n\n" + get_prompt(ctx, q),
},
]
for ctx, q in zip(ctxs, questions)
]
else:
messages = [
[
{"role": "system", "content": SYSTEM_TEMPLATE},
{"role": "user", "content": get_prompt(ctx, q)},
]
for ctx, q in zip(ctxs, questions)
]
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 1.0
),
),
# Get dataset configs from config or CLI args
config = load_config(args.config) if args.config else None
if args.ds_names or args.config:
dataset_configs = get_dataset_configs(
ds_names=args.ds_names,
config=config,
split=args.split,
)
samples = []
for ctx, q, completion in zip(ctxs, questions, completions):
response = completion.outputs[0].text
samples.append(
{
"context": ctx,
"prompt": q,
"response": response,
}
# 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)
else:
assert args.glob_pattern, (
"glob_pattern must be provided if no ds_names or config"
)
files = glob(args.glob_pattern)
for file in files:
print(f"Processing file: {file}")
self_generate(
ds_name=None,
parquet_file=file,
split=args.split,
args=args,
llm=llm,
system_template=system_template,
)
if debug:
print(f"{ctx=}")
print(f"{q=}")
print(f"{response=}")
print("=" * 80)
print(f"Generated {len(samples)} samples")
random.shuffle(samples)
df = pd.DataFrame(samples)
ds = Dataset.from_pandas(df)
shard_name = path.split("/")[-1].split(".")[0]
ds.to_parquet(f"data/raw_datasets/fw_qa_3/{vllm_model}/{shard_name}.parquet")
print(f"Saved to data/raw_datasets/fw_qa_3/{vllm_model}/{shard_name}.parquet")

2
src/ctx_to_lora/data/definitions.py
View file

@ -7,6 +7,8 @@ FW_QA_PATHS = [
]
TRANSFORMED_DATA_DIR = "data/processed_datasets"
SELF_GEN_DATA_DIR = "data/raw_datasets/self_gen/"
REPEAT_PROMPTS = [
"Repeat the text.",

59
src/ctx_to_lora/data/processing.py
View file

@ -3,6 +3,7 @@ import json
import logging
import random
from collections.abc import Callable, Iterator
from glob import glob
from os import path
from typing import Any
@ -17,6 +18,7 @@ from ctx_to_lora.data.definitions import (
EVAL_INTX_TEMPLATES,
IGNORE_INDEX,
REPEAT_PROMPTS,
SELF_GEN_DATA_DIR,
TRANSFORMED_DATA_DIR,
)
from ctx_to_lora.utils import TRAINING_TASK
@ -249,9 +251,20 @@ def get_ds_kwargs(ds_name: str, split: str) -> dict[str, Any]:
f"No dataset kwargs found for '{ds_name}' with split '{split}'.\n"
f"Using default kwargs: {kwargs}"
)
# return kwargs
elif ds_name.startswith("self_gen/"):
# e.g., "self_gen/google/gemma-2-2b-it/pwc/val"
base_model_name = "/".join(ds_name.split("/")[1:-1])
base_ds = ds_name.split("/")[-1]
files = glob(
f"{SELF_GEN_DATA_DIR}/{base_model_name}/{base_ds}/{split}/*.parquet"
)
if not files:
raise FileNotFoundError(
f"No self-gen files found for base model {base_model_name} "
f"in {SELF_GEN_DATA_DIR}/{base_model_name}/{base_ds}/"
)
kwargs = dict(path="parquet", data_files=files, split="train")
else:
# return DS_KWARGS[ds_name][split]
kwargs = DS_KWARGS[ds_name][split]
# # custom logic for slicing iterable datasets
@ -384,16 +397,32 @@ def filter_none(samples):
return out
def _load_and_process_dataset(
def load_and_process_dataset(
ds_name: str,
split: str,
add_negative_prompt: bool,
add_repeat_prompt: bool,
repeat_prob: float,
streaming: bool,
ds_path: str,
num_proc: int,
):
load_and_process_kwargs = dict(
ds_name=ds_name,
split=split,
add_negative_prompt=add_negative_prompt,
add_repeat_prompt=add_repeat_prompt,
repeat_prob=repeat_prob,
streaming=streaming,
)
ds_hash = hashlib.sha256(json.dumps(load_and_process_kwargs).encode()).hexdigest()
ds_path = f"{TRANSFORMED_DATA_DIR}/{ds_hash}"
if path.exists(ds_path) and split == "train":
# load the cached ds
logger.info(f"Loaded processed dataset from {ds_path}")
return datasets.load_from_disk(ds_path)
else:
logger.info(f"Loading dataset {ds_name} with split {split}...")
try:
ds_kwargs = get_ds_kwargs(ds_name, split)
skip = ds_kwargs.pop("skip", None)
@ -457,6 +486,8 @@ def _load_and_process_dataset(
num_proc=num_proc,
)
ds.save_to_disk(ds_path, num_proc=num_proc)
# force loading the cached version for newly created ds
ds = datasets.load_from_disk(ds_path)
return ds
@ -493,22 +524,10 @@ def get_tokenized_dataset(
)
num_proc = None if streaming and split == "train" else 8
ds_hash = hashlib.sha256(json.dumps(load_and_process_kwargs).encode()).hexdigest()
ds_path = f"{TRANSFORMED_DATA_DIR}/{ds_hash}"
if path.exists(ds_path) and split == "train":
# load the cached ds
logger.info(f"Loaded processed dataset from {ds_path}")
else:
logger.info(f"Loading dataset {ds_name} with split {split}...")
ds = _load_and_process_dataset(
**load_and_process_kwargs,
ds_path=ds_path,
num_proc=num_proc,
)
if split == "train":
# force loading the cached version for newly created ds
ds = datasets.load_from_disk(ds_path)
ds = load_and_process_dataset(
**load_and_process_kwargs,
num_proc=num_proc,
)
tokenized_ds = construct_and_tokenize_ctx_qa(
base_model_max_len,