| 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() |
|
|