AxiomAI-1.0 / generate_multilingual.py
yannyoper's picture
Upload 2 files
cf6354c verified
import json
import argparse
import multiprocessing as mp
from faker import Faker
import os
import random
def generate_chunk(chunk_id, num_examples, languages, output_dir, chunk_size=10000):
"""
Генерирует чанк датасета.
Если languages — список, то для каждого примера случайно выбирается язык.
"""
fakers = {lang: Faker(lang) for lang in languages}
filename = os.path.join(output_dir, f"chunk_{chunk_id:05d}.jsonl")
with open(filename, 'w', encoding='utf-8') as f:
for _ in range(num_examples):
lang = random.choice(languages)
fake = fakers[lang]
text = fake.paragraph(nb_sentences=random.randint(5, 15), variable_nb_sentences=True)
record = {"text": text, "language": lang}
f.write(json.dumps(record, ensure_ascii=False) + '\n')
return filename
def main():
parser = argparse.ArgumentParser(description="Генерация мультиязычного синтетического датасета в формате JSONL")
parser.add_argument("--total_examples", type=int, default=1_000_000,
help="Общее количество примеров (текстов) для генерации")
parser.add_argument("--languages", type=str, nargs='+', default=["en_US", "ru_RU", "zh_CN"],
help="Список локалей через пробел (например, en_US ru_RU zh_CN)")
parser.add_argument("--output_dir", type=str, default="./multilingual_dataset",
help="Директория для сохранения чанков")
parser.add_argument("--chunk_size", type=int, default=10000,
help="Количество примеров в одном файле-чанке")
parser.add_argument("--num_workers", type=int, default=4,
help="Количество параллельных процессов для генерации")
args = parser.parse_args()
os.makedirs(args.output_dir, exist_ok=True)
num_chunks = (args.total_examples + args.chunk_size - 1) // args.chunk_size
chunks = []
for i in range(num_chunks):
start = i * args.chunk_size
end = min(start + args.chunk_size, args.total_examples)
num_in_chunk = end - start
chunks.append((i, num_in_chunk, args.languages, args.output_dir, args.chunk_size))
pool = mp.Pool(processes=args.num_workers)
results = []
for chunk in chunks:
results.append(pool.apply_async(generate_chunk, chunk))
pool.close()
pool.join()
for r in results:
r.get()
print(f"Генерация завершена. Чанки сохранены в {args.output_dir}")
if __name__ == "main":
main()