padoc/data.py

"""Portable PaDoc tree records and JSON/JSONL data loading."""

from __future__ import annotations

import json
from collections.abc import Iterator
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any


@dataclass
class Node:
    prefix: str
    children: list[Node] = field(default_factory=list)


@dataclass
class DataPiece:
    query: str
    response: Node
    images: list[str] = field(default_factory=list)


@dataclass(frozen=True)
class TrainingRecord:
    """One raw record plus the directory used for relative image paths."""

    value: dict[str, Any]
    image_dir: Path


@dataclass(frozen=True)
class DataConfig:
    records: list[TrainingRecord]
    seed: int = 42


def parse_node(value: dict[str, Any]) -> Node:
    if not isinstance(value, dict) or not isinstance(value.get("prefix"), str):
        raise ValueError("Every response node requires a string 'prefix'.")
    children = value.get("children", [])
    if not isinstance(children, list):
        raise ValueError("Node 'children' must be a list.")
    return Node(prefix=value["prefix"], children=[parse_node(child) for child in children])


def parse_data_piece(value: dict[str, Any]) -> DataPiece:
    if not isinstance(value, dict) or not isinstance(value.get("query"), str):
        raise ValueError("Every training record requires a string 'query'.")
    images = value.get("images", [])
    if not isinstance(images, list) or not all(isinstance(item, str) for item in images):
        raise ValueError("Record 'images' must be a list of paths.")
    return DataPiece(
        query=value["query"],
        response=parse_node(value["response"]),
        images=list(images),
    )


def _iter_file(path: Path) -> Iterator[dict[str, Any]]:
    if path.suffix.lower() == ".jsonl":
        with path.open(encoding="utf-8") as handle:
            for line_number, line in enumerate(handle, 1):
                if not line.strip():
                    continue
                try:
                    value = json.loads(line)
                except json.JSONDecodeError as exc:
                    raise ValueError(f"{path}:{line_number}: invalid JSON: {exc}") from exc
                if not isinstance(value, dict):
                    raise ValueError(f"{path}:{line_number}: expected a JSON object")
                yield value
        return

    with path.open(encoding="utf-8") as handle:
        values = json.load(handle)
    if not isinstance(values, list):
        raise ValueError(f"{path}: expected a top-level JSON array")
    for index, value in enumerate(values):
        if not isinstance(value, dict):
            raise ValueError(f"{path}: record {index} is not an object")
        yield value


def load_records(path: str | Path) -> list[dict[str, Any]]:
    """Load records from one JSON/JSONL file or a directory of such files."""

    path = Path(path).expanduser().resolve()
    if path.is_file():
        return list(_iter_file(path))
    if not path.is_dir():
        raise FileNotFoundError(f"Data path does not exist: {path}")
    files = sorted([*path.glob("*.json"), *path.glob("*.jsonl")])
    if not files:
        raise FileNotFoundError(f"No JSON or JSONL files found under {path}")
    records: list[dict[str, Any]] = []
    for file_path in files:
        records.extend(_iter_file(file_path))
    return records


def records_from_path(
    data_path: str | Path,
    *,
    image_dir: str | Path = ".",
) -> list[TrainingRecord]:
    root = Path(image_dir).expanduser().resolve()
    return [TrainingRecord(value, root) for value in load_records(data_path)]


def load_data_config(path: str | Path) -> DataConfig:
    """Load a portable YAML mix whose paths are relative to the YAML file.

    Schema::

        seed: 42
        sources:
          - path: ../examples/train.jsonl
            image_dir: ../examples/images
            repeat: 1
    """

    try:
        import yaml
    except ImportError as exc:  # pragma: no cover - declared dependency
        raise RuntimeError("PyYAML is required for --data-config") from exc

    config_path = Path(path).expanduser().resolve()
    with config_path.open(encoding="utf-8") as handle:
        raw = yaml.safe_load(handle)
    if not isinstance(raw, dict):
        raise ValueError(f"{config_path}: top-level YAML must be a mapping")
    unknown = set(raw) - {"seed", "sources"}
    if unknown:
        raise ValueError(f"{config_path}: unknown keys: {sorted(unknown)}")
    sources = raw.get("sources")
    if not isinstance(sources, list) or not sources:
        raise ValueError(f"{config_path}: 'sources' must be a non-empty list")

    records: list[TrainingRecord] = []
    for index, source in enumerate(sources):
        if not isinstance(source, dict):
            raise ValueError(f"{config_path}: sources[{index}] must be a mapping")
        source_unknown = set(source) - {"path", "image_dir", "repeat"}
        if source_unknown:
            raise ValueError(
                f"{config_path}: sources[{index}] unknown keys: {sorted(source_unknown)}"
            )
        if not isinstance(source.get("path"), str):
            raise ValueError(f"{config_path}: sources[{index}] requires string 'path'")
        repeat = source.get("repeat", 1)
        if not isinstance(repeat, int) or repeat < 1:
            raise ValueError(f"{config_path}: sources[{index}].repeat must be >= 1")

        source_path = (config_path.parent / source["path"]).resolve()
        image_value = source.get("image_dir")
        image_dir = (
            (config_path.parent / image_value).resolve()
            if isinstance(image_value, str)
            else (source_path.parent if source_path.is_file() else source_path)
        )
        loaded = [TrainingRecord(value, image_dir) for value in load_records(source_path)]
        for _ in range(repeat):
            records.extend(loaded)

    return DataConfig(records=records, seed=int(raw.get("seed", 42)))