File size: 3,526 Bytes
b233cf7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 | from __future__ import annotations
import argparse
import json
import logging
from pathlib import Path
from typing import Any
import pyarrow as pa
import pyarrow.parquet as pq
from .dataset_builder import DatasetBuilder
logger = logging.getLogger("pino.factory")
GENRES = ["citrus_cologne", "fougere", "floral_woody", "amber_oriental", "wildcard"]
def _record_to_parquet_row(record: dict[str, Any]) -> dict[str, Any]:
"""Flatten a record into Parquet-compatible columns."""
metadata = record.get("metadata", {})
return {
"formula_id": record.get("formula_id", ""),
"generation_strategy": metadata.get("generation_strategy", ""),
"status": record.get("status", ""),
"active_components_count": metadata.get("active_components_count", 0),
"formula": json.dumps(record.get("formula", [])),
"depletion_rates": json.dumps(record.get("depletion_rates", {})),
"trajectory": json.dumps(record.get("trajectory", [])),
}
def write_jsonl_to_parquet(jsonl_path: Path, parquet_path: Path) -> int:
"""Convert an existing JSON-L dataset to Parquet."""
rows = []
with jsonl_path.open("r") as f:
for line in f:
try:
record = json.loads(line)
rows.append(_record_to_parquet_row(record))
except Exception:
continue
schema = pa.schema([
pa.field("formula_id", pa.string()),
pa.field("generation_strategy", pa.string()),
pa.field("status", pa.string()),
pa.field("active_components_count", pa.int64()),
pa.field("formula", pa.string()),
pa.field("depletion_rates", pa.string()),
pa.field("trajectory", pa.string()),
])
arrays = {k: [r[k] for r in rows] for k in rows[0].keys()}
table = pa.table(arrays, schema=schema)
pq.write_table(table, parquet_path)
return len(rows)
def main() -> None:
parser = argparse.ArgumentParser(description="PINO Synthetic Data Factory")
parser.add_argument("--output", default="data/trajectories.parquet", help="Output Parquet path")
parser.add_argument("--count", type=int, default=5000, help="Total number of records")
parser.add_argument("--rules", default="data/genre_rules.json")
parser.add_argument("--registry", default="src/pino/registry.db")
parser.add_argument("--workers", type=int, default=16)
parser.add_argument("--seed", type=int, default=2026)
parser.add_argument("--from-jsonl", default="data/synthetic_dataset_v2.jsonl", help="Convert existing JSON-L instead of regenerating")
parser.add_argument("--regenerate", action="store_true", help="Regenerate from scratch")
args = parser.parse_args()
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s")
output_path = Path(args.output)
output_path.parent.mkdir(parents=True, exist_ok=True)
jsonl_path = Path(args.from_jsonl)
if args.regenerate or not jsonl_path.exists() or jsonl_path.stat().st_size == 0:
per_genre = args.count // len(GENRES)
builder = DatasetBuilder(
output_path=jsonl_path,
rules_path=args.rules,
registry_path=args.registry,
per_genre=per_genre,
max_workers=args.workers,
seed=args.seed,
)
builder.build()
count = write_jsonl_to_parquet(jsonl_path, output_path)
logger.info("Wrote %d records to %s", count, output_path)
if __name__ == "__main__":
main()
|