wehub-resource-sync
148 lines
4.6 KiB
Python
148 lines
4.6 KiB
Python
#!/usr/bin/env python3
|
|
"""Generate fake (query, chunk_path) pairs for training pipeline validation.
|
|
|
|
Randomly samples chunks from kiwix_tiles, uses article titles to create
|
|
synthetic queries, and splits into train/eval JSONL files.
|
|
|
|
Usage:
|
|
uv run python training/fake_data.py \
|
|
--tiles-dir /opt/dlami/nvme/kiwix_tiles \
|
|
--articles-json /opt/dlami/nvme/kiwix/wikipedia_en_all_maxi_2025-08.zim.articles.json \
|
|
--output-dir training/data \
|
|
--num-articles 1000
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import logging
|
|
import random
|
|
import re
|
|
from pathlib import Path
|
|
|
|
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
|
|
logger = logging.getLogger(__name__)
|
|
|
|
QUERY_TEMPLATES = [
|
|
"What is {title}?",
|
|
"Tell me about {title}",
|
|
"{title} overview",
|
|
]
|
|
|
|
|
|
def find_chunk_paths(tiles_dir: Path, article_id: int) -> list[str]:
|
|
"""Find all chunk PNG paths for a given article ID."""
|
|
shard_size = 8284
|
|
top_shard = article_id // shard_size
|
|
top_dir = tiles_dir / f"shard_{top_shard:03d}"
|
|
if not top_dir.exists():
|
|
return []
|
|
tile_dir_name = f"{article_id}.png.tiles"
|
|
# Search sub-shards
|
|
for sub in top_dir.iterdir():
|
|
if not sub.is_dir() or not sub.name.startswith("shard_"):
|
|
continue
|
|
candidate = sub / tile_dir_name
|
|
if candidate.exists():
|
|
chunks_json = candidate / "chunks.json"
|
|
if chunks_json.exists():
|
|
try:
|
|
chunks = json.loads(chunks_json.read_text())
|
|
return [
|
|
str(candidate / c["file"])
|
|
for c in chunks.get("chunks", [])
|
|
if (candidate / c["file"]).exists()
|
|
]
|
|
except (json.JSONDecodeError, KeyError):
|
|
pass
|
|
return []
|
|
|
|
|
|
def title_from_slug(slug: str) -> str:
|
|
"""Convert URL slug to readable title."""
|
|
title = slug.replace("_", " ")
|
|
# Remove URL encoding
|
|
title = re.sub(r"%[0-9A-Fa-f]{2}", " ", title)
|
|
return title.strip()
|
|
|
|
|
|
def generate_queries(title: str) -> list[str]:
|
|
"""Generate fake queries from article title."""
|
|
return [t.format(title=title) for t in QUERY_TEMPLATES]
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Generate fake training data")
|
|
parser.add_argument(
|
|
"--tiles-dir", type=Path, default=Path("/opt/dlami/nvme/kiwix_tiles")
|
|
)
|
|
parser.add_argument(
|
|
"--articles-json",
|
|
type=Path,
|
|
default=Path(
|
|
"/opt/dlami/nvme/kiwix/wikipedia_en_all_maxi_2025-08.zim.articles.json"
|
|
),
|
|
)
|
|
parser.add_argument("--output-dir", type=Path, default=Path("training/data"))
|
|
parser.add_argument("--num-articles", type=int, default=1000)
|
|
parser.add_argument("--seed", type=int, default=42)
|
|
args = parser.parse_args()
|
|
|
|
random.seed(args.seed)
|
|
args.output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
logger.info("Loading articles.json...")
|
|
articles = json.loads(args.articles_json.read_text())
|
|
num_articles = len(articles)
|
|
logger.info(f"Loaded {num_articles} article slugs")
|
|
|
|
# Sample random article IDs and find ones with chunks
|
|
pairs = []
|
|
sampled = 0
|
|
indices = list(range(num_articles))
|
|
random.shuffle(indices)
|
|
|
|
for aid in indices:
|
|
if len(pairs) >= args.num_articles * 3: # 3 queries per article
|
|
break
|
|
sampled += 1
|
|
slug = articles[aid]
|
|
if not slug or slug.startswith("_"):
|
|
continue
|
|
chunk_paths = find_chunk_paths(args.tiles_dir, aid)
|
|
if not chunk_paths:
|
|
continue
|
|
title = title_from_slug(slug)
|
|
queries = generate_queries(title)
|
|
for query in queries:
|
|
# Pick a random chunk for this query
|
|
chunk_path = random.choice(chunk_paths)
|
|
pairs.append({"query": query, "chunk_path": chunk_path})
|
|
|
|
if len(pairs) % 300 == 0:
|
|
logger.info(
|
|
f"Generated {len(pairs)} pairs from {sampled} sampled articles..."
|
|
)
|
|
|
|
random.shuffle(pairs)
|
|
logger.info(f"Total pairs: {len(pairs)} from {sampled} sampled articles")
|
|
|
|
# Split 80/20
|
|
split = int(len(pairs) * 0.8)
|
|
train_pairs = pairs[:split]
|
|
eval_pairs = pairs[split:]
|
|
|
|
train_path = args.output_dir / "train.jsonl"
|
|
eval_path = args.output_dir / "eval.jsonl"
|
|
|
|
for path, data in [(train_path, train_pairs), (eval_path, eval_pairs)]:
|
|
with open(path, "w") as f:
|
|
for item in data:
|
|
f.write(json.dumps(item) + "\n")
|
|
logger.info(f"Wrote {len(data)} pairs to {path}")
|
|
|
|
logger.info("Done!")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|