mirror of
https://github.com/SakanaAI/doc-to-lora.git
synced 2026-07-23 17:01:04 +02:00
99 lines
3.2 KiB
Python
99 lines
3.2 KiB
Python
import itertools
|
|
import json
|
|
import os
|
|
import random
|
|
from typing import Dict, List
|
|
|
|
|
|
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.choices(range(len(numbers)), k=n) for _ in range(k)]
|
|
return zip(*[[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"What's your top-{k} favourite numbers in [0-999]? Answer with only the numbers separated by commas."
|
|
|
|
# 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 range(2, 129):
|
|
save_dir = f"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, 1024, 2048]:
|
|
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}")
|