diff --git a/.github/workflows/data-snapshot.yml b/.github/workflows/data-snapshot.yml new file mode 100644 index 00000000..96952a5f --- /dev/null +++ b/.github/workflows/data-snapshot.yml @@ -0,0 +1,43 @@ +name: Data Snapshot + +on: + push: + branches: [main] + paths: + - "src/ai4data/**" + - "tests/data_snapshot/**" + - "tests/conftest.py" + - "pyproject.toml" + - ".github/workflows/data-snapshot.yml" + pull_request: + paths: + - "src/ai4data/**" + - "tests/data_snapshot/**" + - "tests/conftest.py" + - "pyproject.toml" + - ".github/workflows/data-snapshot.yml" + +jobs: + test: + runs-on: ubuntu-latest + strategy: + matrix: + python-version: ["3.11", "3.12"] + steps: + - uses: actions/checkout@v4 + with: + fetch-depth: 0 + - uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + - name: Build wheel and source distribution + run: | + pip install build + python -m build + - name: Install wheel with test dependencies + run: | + for wheel in dist/*.whl; do + pip install "${wheel}[data-snapshot-metadata]" pytest + done + - name: Test offline from the installed package + run: python -m pytest tests/data_snapshot -q diff --git a/.gitignore b/.gitignore index 29e9ae7e..9fb02e15 100644 --- a/.gitignore +++ b/.gitignore @@ -2,6 +2,7 @@ _build/ sample.doc.json *.DS_Store data/ +/output/ # Byte-compiled / optimized / DLL files __pycache__/ @@ -127,6 +128,7 @@ celerybeat.pid # Environments .env .venv +/.venv-tfid/ env/ venv/ ENV/ diff --git a/README.md b/README.md index f02c298c..c1ebe461 100644 --- a/README.md +++ b/README.md @@ -211,6 +211,10 @@ Each capability is an optional extra. | `metadata` | Schema mapping, PDF-to-metadata extraction, and metadata augmentation | | `discovery` | NADA catalog access, metadata templates, and PDF helpers | | `metadata-reviewer` | The five-agent review pipeline, combined with one provider extra: `openai`, `anthropic`, `azure`, or `ollama` | +| `data-snapshot` | Figure and table images from PDFs | +| `data-snapshot-metadata` | Validated Data Snapshot Metadata Schema v1.4 metadata from snapshot images through OpenRouter or OpenAI | + +For the Data Snapshot pipeline, combine `data-snapshot` and `data-snapshot-metadata`. These extras are absent from `ai4data==0.1.0`; until a release includes them, follow the [source installation instructions](src/ai4data/data_snapshot/README.md#installation). In this checkout, `all` includes DocLayout-YOLO and snapshot metadata extraction. TF-ID is installed separately; the [usage guide](src/ai4data/data_snapshot/README.md) describes its dependency requirements and the three commands.
Example: review a metadata record @@ -330,7 +334,7 @@ Synthetic tabular and relational data generation is in a separate repository, [R ``` ai4data/ -├── src/ai4data/ Python package: anomaly, data_use, discovery, metadata (augmentation, reviewer) +├── src/ai4data/ Python package: anomaly, data_use, data_snapshot, discovery, metadata (augmentation, reviewer) ├── packages/ai4data/ JavaScript packages: @ai4data/core, @ai4data/search ├── reference/end_to_end/ Reference implementation with committed sample runs ├── cookbook/ Cookbook sources and their scripts diff --git a/experimental/data-snapshot/README.md b/experimental/data-snapshot/README.md index 056165ee..fd053250 100755 --- a/experimental/data-snapshot/README.md +++ b/experimental/data-snapshot/README.md @@ -1,5 +1,9 @@ # data-snapshot +This is the research and evaluation package. For the practical PDF-to-snapshot +and metadata workflow, use [`ai4data.data_snapshot`](../../src/ai4data/data_snapshot/README.md). +Keep this package's Python 3.10 environment separate from the production module. + Data snapshot toolkit for PDF documents. This package provides the source code for [Benchmarking Open-Source Layout Detection Models for Data Snapshot Extraction from Institutional Documents](https://arxiv.org/abs/2606.06242). diff --git a/pyproject.toml b/pyproject.toml index 089e4590..9800dcd5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -153,7 +153,43 @@ ollama = [ "autogen-ext[ollama]>=0.7.5", ] -# All capabilities +# Shared dependencies for metadata-only and alternative extractor installs. +data-snapshot-core = [ + "pillow>=11,<13", + "pydantic>=2.13.5,<3", + "pymupdf>=1.25,<2", + "python-dotenv>=1,<2", + "tqdm>=4.67,<5", +] +data-snapshot-metadata = [ + "ai4data[data-snapshot-core]", + "openai>=2.53,<3", +] +# The default Data Snapshot extra bundles DocLayout-YOLO. +data-snapshot = [ + "ai4data[data-snapshot-core]", + "doclayout-yolo==0.0.4", + "huggingface-hub>=0.23,<2", +] +data-snapshot-yolo11 = [ + "ai4data[data-snapshot-core]", + "huggingface-hub>=0.23,<2", + "ultralytics>=8.4,<9", +] +data-snapshot-yolo26 = [ + "ai4data[data-snapshot-yolo11]", +] +# TF-ID uses an older Transformers stack; install separately (see module README). +data-snapshot-tfid = [ + "ai4data[data-snapshot-core]", + "accelerate>=1,<2", + "einops>=0.8,<1", + "huggingface-hub>=0.23.2,<1", + "timm>=1,<2", + "transformers==4.46.0", +] + +# All capabilities with the default snapshot extractor; TF-ID stays opt-in. all = [ "ai4data[datause]", "ai4data[harmonization]", @@ -162,6 +198,7 @@ all = [ "ai4data[metadata]", "ai4data[discovery]", "ai4data[metadata-reviewer]", + "ai4data[data-snapshot,data-snapshot-metadata]", ] # Documentation (Docusaurus site in website/; notebooks are converted to Markdown diff --git a/src/ai4data/README.md b/src/ai4data/README.md index 306b9018..385c68d5 100644 --- a/src/ai4data/README.md +++ b/src/ai4data/README.md @@ -26,6 +26,15 @@ Each capability is an optional extra. | `metadata` | Schema mapping, PDF-to-metadata extraction, and metadata augmentation | | `discovery` | NADA catalog access, metadata templates, and PDF helpers | | `metadata-reviewer` | The five-agent metadata review pipeline, combined with one provider extra: `openai`, `anthropic`, `azure`, or `ollama` | +| `data-snapshot` | Figure and table images from PDFs | +| `data-snapshot-metadata` | Validated Schema v1.4 metadata from snapshot images through OpenRouter or OpenAI | + +For the Data Snapshot pipeline, combine `data-snapshot` and +`data-snapshot-metadata`. These extras are absent from `ai4data==0.1.0`; until a +release includes them, install the checkout as described in the +[Data Snapshot usage guide](https://github.com/worldbank/ai4data/blob/main/src/ai4data/data_snapshot/README.md). +In this checkout, `all` includes DocLayout-YOLO and snapshot metadata extraction. +TF-ID requires a separate environment; the guide lists its installation steps. ## Examples @@ -53,6 +62,9 @@ issues = job.wait_sync(timeout=300) The documentation site describes each module: monitoring of data use, anomaly detection and explanation, data discoverability, metadata augmentation, and the Metadata Reviewer user manual. +Data Snapshot documentation currently covers repository installation and usage; +concepts, methods, worked examples, and user guides for the published site are +planned in a separate PR. ## License diff --git a/src/ai4data/__init__.py b/src/ai4data/__init__.py index 1575125a..a2874725 100644 --- a/src/ai4data/__init__.py +++ b/src/ai4data/__init__.py @@ -14,6 +14,9 @@ - discovery: NADA catalog, metadata templates, and related discovery helpers Install with: uv pip install ai4data[discovery] +- data_snapshot: Figure/table snapshots and Schema v1.4 metadata from PDFs + Install with: python -m pip install "ai4data[data-snapshot,data-snapshot-metadata]" + For all capabilities: uv pip install ai4data[all] """ @@ -37,6 +40,7 @@ from . import anomaly from . import discovery from . import metadata +from . import data_snapshot __all__ = [ "__version__", @@ -44,4 +48,5 @@ "anomaly", "discovery", "metadata", + "data_snapshot", ] diff --git a/src/ai4data/data_snapshot/README.md b/src/ai4data/data_snapshot/README.md new file mode 100644 index 00000000..209663ea --- /dev/null +++ b/src/ai4data/data_snapshot/README.md @@ -0,0 +1,236 @@ +# Data Snapshot + +Extract figure and table images from PDFs, then extract validated Schema v1.4 +metadata from those images. DocLayout-YOLO is the default extractor. Value +transcription and submission to external services are not implemented. + +## Installation + +Requires Python 3.11 or newer. Use your preferred Python environment and package +manager. Data Snapshot is part of `ai4data`, with optional extractor and metadata +dependencies. + +**Release availability:** `ai4data==0.1.0` does not include these extras. Once a +release containing this module is published, the standard installation is: + +```bash +python -m pip install "ai4data[data-snapshot,data-snapshot-metadata]" +``` + +Until then, install the checkout from the repository root: + +```bash +python -m pip install ".[data-snapshot,data-snapshot-metadata]" +``` + +This builds and installs the package normally. Reinstall after changing the +source. Contributors who want edits to take effect immediately can use the +editable installation described under [Development](#development). + +The `data-snapshot` extra includes DocLayout-YOLO. For snapshot extraction only, +omit `data-snapshot-metadata`. For metadata from existing images, install only +`data-snapshot-metadata`; it supports both OpenRouter and OpenAI without a model +runtime or the reviewer's provider extras. + +Install only the extractor you need. Weights download on the first extractor +run and are cached under `.cache/models/`, relative to the working directory. + +| Extractor | Extra | Weights | Inference size | +| --- | --- | --- | ---: | +| `doclayoutyolo` (default) | `data-snapshot` | `juliozhao/DocLayout-YOLO-DocStructBench` | 1024 | +| `yolo11` | `data-snapshot-yolo11` | `Armaggheddon/yolo11-document-layout`, `yolo11m_doc_layout.pt` | 1280 | +| `yolo26` | `data-snapshot-yolo26` | `Armaggheddon/yolo26-document-layout`, `yolo26m_doc_layout.pt` | 1280 | +| `tfid` | `data-snapshot-tfid` | `yifeihu/TF-ID-large` | Model-defined | + +The shared `data-snapshot-core` extra supplies rendering and schema dependencies +without a model runtime or metadata API SDK. It is included by every snapshot +extra. `ai4data[all]` includes DocLayout-YOLO and snapshot metadata extraction; +it excludes TF-ID. + +### Model configuration + +Default model locations are stored in the packaged +[`config/models.json`](config/models.json), rather than in the adapters' Python +code. No local configuration file is required to use these defaults. + +Override model locations with environment variables or a `.env` file in the +working directory. Both PDF commands load that file before creating a detector; +exported environment variables take precedence over `.env`. When running from +the repository root, put `.env` alongside `pyproject.toml`. It is Git-ignored. + +Each row below lists the supported variable names and their defaults: + +| Extractor | Environment variable | Default | +| --- | --- | --- | +| DocLayout-YOLO | `AI4DATA_SNAPSHOT_DOCLAYOUTYOLO_REPO_ID` | `juliozhao/DocLayout-YOLO-DocStructBench` | +| | `AI4DATA_SNAPSHOT_DOCLAYOUTYOLO_FILENAME` | `doclayout_yolo_docstructbench_imgsz1024.pt` | +| | `AI4DATA_SNAPSHOT_DOCLAYOUTYOLO_CACHE_DIR` | `.cache/models` | +| YOLO11 | `AI4DATA_SNAPSHOT_YOLO11_REPO_ID` | `Armaggheddon/yolo11-document-layout` | +| | `AI4DATA_SNAPSHOT_YOLO11_FILENAME` | `yolo11m_doc_layout.pt` | +| | `AI4DATA_SNAPSHOT_YOLO11_CACHE_DIR` | `.cache/models` | +| YOLO26 | `AI4DATA_SNAPSHOT_YOLO26_REPO_ID` | `Armaggheddon/yolo26-document-layout` | +| | `AI4DATA_SNAPSHOT_YOLO26_FILENAME` | `yolo26m_doc_layout.pt` | +| | `AI4DATA_SNAPSHOT_YOLO26_CACHE_DIR` | `.cache/models` | +| TF-ID | `AI4DATA_SNAPSHOT_TFID_MODEL_ID` | `yifeihu/TF-ID-large` | +| | `AI4DATA_SNAPSHOT_TFID_REVISION` | `a91ffef47039768b2deff82681b75d77252703f2` | +| | `AI4DATA_SNAPSHOT_TFID_CACHE_DIR` | `.cache/models/tfid` | + +For example, save DocLayout-YOLO weights in a shared cache: + +```dotenv +AI4DATA_SNAPSHOT_DOCLAYOUTYOLO_CACHE_DIR=/path/to/model-cache +``` + +Unset or empty variables use the packaged JSON defaults. Environment overrides +are read whenever a configuration is constructed, rather than at module import. +Precedence is explicit Python arguments, then environment variables, then JSON +defaults. For a local customization, use an environment override rather than +editing the installed JSON, which can be replaced during package upgrades: + +```python +from ai4data.data_snapshot.adapters.doclayoutyolo import DocLayoutYOLOConfig + +config = DocLayoutYOLOConfig(cache_dir="/path/to/model-cache") +``` + +Python callers using `.env` should call `dotenv.load_dotenv` with the file path +before constructing a configuration or running extraction; automatic loading +of model settings from `.env` is provided by the PDF commands. + +### Metadata API setup + +Required only for end-to-end and metadata-only runs. Set the appropriate key +in the environment or a `.env` file in the working directory. When running +commands from the repository root, put `.env` in that root, alongside +`pyproject.toml`. It is Git-ignored: + +```dotenv +OPENROUTER_API_KEY=your-key +OPENAI_API_KEY=your-key +``` + +Only the selected provider's key is required. Existing environment variables +take precedence over `.env`. OpenRouter is the default; select OpenAI with +`--metadata-provider=openai`. Both packaged presets use GPT-5.6 Luna through the +Responses API. Requests consume provider credits; automatic SDK retries are +disabled. Snapshot extraction alone makes no metadata API requests. + +### TF-ID installation + +**Treat TF-ID as incompatible with the rest of the repository's model stack.** +It pins `transformers==4.46.0` and `huggingface-hub<1`; installing it alongside +other capabilities can downgrade shared libraries and break those workflows. +Install it in its own environment, without `all` or other capability extras: + +```bash +python -m venv .venv-tfid +.venv-tfid/bin/python -m pip install ".[data-snapshot-tfid]" +.venv-tfid/bin/python -m ai4data.data_snapshot.extract_data_snapshot \ + --model=tfid --input=data/example.pdf +``` + +For end-to-end TF-ID runs, also install `data-snapshot-metadata` in that same +environment. TF-ID executes model-provided code at the configured revision, +defaulting to the pinned commit `a91ffef47039768b2deff82681b75d77252703f2`. +Changing `AI4DATA_SNAPSHOT_TFID_MODEL_ID` or `AI4DATA_SNAPSHOT_TFID_REVISION` +changes which remote model code is executed. + +## Commands + +After installation, use Python from that environment. Paths are relative to +the working directory; inputs can also be absolute paths. Add `--help` to any +command for options. + +### End-to-end + +Extract snapshots, then metadata: + +```bash +python -m ai4data.data_snapshot --input=data/example.pdf +``` + +### Snapshot-only + +```bash +python -m ai4data.data_snapshot.extract_data_snapshot --input=data/example.pdf +``` + +Both PDF commands accept a PDF or a directory searched recursively for PDFs. +Use `--model=yolo11`, `--model=yolo26`, or `--model=tfid` to select another +installed extractor. Exactly one extractor runs per invocation. + +### Metadata-only + +Fill in missing metadata without rerunning the PDF extractor: + +```bash +python -m ai4data.data_snapshot.extract_metadata \ + --input=output/doclayoutyolo/snapshots +``` + +Accepts a PNG, JPEG, GIF, or WebP image, or a directory searched recursively for +images. A `snapshots` input directory, or an image directly inside one, defaults +to its sibling `metadata` directory. For other locations, supply `--output-dir`: + +```bash +python -m ai4data.data_snapshot.extract_metadata \ + --input=path/to/snapshot.png --output-dir=path/to/metadata \ + --metadata-provider=openai +``` + +### Outputs and reruns + +PDF commands default to `output//`, with `snapshots/` and `metadata/` +subdirectories. `--output-dir` changes the root. For `example.pdf`: + +```text +output/doclayoutyolo/ +├── snapshots/example_figure_000.png +├── snapshots/example_table_000.png +├── metadata/example_figure_000.json +└── metadata/example_table_000.json +``` + +Counters start at zero and are independent for figures and tables. PDF stems +must be unique within a run, including across subdirectories. Metadata is +validated against the canonical Pydantic Schema v1.4 model. Requests use only +image-visible evidence, excluding filenames and source-document metadata. + +| Command | Default | With `--force` / `-f` | +| --- | --- | --- | +| End-to-end | Skip a PDF if any matching snapshot or metadata exists, including partial output. | Delete matching snapshots and metadata; rerun both stages. | +| Snapshot-only | Same PDF-level skip behavior. | Delete matching snapshots and metadata; extract snapshots only. | +| Metadata-only | Skip images with an existing JSON output. | Delete and regenerate JSON; preserve images. | + +**Forced replacement deletes existing files before extraction starts. Failure +does not restore them.** Use metadata-only to fill missing JSON after a failed +metadata stage. Fully skipped PDF runs do not load the extractor. + +## Python API + +```python +from ai4data.data_snapshot import run_extraction, run_pipeline + +# DocLayout-YOLO is the default; snapshot-only requires no metadata API key. +snapshots = run_extraction("data/example.pdf", "output") + +# Use fresh input/output or the metadata-only command after snapshot-only runs: +paths = run_pipeline("data/another.pdf", "output", metadata_provider="openai") +``` + +## Development + +Install `data-snapshot-metadata` and `pytest` to run the offline tests; model +runtimes and weights are unnecessary. From the repository root, use an editable +installation so source changes take effect without reinstalling: + +```bash +python -m pip install -e ".[data-snapshot-metadata]" pytest +python -m pytest tests/data_snapshot -q +``` + +Ported from `ajdajd/data-snapshot-pipeline` at +`b12c27780a5082d4a96860de868da9af6b03c094`. This namespace replaces +`data_snapshot_pipeline`; no old import or command aliases are provided. +The [experimental package](../../../experimental/data-snapshot/README.md) +remains the separate research and evaluation implementation. diff --git a/src/ai4data/data_snapshot/__init__.py b/src/ai4data/data_snapshot/__init__.py new file mode 100644 index 00000000..5c28c3ed --- /dev/null +++ b/src/ai4data/data_snapshot/__init__.py @@ -0,0 +1,47 @@ +"""Extract figure/table snapshots and metadata with optional model dependencies.""" + +from importlib import import_module +from typing import Any + +_EXPORT_MODULES = { + "Detection": "detection", + "DataSnapshotMetadata": "metadata_schema", + "ExtractionResult": "metadata_extraction", + "SnapshotDetector": "detection", + "build_snapshot_stem": "naming", + "discover_pdfs": "inputs", + "discover_snapshot_images": "inputs", + "extract_snapshots": "extraction", + "extract_metadata": "metadata_extraction", + "iter_pdf_pages": "rendering", + "run_extraction": "pipeline", + "run_pipeline": "pipeline", + "submit_snapshot": "submission", +} + +__all__ = list(_EXPORT_MODULES) + + +def __getattr__(name: str) -> Any: + """Resolve and cache public objects without importing optional runtimes. + + Parameters + ---------- + name : str + Public attribute requested from the module. + + Returns + ------- + Any + The corresponding public function or class. + + Raises + ------ + AttributeError + If the name is not a public export. + """ + if name not in _EXPORT_MODULES: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + value = getattr(import_module(f"{__name__}.{_EXPORT_MODULES[name]}"), name) + globals()[name] = value + return value diff --git a/src/ai4data/data_snapshot/__main__.py b/src/ai4data/data_snapshot/__main__.py new file mode 100644 index 00000000..49890106 --- /dev/null +++ b/src/ai4data/data_snapshot/__main__.py @@ -0,0 +1,5 @@ +"""Run the pipeline with ``python -m ai4data.data_snapshot``.""" + +from ai4data.data_snapshot.cli import main + +raise SystemExit(main()) diff --git a/src/ai4data/data_snapshot/adapters/__init__.py b/src/ai4data/data_snapshot/adapters/__init__.py new file mode 100644 index 00000000..333cd411 --- /dev/null +++ b/src/ai4data/data_snapshot/adapters/__init__.py @@ -0,0 +1,43 @@ +"""Create snapshot detector adapters for data snapshot extraction.""" + +from ai4data.data_snapshot.detection import SnapshotDetector + + +def create_adapter(model_name: str) -> SnapshotDetector: + """Create the one snapshot detector selected for a pipeline run. + + Parameters + ---------- + model_name : str + Extractor identifier, currently ``"doclayoutyolo"``, ``"tfid"``, + ``"yolo11"``, or ``"yolo26"``. + + Returns + ------- + SnapshotDetector + Selected snapshot detector adapter. + + Raises + ------ + ValueError + If the model identifier is not implemented. + """ + if model_name == "doclayoutyolo": + from ai4data.data_snapshot.adapters.doclayoutyolo import ( + DocLayoutYOLOAdapter, + ) + + return DocLayoutYOLOAdapter() + if model_name == "tfid": + from ai4data.data_snapshot.adapters.tfid import TFIDAdapter + + return TFIDAdapter() + if model_name == "yolo11": + from ai4data.data_snapshot.adapters.yolo11 import YOLO11Adapter + + return YOLO11Adapter() + if model_name == "yolo26": + from ai4data.data_snapshot.adapters.yolo26 import YOLO26Adapter + + return YOLO26Adapter() + raise ValueError(f"Unsupported model: {model_name}") diff --git a/src/ai4data/data_snapshot/adapters/doclayoutyolo.py b/src/ai4data/data_snapshot/adapters/doclayoutyolo.py new file mode 100644 index 00000000..16643101 --- /dev/null +++ b/src/ai4data/data_snapshot/adapters/doclayoutyolo.py @@ -0,0 +1,80 @@ +"""Detect figure and table snapshots with DocLayout-YOLO.""" + +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +from ai4data.data_snapshot.adapters.yolo26 import YOLO26Adapter, YOLO26Config +from ai4data.data_snapshot.model_config import get_model_setting + + +@dataclass(frozen=True) +class DocLayoutYOLOConfig(YOLO26Config): + """Configure DocLayout-YOLO inference with official defaults. + + Notes + ----- + Inherits the YOLO26 configuration parameters. Omitted model location fields + use ``AI4DATA_SNAPSHOT_DOCLAYOUTYOLO_*`` environment variables at construction + time. Explicit arguments take precedence. + """ + + repo_id: str = field( + default_factory=lambda: get_model_setting("doclayoutyolo", "repo_id") + ) + filename: str = field( + default_factory=lambda: get_model_setting("doclayoutyolo", "filename") + ) + cache_dir: str | Path = field( + default_factory=lambda: Path(get_model_setting("doclayoutyolo", "cache_dir")) + ) + conf: float = 0.2 + imgsz: int = 1024 + + +class DocLayoutYOLOAdapter(YOLO26Adapter): + """Convert DocLayout-YOLO predictions into canonical detections. + + Parameters + ---------- + config : DocLayoutYOLOConfig | None + Inference configuration. Defaults are used when omitted. + model : Any | None + Injected DocLayout-YOLO-compatible model, primarily for offline tests. + """ + + def __init__( + self, + config: DocLayoutYOLOConfig | None = None, + *, + model: Any | None = None, + ) -> None: + resolved_config = config or DocLayoutYOLOConfig() + super().__init__( + resolved_config, + model=model if model is not None else _load_model(resolved_config), + ) + + +def _load_model(config: DocLayoutYOLOConfig) -> Any: + """Load DocLayout-YOLO from the local cache or Hugging Face.""" + from doclayout_yolo import YOLOv10 + + if config.model_path is not None: + model_path = Path(config.model_path) + if not model_path.is_file(): + raise FileNotFoundError(f"Model weights not found: {model_path}") + else: + model_path = Path(config.cache_dir) / config.filename + if not model_path.is_file(): + from huggingface_hub import hf_hub_download + + model_path = Path( + hf_hub_download( + repo_id=config.repo_id, + filename=config.filename, + repo_type="model", + local_dir=config.cache_dir, + ) + ) + return YOLOv10(str(model_path)) diff --git a/src/ai4data/data_snapshot/adapters/tfid.py b/src/ai4data/data_snapshot/adapters/tfid.py new file mode 100644 index 00000000..cda3e469 --- /dev/null +++ b/src/ai4data/data_snapshot/adapters/tfid.py @@ -0,0 +1,188 @@ +"""Detect figure and table snapshots with TF-ID Large.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +from PIL import Image + +from ai4data.data_snapshot.adapters.yolo26 import _clamp, _normalize_bbox +from ai4data.data_snapshot.model_config import get_model_setting +from ai4data.data_snapshot.detection import Detection + +_LABELS = { + "1": "figure", + "2": "table", + "chart": "figure", + "diagram": "figure", + "fig": "figure", + "figure": "figure", + "image": "figure", + "table": "table", + "tbl": "table", +} + + +@dataclass(frozen=True) +class TFIDConfig: + """Configure TF-ID Large inference. + + Parameters + ---------- + model_id : str + Hugging Face model identifier. + revision : str + Hugging Face model revision containing the remote model code. + Defaults to a pinned commit. + cache_dir : str | Path + Repository-local Hugging Face model cache. + device : str + Torch device string, such as ``"cpu"`` or ``"cuda:0"``. + min_area : float + Minimum normalized bounding-box area retained as a snapshot. + + Notes + ----- + Omitted model location fields are read from ``AI4DATA_SNAPSHOT_TFID_*`` + environment variables at construction time. Explicit arguments take precedence. + + Raises + ------ + ValueError + If the minimum area is outside the inclusive unit interval. + """ + + model_id: str = field(default_factory=lambda: get_model_setting("tfid", "model_id")) + revision: str = field(default_factory=lambda: get_model_setting("tfid", "revision")) + cache_dir: str | Path = field( + default_factory=lambda: Path(get_model_setting("tfid", "cache_dir")) + ) + device: str = "cpu" + min_area: float = 0.008 + + def __post_init__(self) -> None: + """Validate TF-ID inference settings.""" + if not 0.0 <= self.min_area <= 1.0: + raise ValueError("Minimum area must be between zero and one.") + + +class TFIDAdapter: + """Convert TF-ID Large predictions into canonical detections. + + Parameters + ---------- + config : TFIDConfig | None + Inference configuration. Defaults are used when omitted. + model : Any | None + Injected Transformers-compatible model for offline tests. + processor : Any | None + Injected Transformers-compatible processor for offline tests. + + Raises + ------ + ValueError + If only one of ``model`` and ``processor`` is supplied. + """ + + def __init__( + self, + config: TFIDConfig | None = None, + *, + model: Any | None = None, + processor: Any | None = None, + ) -> None: + self.config = config or TFIDConfig() + if (model is None) != (processor is None): + raise ValueError("Model and processor must be supplied together.") + if model is None: + model, processor = _load_components(self.config) + self._model = model + self._processor = processor + + def detect(self, image: Image.Image) -> list[Detection]: + """Detect figure and table snapshots in one rendered PDF page. + + Parameters + ---------- + image : PIL.Image.Image + RGB page image produced by the pipeline renderer. + + Returns + ------- + list[Detection] + Canonical detections in model output order. + """ + inputs = self._processor(text="", images=image, return_tensors="pt") + inputs = inputs.to(self._model.device) + generated_ids = self._model.generate( + input_ids=inputs["input_ids"], + pixel_values=inputs["pixel_values"], + max_new_tokens=1024, + do_sample=False, + num_beams=3, + ) + generated_text = self._processor.batch_decode( + generated_ids, skip_special_tokens=False + )[0] + result = self._processor.post_process_generation( + generated_text, + task="", + image_size=(image.width, image.height), + ).get("", {}) + boxes = result.get("bboxes", []) + labels = result.get("labels", result.get("classes", [])) + scores = result.get("scores") + if not isinstance(boxes, list) or not isinstance(labels, list): + return [] + + detections: list[Detection] = [] + for index, (box, raw_label) in enumerate(zip(boxes, labels)): + label = _coerce_label(raw_label) + bbox = _normalize_bbox(box, image.width, image.height) + if label is None or bbox is None: + continue + if (bbox[2] - bbox[0]) * (bbox[3] - bbox[1]) < self.config.min_area: + continue + score = _score_at(scores, index) + detections.append(Detection(label=label, bbox=bbox, score=score)) + return detections + + +def _load_components(config: TFIDConfig) -> tuple[Any, Any]: + """Load the pinned TF-ID model and processor from Hugging Face.""" + from transformers import AutoModelForCausalLM, AutoProcessor + + common = { + "cache_dir": config.cache_dir, + "revision": config.revision, + "trust_remote_code": True, + } + model = AutoModelForCausalLM.from_pretrained( + config.model_id, + device_map=config.device, + **common, + ) + model.eval() + processor = AutoProcessor.from_pretrained(config.model_id, **common) + return model, processor + + +def _coerce_label(raw_label: Any) -> str | None: + """Normalize one TF-ID label to a supported snapshot label.""" + if raw_label is None: + return None + if isinstance(raw_label, (int, float)): + raw_label = int(raw_label) + return _LABELS.get(str(raw_label).strip().lower()) + + +def _score_at(scores: Any, index: int) -> float: + """Return one optional TF-ID score, defaulting to full confidence.""" + if not isinstance(scores, list) or index >= len(scores): + return 1.0 + try: + return _clamp(float(scores[index])) + except (TypeError, ValueError): + return 1.0 diff --git a/src/ai4data/data_snapshot/adapters/yolo11.py b/src/ai4data/data_snapshot/adapters/yolo11.py new file mode 100644 index 00000000..7869268c --- /dev/null +++ b/src/ai4data/data_snapshot/adapters/yolo11.py @@ -0,0 +1,45 @@ +"""Detect figure and table snapshots with YOLO11 Document Layout.""" + +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +from ai4data.data_snapshot.adapters.yolo26 import YOLO26Adapter, YOLO26Config +from ai4data.data_snapshot.model_config import get_model_setting + + +@dataclass(frozen=True) +class YOLO11Config(YOLO26Config): + """Configure medium YOLO11 Document Layout inference. + + Notes + ----- + Inherits the YOLO26 configuration parameters. Omitted model location fields + use ``AI4DATA_SNAPSHOT_YOLO11_*`` environment variables at construction time. + Explicit arguments take precedence. + """ + + repo_id: str = field(default_factory=lambda: get_model_setting("yolo11", "repo_id")) + filename: str = field( + default_factory=lambda: get_model_setting("yolo11", "filename") + ) + cache_dir: str | Path = field( + default_factory=lambda: Path(get_model_setting("yolo11", "cache_dir")) + ) + + +class YOLO11Adapter(YOLO26Adapter): + """Convert YOLO11 model predictions into canonical detections. + + Parameters + ---------- + config : YOLO11Config | None + Inference configuration. Defaults are used when omitted. + model : Any | None + Injected Ultralytics-compatible model, primarily for offline tests. + """ + + def __init__( + self, config: YOLO11Config | None = None, *, model: Any | None = None + ) -> None: + super().__init__(config or YOLO11Config(), model=model) diff --git a/src/ai4data/data_snapshot/adapters/yolo26.py b/src/ai4data/data_snapshot/adapters/yolo26.py new file mode 100644 index 00000000..abead25c --- /dev/null +++ b/src/ai4data/data_snapshot/adapters/yolo26.py @@ -0,0 +1,202 @@ +"""Detect figure and table snapshots with YOLO26 Document Layout.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +from PIL import Image + +from ai4data.data_snapshot.detection import Detection +from ai4data.data_snapshot.model_config import get_model_setting + +_LABELS = { + "figure": "figure", + "picture": "figure", + "table": "table", +} + + +@dataclass(frozen=True) +class YOLO26Config: + """Configure YOLO26 Document Layout inference. + + Parameters + ---------- + model_path : str | Path | None + Optional local model weights. The configured Hugging Face model is + downloaded when omitted. + repo_id : str + Hugging Face model repository identifier. + filename : str + Model weights filename in the repository. + cache_dir : str | Path + Repository-local directory used to cache downloaded model weights. + device : str + Ultralytics device identifier, such as ``"cpu"`` or ``"cuda:0"``. + conf : float + Minimum confidence passed to YOLO prediction. + iou : float + Intersection-over-union threshold used for non-maximum suppression. + imgsz : int + YOLO inference image size. + min_area : float + Minimum normalized bounding-box area retained as a snapshot. + + Notes + ----- + Omitted model location fields are read from ``AI4DATA_SNAPSHOT_YOLO26_*`` + environment variables at construction time. Explicit arguments take precedence. + + Raises + ------ + ValueError + If a numeric configuration value is outside its supported range. + """ + + model_path: str | Path | None = None + repo_id: str = field(default_factory=lambda: get_model_setting("yolo26", "repo_id")) + filename: str = field( + default_factory=lambda: get_model_setting("yolo26", "filename") + ) + cache_dir: str | Path = field( + default_factory=lambda: Path(get_model_setting("yolo26", "cache_dir")) + ) + device: str = "cpu" + conf: float = 0.25 + iou: float = 0.7 + imgsz: int = 1280 + min_area: float = 0.008 + + def __post_init__(self) -> None: + """Validate inference settings before loading the model.""" + if not 0.0 <= self.conf <= 1.0: + raise ValueError("Confidence threshold must be between zero and one.") + if not 0.0 <= self.iou <= 1.0: + raise ValueError("IoU threshold must be between zero and one.") + if self.imgsz <= 0: + raise ValueError("Inference image size must be greater than zero.") + if not 0.0 <= self.min_area <= 1.0: + raise ValueError("Minimum area must be between zero and one.") + + +class YOLO26Adapter: + """Convert YOLO26 model predictions into canonical detections. + + Parameters + ---------- + config : YOLO26Config | None + Inference configuration. Defaults are used when omitted. + model : Any | None + Injected Ultralytics-compatible model, primarily for offline tests. + """ + + def __init__( + self, config: YOLO26Config | None = None, *, model: Any | None = None + ) -> None: + self.config = config or YOLO26Config() + self._model = model if model is not None else _load_model(self.config) + + def detect(self, image: Image.Image) -> list[Detection]: + """Detect figure and table snapshots in one rendered PDF page. + + Parameters + ---------- + image : PIL.Image.Image + RGB page image produced by the pipeline renderer. + + Returns + ------- + list[Detection] + Canonical detections in model output order. + """ + results = self._model.predict( + image, + imgsz=self.config.imgsz, + conf=self.config.conf, + iou=self.config.iou, + device=self.config.device, + verbose=False, + ) + if not results: + return [] + + result = results[0] + boxes = result.boxes.xyxy.cpu().tolist() + scores = result.boxes.conf.cpu().tolist() + classes = result.boxes.cls.cpu().tolist() + detections: list[Detection] = [] + + for box, score, class_index in zip(boxes, scores, classes): + label = _coerce_label(_class_name(result.names, int(class_index))) + bbox = _normalize_bbox(box, image.width, image.height) + if label is None or bbox is None: + continue + if (bbox[2] - bbox[0]) * (bbox[3] - bbox[1]) < self.config.min_area: + continue + detections.append( + Detection(label=label, bbox=bbox, score=_clamp(float(score))) + ) + + return detections + + +def _load_model(config: YOLO26Config) -> Any: + """Load YOLO26 from a local path or Hugging Face.""" + from ultralytics import YOLO + + if config.model_path is not None: + model_path = Path(config.model_path) + if not model_path.is_file(): + raise FileNotFoundError(f"Model weights not found: {model_path}") + else: + model_path = Path(config.cache_dir) / config.filename + if not model_path.is_file(): + from huggingface_hub import hf_hub_download + + model_path = Path( + hf_hub_download( + repo_id=config.repo_id, + filename=config.filename, + repo_type="model", + local_dir=config.cache_dir, + ) + ) + return YOLO(str(model_path)) + + +def _class_name(names: Mapping[int, str] | Sequence[str], index: int) -> str | None: + """Return a class name from Ultralytics mapping or sequence metadata.""" + if isinstance(names, Mapping): + return names.get(index) + if 0 <= index < len(names): + return names[index] + return None + + +def _coerce_label(raw_label: str | None) -> str | None: + """Normalize one YOLO class name to a supported snapshot label.""" + if raw_label is None: + return None + return _LABELS.get(raw_label.strip().lower()) + + +def _normalize_bbox( + box: Sequence[float], width: int, height: int +) -> tuple[float, float, float, float] | None: + """Normalize and clamp one pixel-space YOLO bounding box.""" + if len(box) != 4 or width <= 0 or height <= 0: + return None + x1, y1, x2, y2 = (float(value) for value in box) + nx1, nx2 = sorted((_clamp(x1 / width), _clamp(x2 / width))) + ny1, ny2 = sorted((_clamp(y1 / height), _clamp(y2 / height))) + if nx2 <= nx1 or ny2 <= ny1: + return None + return nx1, ny1, nx2, ny2 + + +def _clamp(value: float) -> float: + """Clamp a floating-point value to the inclusive unit interval.""" + return max(0.0, min(1.0, value)) diff --git a/src/ai4data/data_snapshot/cli.py b/src/ai4data/data_snapshot/cli.py new file mode 100644 index 00000000..5a40af34 --- /dev/null +++ b/src/ai4data/data_snapshot/cli.py @@ -0,0 +1,231 @@ +"""Provide command-line interfaces for each pipeline stage.""" + +import argparse +import sys +from pathlib import Path + +from dotenv import load_dotenv + +from ai4data.data_snapshot.inputs import discover_snapshot_images +from ai4data.data_snapshot.metadata_extraction import extract_metadata_files +from ai4data.data_snapshot.pipeline import run_extraction, run_pipeline + + +def _positive_int(value: str) -> int: + """Parse a positive integer command-line value.""" + parsed = int(value) + if parsed < 1: + raise argparse.ArgumentTypeError("must be a positive integer") + return parsed + + +def main(argv: list[str] | None = None) -> int: + """Run the end-to-end pipeline from command-line arguments. + + Parameters + ---------- + argv : list[str] | None, optional + Arguments to parse, excluding the executable name. Uses ``sys.argv`` + when omitted. + + Returns + ------- + int + Zero when the pipeline completes successfully, otherwise one. + + Raises + ------ + SystemExit + If arguments or local inputs are invalid. + """ + parser = _snapshot_parser( + "python -m ai4data.data_snapshot", + "Extract data snapshots and Schema v1.4 metadata from PDFs.", + ) + _add_metadata_provider_argument(parser) + return _run_snapshot_command(parser, argv, end_to_end=True) + + +def snapshot_main(argv: list[str] | None = None) -> int: + """Run data snapshot extraction from command-line arguments. + + Parameters + ---------- + argv : list[str] | None, optional + Arguments to parse, excluding the executable name. Uses ``sys.argv`` + when omitted. + + Returns + ------- + int + Zero when extraction completes successfully, otherwise one. + + Raises + ------ + SystemExit + If arguments or local inputs are invalid. + """ + parser = _snapshot_parser( + "python -m ai4data.data_snapshot.extract_data_snapshot", + "Extract data snapshots from PDFs without metadata extraction.", + ) + return _run_snapshot_command(parser, argv, end_to_end=False) + + +def metadata_main(argv: list[str] | None = None) -> int: + """Run metadata extraction from command-line arguments. + + Parameters + ---------- + argv : list[str] | None, optional + Arguments to parse, excluding the executable name. Uses ``sys.argv`` + when omitted. + + Returns + ------- + int + Zero when extraction completes successfully, otherwise one. + + Raises + ------ + SystemExit + If arguments or local inputs are invalid. + """ + parser = argparse.ArgumentParser( + prog="python -m ai4data.data_snapshot.extract_metadata", + description="Extract Schema v1.4 metadata from snapshot images.", + ) + parser.add_argument( + "--input", + required=True, + type=Path, + help="A snapshot image or directory containing snapshot images.", + ) + parser.add_argument( + "--output-dir", + type=Path, + help=( + "Metadata output directory. By default, infer the metadata sibling " + "of an input snapshots directory." + ), + ) + parser.add_argument( + "-f", + "--force", + action="store_true", + help="Replace existing metadata without modifying snapshots.", + ) + _add_metadata_provider_argument(parser) + args = parser.parse_args(argv) + + try: + snapshot_paths = discover_snapshot_images(args.input) + output_dir = args.output_dir or _infer_metadata_output_dir(args.input) + paths = extract_metadata_files( + snapshot_paths, + output_dir, + force=args.force, + metadata_provider=args.metadata_provider, + ) + except (FileNotFoundError, ValueError) as error: + parser.error(str(error)) + except RuntimeError as error: + print(f"error: {error}", file=sys.stderr) + return 1 + + for path in paths: + print(path) + return 0 + + +def _snapshot_parser(prog: str, description: str) -> argparse.ArgumentParser: + """Build the shared PDF extraction argument parser.""" + parser = argparse.ArgumentParser(prog=prog, description=description) + parser.add_argument( + "--model", + default="doclayoutyolo", + help="One extractor: doclayoutyolo (default), tfid, yolo11, or yolo26.", + ) + parser.add_argument( + "--input", + required=True, + type=Path, + help="A PDF file or directory containing PDFs.", + ) + parser.add_argument( + "--output-dir", + type=Path, + default=Path("output"), + help="Output root (default: output).", + ) + parser.add_argument( + "--dpi", + type=_positive_int, + default=300, + help="PDF rendering resolution (default: 300).", + ) + parser.add_argument( + "-f", + "--force", + action="store_true", + help="Delete matching crops and metadata before extraction.", + ) + return parser + + +def _add_metadata_provider_argument(parser: argparse.ArgumentParser) -> None: + """Add the shared metadata provider option to a parser.""" + parser.add_argument( + "--metadata-provider", + choices=("openrouter", "openai"), + default="openrouter", + help="Metadata API provider (default: openrouter).", + ) + + +def _run_snapshot_command( + parser: argparse.ArgumentParser, + argv: list[str] | None, + *, + end_to_end: bool, +) -> int: + """Run one PDF command and print its output paths.""" + args = parser.parse_args(argv) + load_dotenv(Path.cwd() / ".env", override=False) + try: + if end_to_end: + paths = run_pipeline( + args.input, + args.output_dir, + args.model, + dpi=args.dpi, + force=args.force, + metadata_provider=args.metadata_provider, + ) + else: + paths = run_extraction( + args.input, + args.output_dir, + args.model, + dpi=args.dpi, + force=args.force, + ) + except (FileExistsError, FileNotFoundError, ValueError) as error: + parser.error(str(error)) + except RuntimeError as error: + print(f"error: {error}", file=sys.stderr) + return 1 + + for path in paths: + print(path) + return 0 + + +def _infer_metadata_output_dir(input_path: Path) -> Path: + """Infer a sibling metadata directory from a snapshots input path.""" + snapshots_dir = input_path if input_path.is_dir() else input_path.parent + if snapshots_dir.name != "snapshots": + raise ValueError( + "--output-dir is required unless the input is in a snapshots directory." + ) + return snapshots_dir.parent / "metadata" diff --git a/src/ai4data/data_snapshot/config/models.json b/src/ai4data/data_snapshot/config/models.json new file mode 100644 index 00000000..8e397743 --- /dev/null +++ b/src/ai4data/data_snapshot/config/models.json @@ -0,0 +1,22 @@ +{ + "doclayoutyolo": { + "repo_id": "juliozhao/DocLayout-YOLO-DocStructBench", + "filename": "doclayout_yolo_docstructbench_imgsz1024.pt", + "cache_dir": ".cache/models" + }, + "yolo11": { + "repo_id": "Armaggheddon/yolo11-document-layout", + "filename": "yolo11m_doc_layout.pt", + "cache_dir": ".cache/models" + }, + "yolo26": { + "repo_id": "Armaggheddon/yolo26-document-layout", + "filename": "yolo26m_doc_layout.pt", + "cache_dir": ".cache/models" + }, + "tfid": { + "model_id": "yifeihu/TF-ID-large", + "revision": "a91ffef47039768b2deff82681b75d77252703f2", + "cache_dir": ".cache/models/tfid" + } +} diff --git a/src/ai4data/data_snapshot/detection.py b/src/ai4data/data_snapshot/detection.py new file mode 100644 index 00000000..1cc98950 --- /dev/null +++ b/src/ai4data/data_snapshot/detection.py @@ -0,0 +1,61 @@ +"""Define canonical data snapshot detections and detector protocol.""" + +from dataclasses import dataclass +from typing import Literal, Protocol + +from PIL import Image + + +@dataclass(frozen=True) +class Detection: + """Represent one normalized figure or table detection. + + Parameters + ---------- + label : {"figure", "table"} + Canonical snapshot label. + bbox : tuple[float, float, float, float] + Normalized ``(x1, y1, x2, y2)`` coordinates in the range ``[0, 1]``. + score : float + Model confidence in the range ``[0, 1]``. + + Raises + ------ + ValueError + If the label, bounding box, or score is invalid. + """ + + label: Literal["figure", "table"] + bbox: tuple[float, float, float, float] + score: float + + def __post_init__(self) -> None: + """Validate the canonical detection values.""" + if self.label not in {"figure", "table"}: + raise ValueError(f"Unsupported detection label: {self.label}") + x1, y1, x2, y2 = self.bbox + if not all(0.0 <= coordinate <= 1.0 for coordinate in self.bbox): + raise ValueError("Detection coordinates must be between zero and one.") + if x2 <= x1 or y2 <= y1: + raise ValueError("Detection bounding box must have positive area.") + if not 0.0 <= self.score <= 1.0: + raise ValueError("Detection score must be between zero and one.") + + +class SnapshotDetector(Protocol): + """Define the page-image interface implemented by snapshot detectors.""" + + def detect(self, image: Image.Image) -> list[Detection]: + """Return canonical detections for one rendered PDF page. + + Parameters + ---------- + image : PIL.Image.Image + Rendered RGB PDF page. + + Returns + ------- + list[Detection] + Figure and table detections for the page. + """ + ... diff --git a/src/ai4data/data_snapshot/extract_data_snapshot.py b/src/ai4data/data_snapshot/extract_data_snapshot.py new file mode 100644 index 00000000..60a5598a --- /dev/null +++ b/src/ai4data/data_snapshot/extract_data_snapshot.py @@ -0,0 +1,5 @@ +"""Run data snapshot extraction as a Python module.""" + +from ai4data.data_snapshot.cli import snapshot_main + +raise SystemExit(snapshot_main()) diff --git a/src/ai4data/data_snapshot/extract_metadata.py b/src/ai4data/data_snapshot/extract_metadata.py new file mode 100644 index 00000000..a346c933 --- /dev/null +++ b/src/ai4data/data_snapshot/extract_metadata.py @@ -0,0 +1,5 @@ +"""Run metadata extraction as a Python module.""" + +from ai4data.data_snapshot.cli import metadata_main + +raise SystemExit(metadata_main()) diff --git a/src/ai4data/data_snapshot/extraction.py b/src/ai4data/data_snapshot/extraction.py new file mode 100644 index 00000000..1088575f --- /dev/null +++ b/src/ai4data/data_snapshot/extraction.py @@ -0,0 +1,120 @@ +"""Extract data snapshot images from PDF files.""" + +import re +from math import ceil, floor +from pathlib import Path + +from PIL import Image + +from ai4data.data_snapshot.detection import Detection, SnapshotDetector +from ai4data.data_snapshot.inputs import discover_pdfs +from ai4data.data_snapshot.naming import build_snapshot_stem +from ai4data.data_snapshot.rendering import iter_pdf_pages + + +def extract_snapshots( + input_path: str | Path, + output_dir: str | Path, + detector: SnapshotDetector, + *, + dpi: int = 300, + force: bool = False, +) -> list[Path]: + """Detect and save figure and table snapshots from PDF inputs. + + Parameters + ---------- + input_path : str | Path + One PDF file or a directory searched recursively for PDFs. + output_dir : str | Path + Root output directory. Images are written beneath ``snapshots``. + detector : SnapshotDetector + Page-image detector producing canonical detections. + dpi : int, optional + PyMuPDF rendering resolution in dots per inch. + force : bool, optional + Delete existing snapshots for each source PDF before extracting it. + + Returns + ------- + list[Path] + Available snapshot image paths in document, page, and spatial order. + + Raises + ------ + FileExistsError + If a destination snapshot appears while extraction is running. + ValueError + If multiple source PDFs have the same filename stem. + """ + pdf_paths = discover_pdfs(input_path) + validate_unique_stems(pdf_paths) + snapshots_dir = Path(output_dir) / "snapshots" + snapshots_dir.mkdir(parents=True, exist_ok=True) + output_paths: list[Path] = [] + + for pdf_path in pdf_paths: + existing_paths = _existing_snapshot_paths(snapshots_dir, pdf_path.stem) + if existing_paths and not force: + print( + f"Skipping {pdf_path}: existing snapshots found in " + f"{snapshots_dir} (use --force to replace)." + ) + output_paths.extend(existing_paths) + continue + for existing_path in existing_paths: + existing_path.unlink() + + counters = {"figure": 0, "table": 0} + for _, page_image in iter_pdf_pages(pdf_path, dpi=dpi): + detections = sorted( + detector.detect(page_image), + key=lambda item: ( + item.bbox[1], + item.bbox[0], + item.bbox[3], + item.bbox[2], + ), + ) + for detection in detections: + index = counters[detection.label] + counters[detection.label] += 1 + output_path = snapshots_dir / ( + f"{build_snapshot_stem(pdf_path, detection.label, index)}.png" + ) + if output_path.exists(): + raise FileExistsError(f"Snapshot already exists: {output_path}") + _crop_snapshot(page_image, detection).save(output_path, format="PNG") + output_paths.append(output_path) + + return output_paths + + +def _existing_snapshot_paths(snapshots_dir: Path, pdf_stem: str) -> list[Path]: + """Return extraction outputs belonging to one source PDF stem.""" + pattern = re.compile(rf"{re.escape(pdf_stem)}_(?:figure|table)_\d{{3,}}\.png") + return sorted( + path + for path in snapshots_dir.iterdir() + if path.is_file() and pattern.fullmatch(path.name) + ) + + +def validate_unique_stems(pdf_paths: list[Path]) -> None: + """Reject source documents that would share output filenames.""" + stems = [path.stem for path in pdf_paths] + duplicates = sorted({stem for stem in stems if stems.count(stem) > 1}) + if duplicates: + raise ValueError(f"Source PDF filename stems must be unique: {duplicates}") + + +def _crop_snapshot(image: Image.Image, detection: Detection) -> Image.Image: + """Crop a normalized detection from its rendered page image.""" + x1, y1, x2, y2 = detection.bbox + box = ( + floor(x1 * image.width), + floor(y1 * image.height), + ceil(x2 * image.width), + ceil(y2 * image.height), + ) + return image.crop(box) diff --git a/src/ai4data/data_snapshot/inputs.py b/src/ai4data/data_snapshot/inputs.py new file mode 100644 index 00000000..719c6581 --- /dev/null +++ b/src/ai4data/data_snapshot/inputs.py @@ -0,0 +1,84 @@ +"""Discover PDF inputs for the data snapshot pipeline.""" + +from pathlib import Path + +_IMAGE_SUFFIXES = {".gif", ".jpeg", ".jpg", ".png", ".webp"} + + +def discover_pdfs(input_path: str | Path) -> list[Path]: + """Return PDF files represented by a file or directory input. + + Directory inputs are searched recursively and returned in deterministic + path order. + + Parameters + ---------- + input_path : str | Path + A PDF file or directory containing PDF files. + + Returns + ------- + list[Path] + Discovered PDF paths. + + Raises + ------ + FileNotFoundError + If the input does not exist or a directory contains no PDF files. + ValueError + If a file input does not have a ``.pdf`` extension. + """ + path = Path(input_path) + if not path.exists(): + raise FileNotFoundError(f"Input does not exist: {path}") + if path.is_file(): + if path.suffix.lower() != ".pdf": + raise ValueError(f"Input file is not a PDF: {path}") + return [path] + + pdf_paths = sorted( + candidate + for candidate in path.rglob("*") + if candidate.is_file() and candidate.suffix.lower() == ".pdf" + ) + if not pdf_paths: + raise FileNotFoundError(f"No PDF files found under: {path}") + return pdf_paths + + +def discover_snapshot_images(input_path: str | Path) -> list[Path]: + """Return supported snapshot images represented by a file or directory. + + Parameters + ---------- + input_path : str | Path + One snapshot image or a directory searched recursively for images. + + Returns + ------- + list[Path] + Discovered image paths in deterministic order. + + Raises + ------ + FileNotFoundError + If the input does not exist or a directory contains no images. + ValueError + If a file input does not use a supported image extension. + """ + path = Path(input_path) + if not path.exists(): + raise FileNotFoundError(f"Input does not exist: {path}") + if path.is_file(): + if path.suffix.lower() not in _IMAGE_SUFFIXES: + raise ValueError(f"Input file is not a supported image: {path}") + return [path] + + image_paths = sorted( + candidate + for candidate in path.rglob("*") + if candidate.is_file() and candidate.suffix.lower() in _IMAGE_SUFFIXES + ) + if not image_paths: + raise FileNotFoundError(f"No snapshot images found under: {path}") + return image_paths diff --git a/src/ai4data/data_snapshot/metadata_extraction/__init__.py b/src/ai4data/data_snapshot/metadata_extraction/__init__.py new file mode 100644 index 00000000..139acc08 --- /dev/null +++ b/src/ai4data/data_snapshot/metadata_extraction/__init__.py @@ -0,0 +1,17 @@ +"""Extract Schema v1.4 metadata from individual snapshot images.""" + +from ai4data.data_snapshot.metadata_extraction.batch import extract_metadata_files +from ai4data.data_snapshot.metadata_extraction.extraction import ( + ExtractionResult, + extract_metadata, + get_extraction_config_path, + load_extraction_config, +) + +__all__ = [ + "ExtractionResult", + "extract_metadata", + "extract_metadata_files", + "get_extraction_config_path", + "load_extraction_config", +] diff --git a/src/ai4data/data_snapshot/metadata_extraction/batch.py b/src/ai4data/data_snapshot/metadata_extraction/batch.py new file mode 100644 index 00000000..2092bebb --- /dev/null +++ b/src/ai4data/data_snapshot/metadata_extraction/batch.py @@ -0,0 +1,92 @@ +"""Write validated metadata JSON for extracted snapshot images.""" + +from pathlib import Path +from typing import Iterable + +from tqdm.auto import tqdm + +from ai4data.data_snapshot.metadata_extraction.extraction import ( + extract_metadata, + get_extraction_config_path, +) + + +def extract_metadata_files( + snapshot_paths: Iterable[str | Path], + output_dir: str | Path, + *, + force: bool = False, + metadata_provider: str = "openrouter", +) -> list[Path]: + """Extract and save metadata for snapshot images. + + Parameters + ---------- + snapshot_paths : Iterable[str | Path] + Snapshot images to process. + output_dir : str | Path + Directory for one JSON file per snapshot. + force : bool, optional + Delete and replace existing metadata for each snapshot. + metadata_provider : str, optional + API provider whose packaged request settings are used. + + Returns + ------- + list[Path] + Available metadata JSON paths in snapshot order. + + Raises + ------ + RuntimeError + If metadata extraction or validation fails for a snapshot. + """ + metadata_dir = Path(output_dir) + config_path = get_extraction_config_path(metadata_provider) + output_paths: list[Path] = [] + paths = [Path(path) for path in snapshot_paths] + seen_stems: set[str] = set() + for path in paths: + if path.stem in seen_stems: + raise ValueError( + f"Snapshot image filename stems must be unique: {path.stem}" + ) + seen_stems.add(path.stem) + + snapshot_progress = tqdm(paths, desc="Metadata", unit="snapshot") + for snapshot_path in snapshot_progress: + snapshot_progress.set_postfix_str(snapshot_path.name) + metadata_path = metadata_dir / f"{snapshot_path.stem}.json" + output_paths.append(metadata_path) + if metadata_path.exists(): + if force: + print( + f"Existing metadata output for {snapshot_path.name} exists. " + "Cleaning up before continuing...", + end=" ", + flush=True, + ) + metadata_path.unlink() + print("Done!") + else: + print( + f"Existing metadata output for {snapshot_path.name} exists. " + "Skipped metadata extraction." + ) + continue + + result = extract_metadata(snapshot_path, config_path=config_path) + if result.metadata is None: + detail = result.error or "unknown extraction error" + raise RuntimeError( + f"Metadata extraction failed for {snapshot_path}: {detail}" + ) + metadata_dir.mkdir(parents=True, exist_ok=True) + temporary_path = metadata_path.with_suffix(".json.tmp") + temporary_path.write_text( + result.metadata.model_dump_json(indent=2) + "\n", + encoding="utf-8", + ) + temporary_path.replace(metadata_path) + + return output_paths diff --git a/src/ai4data/data_snapshot/metadata_extraction/config/openai.json b/src/ai4data/data_snapshot/metadata_extraction/config/openai.json new file mode 100644 index 00000000..9f798d77 --- /dev/null +++ b/src/ai4data/data_snapshot/metadata_extraction/config/openai.json @@ -0,0 +1,11 @@ +{ + "provider": "openai", + "model": "gpt-5.6-luna", + "max_output_tokens": 25000, + "service_tier": "flex", + "prompt_cache_key": "metadata-extraction-v1-4", + "timeout": 900.0, + "reasoning": { + "effort": "medium" + } +} diff --git a/src/ai4data/data_snapshot/metadata_extraction/config/openrouter.json b/src/ai4data/data_snapshot/metadata_extraction/config/openrouter.json new file mode 100644 index 00000000..72682937 --- /dev/null +++ b/src/ai4data/data_snapshot/metadata_extraction/config/openrouter.json @@ -0,0 +1,10 @@ +{ + "provider": "openrouter", + "model": "openai/gpt-5.6-luna", + "max_output_tokens": 25000, + "service_tier": "flex", + "timeout": 900.0, + "reasoning": { + "effort": "medium" + } +} \ No newline at end of file diff --git a/src/ai4data/data_snapshot/metadata_extraction/extraction.py b/src/ai4data/data_snapshot/metadata_extraction/extraction.py new file mode 100644 index 00000000..36997f97 --- /dev/null +++ b/src/ai4data/data_snapshot/metadata_extraction/extraction.py @@ -0,0 +1,437 @@ +"""Extract validated Schema v1.4 metadata from one snapshot image.""" + +from __future__ import annotations + +import base64 +import json +import mimetypes +import os +import time +from dataclasses import dataclass +from functools import cache +from pathlib import Path +from typing import Any + +from ai4data.data_snapshot.metadata_schema import DataSnapshotMetadata + +_PACKAGE_DIR = Path(__file__).parent +_CONFIG_DIR = _PACKAGE_DIR / "config" +_DEFAULT_CONFIG_PATH = _CONFIG_DIR / "openrouter.json" +_PROMPT_DIR = _PACKAGE_DIR / "prompts" +_COMPLETENESS_PLACEHOLDER = "{{COMPLETENESS_GUIDANCE}}" +_SCHEMA_PLACEHOLDER = "{{MODEL_FACING_SCHEMA_SECTION}}" +_COMPLETENESS_GUIDANCE = ( + "Before returning, re-scan the image against the entire schema and correct any " + "omissions of fields supported by visible evidence. Leave every unsupported " + "field null." +) +_PIPELINE_OWNED_CONFIG_FIELDS = { + "input", + "instructions", + "model", + "provider", + "stream", + "text", + "text_format", +} +_SUPPORTED_PROVIDERS = {"openai", "openrouter"} +_UNSUPPORTED_SCHEMA_KEYWORDS = { + "$schema", + "allOf", + "default", + "dependentRequired", + "dependentSchemas", + "else", + "examples", + "if", + "not", + "then", +} +_SUPPORTED_STRING_FORMATS = { + "date", + "date-time", + "duration", + "email", + "hostname", + "ipv4", + "ipv6", + "time", + "uuid", +} + + +@dataclass(frozen=True) +class ExtractionResult: + """Represent one metadata extraction outcome. + + Parameters + ---------- + metadata : DataSnapshotMetadata | None + Validated Schema v1.4 metadata, or ``None`` after a failure. + model : str | None + Requested model, when configuration loading succeeded. + response_id : str | None + OpenAI response identifier, when available. + api_status : str | None + OpenAI response status, when available. + incomplete_details : dict[str, Any] | None + Structured explanation for an incomplete response, when available. + elapsed_seconds : float | None + API-call duration, or ``None`` when no call was attempted. + usage : dict[str, Any] | None + JSON-compatible token usage returned by the API. + raw_output : str | None + Structured output text returned by the API, when available. + error_type : str | None + Exception class name after a failure. + error : str | None + Error message after a failure. + """ + + metadata: DataSnapshotMetadata | None + model: str | None + response_id: str | None + api_status: str | None + incomplete_details: dict[str, Any] | None + elapsed_seconds: float | None + usage: dict[str, Any] | None + raw_output: str | None + error_type: str | None + error: str | None + + +def load_extraction_config( + config_path: str | Path = _DEFAULT_CONFIG_PATH, +) -> dict[str, Any]: + """Load metadata API settings from a JSON object. + + Parameters + ---------- + config_path : str | Path, optional + JSON file containing ``model``, an optional ``provider``, and API + settings. + + Returns + ------- + dict[str, Any] + Validated request settings. + + Raises + ------ + ValueError + If the config is not an object, lacks a model, or overrides a field + owned by the extraction pipeline. + """ + with Path(config_path).open(encoding="utf-8") as file: + config = json.load(file) + if not isinstance(config, dict): + raise ValueError("Extraction config must be a JSON object.") + model = config.get("model") + if not isinstance(model, str) or not model.strip(): + raise ValueError("Extraction config requires a non-empty 'model' string.") + provider = config.get("provider", "openai") + if provider not in _SUPPORTED_PROVIDERS: + raise ValueError("Extraction config provider must be 'openai' or 'openrouter'.") + forbidden = (set(config) - {"model", "provider"}) & (_PIPELINE_OWNED_CONFIG_FIELDS) + if forbidden: + raise ValueError( + f"Extraction config cannot set pipeline-owned keys: {sorted(forbidden)}." + ) + return config + + +def get_extraction_config_path(provider: str) -> Path: + """Return the packaged configuration for a metadata provider. + + Parameters + ---------- + provider : str + Supported metadata provider name. + + Returns + ------- + pathlib.Path + Path to the provider's packaged JSON configuration. + + Raises + ------ + ValueError + If the provider is not supported. + """ + normalized_provider = provider.strip().lower() + if normalized_provider not in _SUPPORTED_PROVIDERS: + raise ValueError("Metadata provider must be 'openai' or 'openrouter'.") + return _CONFIG_DIR / f"{normalized_provider}.json" + + +def extract_metadata( + image_path: str | Path, + config_path: str | Path = _DEFAULT_CONFIG_PATH, + client: Any | None = None, +) -> ExtractionResult: + """Extract validated Schema v1.4 metadata from one snapshot image. + + Only image-visible evidence is sent to the model. Filenames, directory + names, and source-document metadata are excluded from the prompt. + + Parameters + ---------- + image_path : str | Path + Path to a PNG, JPEG, GIF, or WebP snapshot image. + config_path : str | Path, optional + Responses API configuration file. + client : Any | None, optional + Injected OpenAI-compatible client for offline tests. + + Returns + ------- + ExtractionResult + Validated metadata and call details, or structured failure details. + """ + model: str | None = None + started_at: float | None = None + response: Any | None = None + try: + config = load_extraction_config(config_path) + model = config.pop("model") + provider = config.pop("provider", "openai") + system_prompt = _read_prompt( + "system.md", _COMPLETENESS_PLACEHOLDER, _COMPLETENESS_GUIDANCE + ) + user_prompt = _read_prompt("user.md", _SCHEMA_PLACEHOLDER, "") + image_url = _image_data_url(image_path) + api_client = client if client is not None else _create_api_client(provider) + + started_at = time.perf_counter() + response = api_client.responses.create( + model=model, + text={"format": _response_format()}, + input=[ + { + "role": "system", + "content": [{"type": "input_text", "text": system_prompt}], + }, + { + "role": "user", + "content": [ + {"type": "input_text", "text": user_prompt}, + {"type": "input_image", "image_url": image_url}, + ], + }, + ], + **config, + ) + elapsed_seconds = time.perf_counter() - started_at + status = getattr(response, "status", None) + incomplete_details = _serialize_api_value( + getattr(response, "incomplete_details", None) + ) + if status == "incomplete": + reason = (incomplete_details or {}).get("reason", "unknown reason") + raise ValueError(f"Response was incomplete: {reason}.") + if status not in {None, "completed"}: + detail = _api_error_text(getattr(response, "error", None)) + suffix = f" ({detail})" if detail else "" + raise ValueError(f"Response ended with status: {status}{suffix}.") + raw_output = getattr(response, "output_text", None) + if not isinstance(raw_output, str) or not raw_output: + raise ValueError("Structured output was missing.") + metadata = DataSnapshotMetadata.model_validate_json(raw_output) + return ExtractionResult( + metadata=metadata, + model=model, + response_id=getattr(response, "id", None), + api_status=status, + incomplete_details=incomplete_details, + elapsed_seconds=elapsed_seconds, + usage=_serialize_api_value(getattr(response, "usage", None)), + raw_output=raw_output, + error_type=None, + error=None, + ) + except Exception as error: + elapsed_seconds = ( + time.perf_counter() - started_at if started_at is not None else None + ) + return ExtractionResult( + metadata=None, + model=model, + response_id=getattr(response, "id", None), + api_status=getattr(response, "status", None), + incomplete_details=_serialize_api_value( + getattr(response, "incomplete_details", None) + ), + elapsed_seconds=elapsed_seconds, + usage=_serialize_api_value(getattr(response, "usage", None)), + raw_output=getattr(response, "output_text", None), + error_type=type(error).__name__, + error=str(error), + ) + + +def _read_prompt(filename: str, placeholder: str, replacement: str) -> str: + """Read one prompt and replace its required placeholder.""" + prompt = (_PROMPT_DIR / filename).read_text(encoding="utf-8") + if prompt.count(placeholder) != 1: + raise ValueError(f"Prompt {filename} must contain {placeholder} exactly once.") + return prompt.replace(placeholder, replacement).rstrip() + "\n" + + +def _image_data_url(image_path: str | Path) -> str: + """Encode a supported local image as a base64 data URL.""" + path = Path(image_path) + mime_type, _ = mimetypes.guess_type(path.name) + if mime_type not in {"image/gif", "image/jpeg", "image/png", "image/webp"}: + raise ValueError(f"Unsupported image file: {path}") + encoded = base64.b64encode(path.read_bytes()).decode("ascii") + return f"data:{mime_type};base64,{encoded}" + + +def _openai_compatible_schema(value: Any) -> Any: + """Remove unsupported keywords from canonical Pydantic JSON Schema data.""" + if isinstance(value, list): + return [_openai_compatible_schema(item) for item in value] + if not isinstance(value, dict): + return value + + validation_rules = value.get("x-validation-rules") + result: dict[str, Any] = {} + for key, item in value.items(): + if key in _UNSUPPORTED_SCHEMA_KEYWORDS or key.startswith("x-"): + continue + if key == "format" and item not in _SUPPORTED_STRING_FORMATS: + continue + if key == "anyOf" and _contains_schema_keyword(item, "not"): + continue + result[key] = _openai_compatible_schema(item) + + if isinstance(validation_rules, list) and all( + isinstance(rule, str) for rule in validation_rules + ): + guidance = "Validation rules:\n- " + "\n- ".join(validation_rules) + description = result.get("description") + result["description"] = ( + f"{description}\n\n{guidance}" if description else guidance + ) + properties = result.get("properties") + if result.get("type") == "object" and isinstance(properties, dict): + result["required"] = list(properties) + result["additionalProperties"] = False + return result + + +def _contains_schema_keyword(value: Any, keyword: str) -> bool: + """Return whether nested JSON Schema data contain a named keyword.""" + if isinstance(value, dict): + return keyword in value or any( + _contains_schema_keyword(item, keyword) for item in value.values() + ) + if isinstance(value, list): + return any(_contains_schema_keyword(item, keyword) for item in value) + return False + + +def _inline_ref_siblings(value: Any, root: dict[str, Any]) -> Any: + """Inline local references that have sibling schema keywords.""" + if isinstance(value, list): + return [_inline_ref_siblings(item, root) for item in value] + if not isinstance(value, dict): + return value + + reference = value.get("$ref") + if isinstance(reference, str) and len(value) > 1: + if not reference.startswith("#/"): + raise ValueError(f"Unsupported JSON Schema reference: {reference}") + resolved: Any = root + for part in reference[2:].split("/"): + key = part.replace("~1", "/").replace("~0", "~") + if not isinstance(resolved, dict) or key not in resolved: + raise ValueError(f"Unresolved JSON Schema reference: {reference}") + resolved = resolved[key] + if not isinstance(resolved, dict): + raise ValueError(f"JSON Schema reference is not an object: {reference}") + value = {**resolved, **value} + value.pop("$ref") + return {key: _inline_ref_siblings(item, root) for key, item in value.items()} + + +@cache +def _response_format() -> dict[str, Any]: + """Build the strict OpenAI response format from the canonical model.""" + canonical_schema = DataSnapshotMetadata.model_json_schema(mode="validation") + schema = _openai_compatible_schema(canonical_schema) + schema = _inline_ref_siblings(schema, schema) + return { + "type": "json_schema", + "name": "data_snapshot_metadata_v1_4", + "strict": True, + "schema": schema, + } + + +def _create_api_client(provider: str) -> Any: + """Create the configured API client with automatic retries disabled. + + Parameters + ---------- + provider : str + Supported API provider name. + + Returns + ------- + Any + Configured OpenAI-compatible SDK client. + + Raises + ------ + ValueError + If the selected provider's API key is not set. + """ + from dotenv import load_dotenv + from openai import OpenAI + + load_dotenv(Path.cwd() / ".env") + if provider == "openrouter": + api_key_name = "OPENROUTER_API_KEY" + base_url = "https://openrouter.ai/api/v1" + else: + api_key_name = "OPENAI_API_KEY" + base_url = None + api_key = os.getenv(api_key_name) + if not api_key: + raise ValueError(f"{api_key_name} is not set.") + return OpenAI(api_key=api_key, base_url=base_url, max_retries=0) + + +def _serialize_api_value(value: Any) -> dict[str, Any] | None: + """Convert an OpenAI SDK object to a JSON-compatible dictionary.""" + if value is None: + return None + if hasattr(value, "model_dump"): + return value.model_dump(mode="json", warnings=False) + if isinstance(value, dict): + return value + return { + key: getattr(value, key) + for key in ( + "input_tokens", + "output_tokens", + "reason", + "total_tokens", + ) + if hasattr(value, key) + } + + +def _api_error_text(value: Any) -> str | None: + """Return an API error code and message suitable for CLI output.""" + if value is None: + return None + if isinstance(value, dict): + code = value.get("code") + message = value.get("message") + else: + code = getattr(value, "code", None) + message = getattr(value, "message", None) + detail = ": ".join(str(item) for item in (code, message) if item) + return detail or str(value) diff --git a/src/ai4data/data_snapshot/metadata_extraction/prompts/system.md b/src/ai4data/data_snapshot/metadata_extraction/prompts/system.md new file mode 100644 index 00000000..b875a920 --- /dev/null +++ b/src/ai4data/data_snapshot/metadata_extraction/prompts/system.md @@ -0,0 +1,22 @@ +You extract Schema v1.4 semantic metadata from one data snapshot image. + +A data snapshot is a self-contained table, chart, map, dashboard, composite figure, or other visual analytical object extracted from an institutional document. + +## Evidence boundary + +- Use only evidence visible in the supplied image, including its text, labels, legends, captions, annotations, and visual structure. +- Do not use the filename, directory structure, source-document metadata, or external knowledge. These are not provided as evidence. +- Populate a field only when the image supports it. Leave unsupported optional fields null. +- Preserve source-visible terminology and complete interpretive statements whenever the schema requests source text. +- Apply normalized values only when the mapping is deterministic and permitted by the schema. Do not replace required source-visible text with a normalized label. +- Resolve neither ambiguity nor conflict by guessing. Prefer null or a narrower supported representation. + +## Metadata boundary + +Describe the snapshot and the data represented in it. Do not reconstruct observations, digitize chart values, reproduce table cells, perform general OCR, summarize a parent document, or invent document-level context. + +Use visual form as evidence for visualization metadata, and semantic meaning—not nearby text copied mechanically—to choose fields. Keep variables, dimensions, qualifiers, provenance, temporal coverage, and geographic coverage in their schema-defined roles. + +{{COMPLETENESS_GUIDANCE}} + +Return only the structured metadata required by the response format. diff --git a/src/ai4data/data_snapshot/metadata_extraction/prompts/user.md b/src/ai4data/data_snapshot/metadata_extraction/prompts/user.md new file mode 100644 index 00000000..fd911f02 --- /dev/null +++ b/src/ai4data/data_snapshot/metadata_extraction/prompts/user.md @@ -0,0 +1,33 @@ +Inspect the attached data snapshot and extract all Schema v1.4 metadata explicitly supported by the image. Return the structured metadata only. + +## Field-boundary guidance + +- Preserve the complete visible title or caption in `title`. Also populate + `document_label` when a visible Figure, Table, Annex, or Exhibit label is present. +- Use `source_text` only for exact visible wording. Never invent source wording for + a value inferred from visual form or structure. +- Variables are measured concepts, indicators, or metrics. Dimensions classify, + group, organize, or compare those measures. Do not duplicate a concept as both + unless the image clearly gives it both roles. +- Coefficient, standard error, count, percentage, and confidence interval are + statistical forms or qualifiers on the applicable variable; they are not + standalone measured variables when the measured concept is identifiable. +- Use `category_groups` for explicit group headings and their contained categories. + Use `panel_titles` only for explicit panels in a multi-panel snapshot, not ordinary + section headings within one table. +- Express a visible comparative relationship as one comparison, such as + `Baseline vs Closing Period`, rather than separate unrelated comparison items. +- Put derivation datasets and surveys in `provenance.sources`; put credited agents + in `provenance.attributions`. Do not repeat a source statement in + `interpretive_notes` when its meaning is fully represented elsewhere. +- Temporal coverage describes when the represented data apply, not publication or + footer dates. Normalize visible dates only when the mapping is deterministic. +- Geographic scope is the overall represented area. Use locations for additional + named places and level for the level at which the data are reported. +- Always populate `languages` when the language of the visible snapshot text is + unambiguous. Use the corresponding BCP 47 `tag`; use `source_text` only when the + snapshot explicitly names the language. +- Preserve complete substantive notes that are not fully represented by another + field. Do not extract numerical observations or reproduce table cells. + +{{MODEL_FACING_SCHEMA_SECTION}} diff --git a/src/ai4data/data_snapshot/metadata_schema/__init__.py b/src/ai4data/data_snapshot/metadata_schema/__init__.py new file mode 100644 index 00000000..c7d9e948 --- /dev/null +++ b/src/ai4data/data_snapshot/metadata_schema/__init__.py @@ -0,0 +1,77 @@ +"""Expose the canonical Data Snapshot Metadata Schema v1.4 API.""" + +from ai4data.data_snapshot.metadata_schema.models import ( + AnalyticalRole, + Attribution, + AxisDimension, + AxisPosition, + AxisRole, + CategoryGroup, + CodedTerm, + Currency, + DataSnapshotMetadata, + Dimension, + EntityReference, + Financing, + GeographicCoverage, + GeographicLevelTerm, + GeographicLevelValue, + GeographicLocation, + Identifier, + Language, + MultiAxisAssignment, + Place, + PresentationRole, + Project, + Provenance, + ReportingIntervalTerm, + ReportingIntervalValue, + StatisticalFormTerm, + StatisticalFormValue, + TemporalCoverage, + TemporalExpression, + TemporalPrecision, + TemporalRelation, + Unit, + Variable, + VisualizationTypeTerm, + VisualizationTypeValue, +) + +__all__ = [ + "AnalyticalRole", + "Attribution", + "AxisDimension", + "AxisPosition", + "AxisRole", + "CategoryGroup", + "CodedTerm", + "Currency", + "DataSnapshotMetadata", + "Dimension", + "EntityReference", + "Financing", + "GeographicCoverage", + "GeographicLevelTerm", + "GeographicLevelValue", + "GeographicLocation", + "Identifier", + "Language", + "MultiAxisAssignment", + "Place", + "PresentationRole", + "Project", + "Provenance", + "ReportingIntervalTerm", + "ReportingIntervalValue", + "StatisticalFormTerm", + "StatisticalFormValue", + "TemporalCoverage", + "TemporalExpression", + "TemporalPrecision", + "TemporalRelation", + "Unit", + "Variable", + "VisualizationTypeTerm", + "VisualizationTypeValue", +] diff --git a/src/ai4data/data_snapshot/metadata_schema/models.py b/src/ai4data/data_snapshot/metadata_schema/models.py new file mode 100644 index 00000000..2e411675 --- /dev/null +++ b/src/ai4data/data_snapshot/metadata_schema/models.py @@ -0,0 +1,2310 @@ +"""Define the canonical Data Snapshot Metadata Schema v1.4 models.""" + +from __future__ import annotations + +import re +import unicodedata +from datetime import date, datetime +from decimal import Decimal +from enum import Enum +from typing import Annotated, Any + +from pydantic import ( + AnyUrl, + BaseModel, + BeforeValidator, + ConfigDict, + Field, + PositiveInt, + StringConstraints, + WithJsonSchema, + field_validator, + model_validator, +) + + +def _normalize_text(value: Any) -> Any: + if isinstance(value, str): + return unicodedata.normalize("NFC", value).strip() + return value + + +NonEmptyText = Annotated[ + str, + StringConstraints(strict=True, min_length=1, pattern=r"\S"), + BeforeValidator(_normalize_text), +] + +# RFC 5646 section 2.1; registry membership remains a separate concern. +_GRANDFATHERED_TAGS = ( + "en-GB-oed", + "i-ami", + "i-bnn", + "i-default", + "i-enochian", + "i-hak", + "i-klingon", + "i-lux", + "i-mingo", + "i-navajo", + "i-pwn", + "i-tao", + "i-tay", + "i-tsu", + "sgn-BE-FR", + "sgn-BE-NL", + "sgn-CH-DE", + "art-lojban", + "cel-gaulish", + "no-bok", + "no-nyn", + "zh-guoyu", + "zh-hakka", + "zh-min", + "zh-min-nan", + "zh-xiang", +) +_BCP47_PATTERN = re.compile( + r"^(?:(?:[A-Za-z]{2,3}(?:-[A-Za-z]{3}){0,3}|[A-Za-z]{4}|[A-Za-z]{5,8})" + r"(?:-[A-Za-z]{4})?(?:-(?:[A-Za-z]{2}|[0-9]{3}))?" + r"(?:-(?:[A-Za-z0-9]{5,8}|[0-9][A-Za-z0-9]{3}))*" + r"(?:-[0-9A-WY-Za-wy-z](?:-[A-Za-z0-9]{2,8})+)*" + r"(?:-[xX](?:-[A-Za-z0-9]{1,8})+)?|[xX](?:-[A-Za-z0-9]{1,8})+|" + + "|".join(_GRANDFATHERED_TAGS) + + r")$" +) +_YEAR_PATTERN = re.compile(r"^[0-9]{4}$") +_MONTH_PATTERN = re.compile(r"^[0-9]{4}-(?:0[1-9]|1[0-2])$") +_DAY_PATTERN = re.compile(r"^[0-9]{4}-(?:0[1-9]|1[0-2])-(?:0[1-9]|[12][0-9]|3[01])$") +_DATETIME_PATTERN = re.compile( + r"^[0-9]{4}-(?:0[1-9]|1[0-2])-(?:0[1-9]|[12][0-9]|3[01])" + r"T(?:[01][0-9]|2[0-3]):[0-5][0-9]" + r"(?::[0-5][0-9](?:[.,][0-9]+)?)?(?:Z|[+-][0-9]{2}:[0-5][0-9])$" +) + + +# RFC 3986 Appendix A. AnyUrl checks IP literals after this lexical check. +_URI_ATOM = r"(?:[A-Za-z0-9._~!$&'()*+,;=-]|%[0-9A-Fa-f]{2})" +_URI_PCHAR = rf"(?:{_URI_ATOM}|[:@])" +_URI_AUTHORITY = rf"(?:{_URI_ATOM}|:)*@" +_URI_PATTERN = re.compile( + rf"^[A-Za-z][A-Za-z0-9+.-]*:" + rf"(?://(?:{_URI_AUTHORITY})?(?:{_URI_ATOM}*|\[[A-Za-z0-9:.!$&'()*+,;=_~-]+\])" + rf"(?::[0-9]*)?(?:/{_URI_PCHAR}*)*" + rf"|/(?:{_URI_PCHAR}+(?:/{_URI_PCHAR}*)*)?" + rf"|{_URI_PCHAR}+(?:/{_URI_PCHAR}*)*|)" + rf"(?:\?(?:{_URI_PCHAR}|[/?])*)?(?:#(?:{_URI_PCHAR}|[/?])*)?$" +) + + +def _validate_uri(value: Any) -> Any: + # Check the original spelling before AnyUrl can strip controls or escape spaces. + if not isinstance(value, (str, AnyUrl)): + raise ValueError("URI must be a string or a validated URL.") + if isinstance(value, str) and not _URI_PATTERN.fullmatch(value): + raise ValueError("URI must use absolute RFC 3986 syntax and valid escapes.") + return value + + +AbsoluteURI = Annotated[ + AnyUrl, + BeforeValidator(_validate_uri), + Field(json_schema_extra={"pattern": _URI_PATTERN.pattern}), +] + + +def _populated(*names: str) -> dict[str, Any]: + return { + "required": list(names), + "properties": {name: {"not": {"type": "null"}} for name in names}, + } + + +def _content_schema(*names: str) -> dict[str, Any]: + return { + "anyOf": [_populated(name) for name in names], + "x-validation-rules": [ + "At least one non-null value is required: " + ", ".join(names) + "." + ], + } + + +def _variable_schema(schema: dict[str, Any]) -> None: + schema.update(_content_schema("name", "unit", "currency", "statistical_forms")) + schema["allOf"] = [ + { + "if": { + "anyOf": [ + _populated("analytical_roles"), + _populated("axis_roles"), + _populated("multi_axis_assignments"), + ] + }, + "then": _populated("name"), + } + ] + schema["x-validation-rules"].append( + "Analytical roles, axis roles, and multi-axis assignments require a " + "non-null variable name." + ) + + +def _temporal_schema(schema: dict[str, Any]) -> None: + schema["anyOf"] = [ + _populated("source_text"), + _populated("start"), + _populated("end"), + ] + schema["allOf"] = [ + { + "if": {"anyOf": [_populated("start"), _populated("end")]}, + "then": _populated("relation", "precision"), + "else": { + "properties": { + name: {"type": "null"} for name in ("relation", "precision") + } + }, + } + ] + for relation, alternatives in { + "point": [("start",)], + "interval": [("start", "end")], + "open_ended_interval": [("start",), ("end",)], + }.items(): + branches = [] + for names in alternatives: + branch = _populated(*names) + branch["properties"].update( + { + name: {"type": "null"} + for name in ("start", "end") + if name not in names + } + ) + branches.append(branch) + schema["allOf"].append( + { + "if": { + "required": ["relation"], + "properties": {"relation": {"const": relation}}, + }, + "then": {"anyOf": branches}, + } + ) + for precision, pattern in { + "year": _YEAR_PATTERN, + "month": _MONTH_PATTERN, + "day": _DAY_PATTERN, + "datetime": _DATETIME_PATTERN, + }.items(): + schema["allOf"].append( + { + "if": { + "required": ["precision"], + "properties": {"precision": {"const": precision}}, + }, + "then": { + "properties": { + name: {"pattern": pattern.pattern} for name in ("start", "end") + } + }, + } + ) + schema["x-validation-rules"] = [ + "At least one of source_text, start, or end must be non-null.", + "Normalized bounds require relation and precision; source-only expressions omit both.", + "point requires start only; interval requires both bounds; open_ended_interval requires exactly one bound.", + "Bounds must match the declared precision. Python additionally validates calendar dates and chronological ordering.", + ] + + +def _standards(*mappings: tuple[str, str]) -> dict[str, object]: + return { + "x-standards": [ + {"term": term, "relationship": relationship} + for term, relationship in mappings + ] + } + + +def _code_list( + authority: str, identifier: str, uri: str, release: str | None = None +) -> dict[str, object]: + metadata: dict[str, str] = { + "authority": authority, + "identifier": identifier, + "uri": uri, + } + if release is not None: + metadata["release"] = release + return {"x-code-list": metadata} + + +class _SchemaModel(BaseModel): + model_config = ConfigDict(extra="forbid") + + @field_validator("*", mode="after") + @classmethod + def _remove_exact_duplicates(cls, value: Any) -> Any: + if not isinstance(value, list): + return value + unique: list[Any] = [] + for item in value: + if item not in unique: + unique.append(item) + return unique + + +class PresentationRole(str, Enum): + """Identify how a dimension organizes a table. + + `row` means the dimension's categories vary downward through table rows. + `column` means they vary horizontally across table columns. The physical + placement of a label in the first column does not by itself make its + dimension a column dimension. + """ + + ROW = "row" + COLUMN = "column" + + +class AnalyticalRole(str, Enum): + """Identify an explicit role in a statistical analysis. + + `outcome` is the dependent or response variable the analysis seeks to + explain or estimate. `predictor` is an explanatory or independent variable + used to explain or predict an outcome. `control` is a covariate included to + adjust an estimate rather than serving as the primary predictor. + `instrumental` is a variable explicitly identified as an instrument in an + instrumental-variable analysis. Assign these roles only when the analytical + table or its labels, title, notes, or caption support them; do not infer them + merely from table position. + """ + + OUTCOME = "outcome" + PREDICTOR = "predictor" + INSTRUMENTAL = "instrumental" + CONTROL = "control" + + +class AxisRole(str, Enum): + """Identify a variable's ordinary Cartesian axis in a figure. + + `x_axis` means the variable is encoded on the horizontal axis. `y_axis` + means it is encoded on the vertical axis. Use `MultiAxisAssignment` instead + when a graph contains distinct axes of the same dimension and the side or + outward position is needed to distinguish them. + """ + + X_AXIS = "x_axis" + Y_AXIS = "y_axis" + + +class AxisDimension(str, Enum): + """Identify a Cartesian axis dimension.""" + + X = "x" + Y = "y" + + +class AxisPosition(str, Enum): + """Identify the side of a plot where an axis appears.""" + + TOP = "top" + BOTTOM = "bottom" + LEFT = "left" + RIGHT = "right" + + +# Developer note: This application-owned relation vocabulary uses the +# open-start and open-end semantics described by EDTF for its open-ended value. +# Reference: https://www.loc.gov/standards/datetime/ +class TemporalRelation(str, Enum): + """Describe how normalized temporal bounds form an expression. + + `point` represents a single time point and requires `start` only. + `interval` represents a bounded period and requires `start` and `end`. + `open_ended_interval` represents a period for which the snapshot + explicitly supports only one temporal boundary. + """ + + POINT = "point" + INTERVAL = "interval" + OPEN_ENDED_INTERVAL = "open_ended_interval" + + +# Developer note: These application-owned precision levels correspond to the +# ISO 8601 forms selected by the normalization profile. +# Reference: https://www.iso.org/iso-8601-date-and-time-format.html +class TemporalPrecision(str, Enum): + """Describe the precision of normalized temporal bounds.""" + + YEAR = "year" + MONTH = "month" + DAY = "day" + DATETIME = "datetime" + + +# Developer note: This application vocabulary combines SDMX statistical +# operations, DDI summary-statistic types, quantile forms, and forms retained +# from schema evidence. It is not a verbatim copy of either external vocabulary. +# References: +# - SDMX CL_STATISTICAL_OPERATION: +# https://sdmx.org/wp-content/uploads/CL_STATISTICAL_OPERATION-1.0.docx +# - DDI Summary Statistic Type 2.1: +# https://rdf-vocabulary.ddialliance.org/ddi-cv/SummaryStatisticType/2.1.0/SummaryStatisticType.html +class StatisticalFormValue(str, Enum): + """Enumerate approved normalized statistical forms.""" + + OBSERVED_VALUE = "observed_value" + COUNT = "count" + ARITHMETIC_MEAN = "arithmetic_mean" + GEOMETRIC_MEAN = "geometric_mean" + WEIGHTED_MEAN = "weighted_mean" + WEIGHTED_VALUE = "weighted_value" + MEDIAN = "median" + SUM = "sum" + WEIGHTED_SUM = "weighted_sum" + MINIMUM = "minimum" + MAXIMUM = "maximum" + RANGE = "range" + INTERQUARTILE_RANGE = "interquartile_range" + VARIANCE = "variance" + STANDARD_DEVIATION = "standard_deviation" + STANDARD_ERROR = "standard_error" + MEAN_ABSOLUTE_DEVIATION = "mean_absolute_deviation" + MEDIAN_ABSOLUTE_DEVIATION = "median_absolute_deviation" + SKEWNESS = "skewness" + KURTOSIS = "kurtosis" + LOWER_BOUND = "lower_bound" + UPPER_BOUND = "upper_bound" + SCORE = "score" + RANK = "rank" + INDEX = "index" + RATE = "rate" + HARMONIC_MEAN = "harmonic_mean" + TRIMMED_MEAN = "trimmed_mean" + MODE = "mode" + COEFFICIENT_OF_VARIATION = "coefficient_of_variation" + VALID_CASES = "valid_cases" + INVALID_CASES = "invalid_cases" + PERCENTAGE_OF_VALID_CASES = "percentage_of_valid_cases" + PERCENTAGE_OF_INVALID_CASES = "percentage_of_invalid_cases" + QUARTILE = "quartile" + QUINTILE = "quintile" + DECILE = "decile" + PERCENTILE = "percentile" + RATIO = "ratio" + PROPORTION = "proportion" + PERCENTAGE = "percentage" + CONFIDENCE_INTERVAL = "confidence_interval" + COEFFICIENT = "coefficient" + P_VALUE = "p_value" + CHANGE = "change" + + +# Developer note: This application-owned vocabulary is grounded where possible +# in Vega-Lite marks and extended with visualization types observed in schema +# evidence. +# Reference: https://vega.github.io/vega-lite/docs/mark.html +class VisualizationTypeValue(str, Enum): + """Enumerate approved normalized visualization types.""" + + TABLE = "table" + BAR_CHART = "bar_chart" + LINE_CHART = "line_chart" + AREA_CHART = "area_chart" + PIE_CHART = "pie_chart" + DONUT_CHART = "donut_chart" + SCATTER_PLOT = "scatter_plot" + BUBBLE_CHART = "bubble_chart" + DOT_PLOT = "dot_plot" + HISTOGRAM = "histogram" + DENSITY_PLOT = "density_plot" + HEATMAP = "heatmap" + BOX_PLOT = "box_plot" + VIOLIN_PLOT = "violin_plot" + ERROR_BAR = "error_bar" + ERROR_BAND = "error_band" + RADAR_CHART = "radar_chart" + TREEMAP = "treemap" + MAP = "map" + CHOROPLETH_MAP = "choropleth_map" + SYMBOL_MAP = "symbol_map" + NETWORK_DIAGRAM = "network_diagram" + FLOW_DIAGRAM = "flow_diagram" + TIMELINE = "timeline" + DIAGRAM = "diagram" + INFOGRAPHIC = "infographic" + DASHBOARD = "dashboard" + COMPOSITE_FIGURE = "composite_figure" + + +# Developer note: This interval vocabulary aligns regular frequencies with +# SDMX CL_FREQ. +# Reference: https://sdmx.org/sdmx_cdcl/ +class ReportingIntervalValue(str, Enum): + """Enumerate approved normalized reporting intervals.""" + + HOURLY = "hourly" + DAILY = "daily" + WEEKLY = "weekly" + MONTHLY = "monthly" + QUARTERLY = "quarterly" + SEMIANNUAL = "semiannual" + ANNUAL = "annual" + + +# Developer note: This application-owned vocabulary implements the geographic +# level design informed by DDI geographic levels and ISO 3166-2 subdivisions. +# References: +# - DDI Lifecycle 3.3 GeographicLevel: +# https://docs.ddialliance.org/DDI-Lifecycle/3.3/ +# - ISO 3166: +# https://www.iso.org/iso-3166-country-codes.html +class GeographicLevelValue(str, Enum): + """Enumerate approved normalized geographic levels.""" + + GLOBAL = "global" + WORLD_REGION = "world_region" + COUNTRY = "country" + ADMINISTRATIVE_AREA_1 = "administrative_area_1" + ADMINISTRATIVE_AREA_2 = "administrative_area_2" + ADMINISTRATIVE_AREA_3 = "administrative_area_3" + LOCALITY = "locality" + SITE = "site" + + +class Identifier(_SchemaModel): + """Represent a verified identifier and its optional authority context. + + Parameters + ---------- + value : str + Identifier exactly as assigned in source evidence or trusted metadata. + scheme : str | None + Verified identifier scheme, when available. + issuer : str | None + Verified issuing agent, when available. + uri : AnyUrl | None + Verified authoritative absolute URI for the identifier. + """ + + value: NonEmptyText = Field( + examples=["P171254", "P178944"], + description="Identifier exactly as assigned in the snapshot or trusted metadata. Do not infer or manufacture it from its apparent pattern or from model knowledge.", + ) + scheme: NonEmptyText | None = Field( + examples=["World Bank project ID"], + default=None, + description="Identifier scheme explicitly supplied by the snapshot, trusted metadata, or configured and verified enrichment. Do not infer it from the identifier's apparent pattern or from model knowledge.", + ) + issuer: NonEmptyText | None = Field( + examples=["World Bank"], + default=None, + description="Issuing agent explicitly supplied by the snapshot, trusted metadata, or configured and verified enrichment. Do not infer it from model knowledge.", + ) + uri: AbsoluteURI | None = Field( + examples=["https://example.org/projects/P171254"], + default=None, + description="Authoritative absolute URI supplied by the snapshot, trusted metadata, or configured and verified enrichment. Do not construct or infer it from model knowledge.", + ) + + +class CodedTerm(_SchemaModel): + """Represent source wording with optional exact vocabulary identifiers. + + Parameters + ---------- + source_text : str | None + Exact source-visible text or symbol that expresses the term. + code : str | None + Code in the named scheme. + scheme : str | None + Code-list or vocabulary identifier. + uri : AnyUrl | None + Authoritative URI for the represented concept. + """ + + model_config = ConfigDict( + json_schema_extra={ + **_content_schema("source_text", "code", "uri"), + "allOf": [ + {"if": _populated("code"), "then": _populated("scheme")}, + {"if": _populated("scheme"), "then": _populated("code")}, + ], + "x-validation-rules": [ + "At least one of source_text, code, or uri must be non-null.", + "Code and scheme must be non-null together.", + ], + } + ) + + source_text: NonEmptyText | None = Field( + examples=["Grant", "District"], + default=None, + description="Exact text or symbol visible in the snapshot that explicitly expresses the value represented by this object. Omit it when the value is inferred from visual form, structure, or context rather than transcribed.", + ) + code: NonEmptyText | None = Field( + examples=["110"], default=None, description="Code in the named scheme." + ) + scheme: NonEmptyText | None = Field( + examples=["IATI Finance Type"], + default=None, + description="Code-list or vocabulary identifier.", + ) + uri: AbsoluteURI | None = Field( + examples=["https://example.org/concepts/TERM-001"], + default=None, + description="Authoritative URI for the represented concept.", + ) + + @model_validator(mode="after") + def _validate_term(self) -> CodedTerm: + if not any((self.source_text, self.code, self.uri)): + raise ValueError("A coded term must contain source text, a code, or a URI.") + if (self.code is None) != (self.scheme is None): + raise ValueError("A coded-term code and scheme must be supplied together.") + return self + + +class _NormalizedTerm(_SchemaModel): + """Validate the shared content rule for normalized term models.""" + + model_config = ConfigDict( + json_schema_extra=_content_schema("source_text", "normalized_value") + ) + + source_text: NonEmptyText | None = None + normalized_value: NonEmptyText | None = None + + @model_validator(mode="after") + def _validate_term(self) -> _NormalizedTerm: + if self.source_text is None and self.normalized_value is None: + raise ValueError("A normalized term must contain source text or a value.") + return self + + +class StatisticalFormTerm(_NormalizedTerm): + """Represent a known or source-only statistical form. + + Parameters + ---------- + source_text : str | None + Exact source-visible expression of the statistical form. + normalized_value : StatisticalFormValue | None + Approved normalized statistical form. + """ + + source_text: NonEmptyText | None = Field( + default=None, + description="Exact text visible in the snapshot that explicitly expresses the mathematical or statistical form of the values, such as count, percent of total, growth rate, confidence interval, index, sum, or mean. When surrounding source wording supports a form, preserve that wording here. A `%` that specifies the values' measurement scale belongs in the variable's `unit`; a `%` that belongs only to an expression such as `95% CI` is part of the statistical-form wording and is not a percent unit.", + examples=[ + "Count", + "Percent of total", + "Growth rate", + "95% CI", + "Index", + "Average", + ], + ) + normalized_value: StatisticalFormValue | None = Field( + examples=[ + "count", + "percentage", + "rate", + "change", + "confidence_interval", + "arithmetic_mean", + ], + default=None, + description="Approved mathematical or statistical form of the values. Prefer explicit source terminology. Use `percentage` for a percentage, share, or proportion per hundred; `rate` for a quantity presented as a rate relative to a base, including incidence, growth, interest, or tariff rates; and `change` for a difference, movement, or return. Do not select a form solely because the unit is `%`.", + ) + + +class VisualizationTypeTerm(_NormalizedTerm): + """Represent a known or source-only visualization type. + + Parameters + ---------- + source_text : str | None + Exact source-visible name of the visualization type. + normalized_value : VisualizationTypeValue | None + Approved normalized visualization type. + """ + + source_text: NonEmptyText | None = Field( + default=None, + description="Exact text visible in the snapshot that explicitly names the visualization type.", + examples=[ + "Bar chart", + "Line chart", + "Table", + "Map", + "Heatmap", + "Composite figure: line charts and map", + ], + ) + normalized_value: VisualizationTypeValue | None = Field( + examples=["bar_chart", "table", "composite_figure"], + default=None, + description="Approved normalized visualization type.", + ) + + +class ReportingIntervalTerm(_NormalizedTerm): + """Represent a known or source-only reporting interval. + + Parameters + ---------- + source_text : str | None + Exact source-visible expression of the reporting interval. + normalized_value : ReportingIntervalValue | None + Approved normalized reporting interval. + """ + + source_text: NonEmptyText | None = Field( + default=None, + description="Exact text visible in the snapshot that explicitly states the reporting interval.", + examples=["Annual", "Monthly", "Quarterly", "Daily"], + ) + normalized_value: ReportingIntervalValue | None = Field( + examples=["annual", "monthly", "quarterly"], + default=None, + description="Approved normalized reporting interval.", + ) + + +class GeographicLevelTerm(_NormalizedTerm): + """Represent a known or source-only geographic reporting level. + + Parameters + ---------- + source_text : str | None + Exact source-visible expression of the geographic level. + normalized_value : GeographicLevelValue | None + Approved normalized geographic level. + """ + + source_text: NonEmptyText | None = Field( + default=None, + description="Exact text visible in the snapshot that explicitly states the geographic level.", + examples=["Country", "Province", "District", "Facility"], + ) + normalized_value: GeographicLevelValue | None = Field( + examples=["country", "administrative_area_1", "site"], + default=None, + description="Approved normalized geographic level.", + ) + + +class EntityReference(_SchemaModel): + """Represent a named project, organization, source, or component. + + Parameters + ---------- + name : str + Source-visible entity name. + identifiers : list[Identifier] | None + Assigned identifiers for the entity. + """ + + name: NonEmptyText = Field( + examples=[ + "World Development Indicators", + "DHS", + "UNHCR Registration Data", + "National Census", + "Map Design Unit", + ], + description="Source-visible entity name.", + ) + identifiers: list[Identifier] | None = Field( + examples=[[{"value": "ENTITY-001", "scheme": "https://example.org/entities"}]], + default=None, + min_length=1, + description="Assigned entity identifiers.", + ) + + +class Attribution(EntityReference): + """Represent a named agent and its optional explicit attribution role. + + Parameters + ---------- + name : str + Source-visible credited-agent name. + identifiers : list[Identifier] | None + Assigned identifiers for the agent. + role : CodedTerm | None + Open, source-grounded attribution role, when explicit. + """ + + name: NonEmptyText = Field( + description="Source-visible entity name.", + examples=["Map Design Unit", "National Statistics Office"], + ) + role: CodedTerm | None = Field( + examples=[{"source_text": "Map maker"}, {"source_text": "Producer"}], + default=None, + description="Explicit source-grounded agent role.", + ) + + +class Unit(_SchemaModel): + """Represent a displayed unit and its optional normalized qualifiers. + + Parameters + ---------- + source_text : str + Complete displayed unit or scale expression. + code : str | None + Exact UN/CEFACT Recommendation 20 code. + multiplier_exponent : int | None + SDMX unit-multiplier exponent. + """ + + source_text: NonEmptyText = Field( + examples=[ + "%", + "People", + "Kilometers", + "US$ millions", + "US$000", + "US$m/SDRm", + ], + description="Exact complete text or symbol visible in the snapshot that states the unit or scale in which the variable's quantitative values are expressed, such as `%`, people, kilometers, `US$ millions`, or `US$000`. Preserve the complete measurement expression, excluding the variable name and annotations that do not specify the values' scale. For example, the `%` in `95% CI` states a confidence level and is not a percent unit. When the expression includes one supported applicable currency, repeat it in `currency`; when it combines multiple currencies or no single currency is supported, preserve the expression here and leave `currency` unset. Record a supported magnitude in `multiplier_exponent`. This intentional repetition preserves source wording while keeping qualifiers queryable. A unit does not determine the statistical form: `%` may accompany a percentage, rate, change, coefficient, or no normalized form.", + ) + code: ( + Annotated[str, StringConstraints(strict=True, pattern=r"^[A-Z0-9]{1,3}$")] + | None + ) = Field( + examples=["P1", "KMT"], + default=None, + description="Exact UN/CEFACT Recommendation 20 common code.", + json_schema_extra=_code_list( + "UNECE", + "Recommendation 20", + "https://unece.org/trade/uncefact/cl-recommendations", + ), + ) + multiplier_exponent: Annotated[int, Field(strict=True)] | None = Field( + examples=[3, 6, 9], + default=None, + description="Base-10 SDMX unit-multiplier exponent supported by the displayed scale, such as `3` for thousand or `6` for million.", + json_schema_extra=_code_list( + "SDMX", + "CL_UNIT_MULT", + "https://registry.sdmx.org/items/codelist.html", + "1.1", + ), + ) + + +class Currency(_SchemaModel): + """Represent a displayed currency and optional ISO 4217 code. + + Parameters + ---------- + source_text : str + Displayed currency expression. + code : str | None + Uppercase ISO 4217 alphabetic code. + """ + + source_text: NonEmptyText = Field( + examples=["US$", "$", "EUR", "JPY"], + description="Exact currency name, abbreviation, or symbol visible in the snapshot. When one applicable currency is supported, record it even when the same token also appears inside the complete `unit.source_text`; the repetition is intentional. If the unit expression combines multiple currencies or an ambiguous symbol cannot be resolved from snapshot evidence, preserve the complete unit expression and leave `currency` unset rather than forcing one value. Do not include noncurrency scale wording such as thousand or million.", + ) + code: ( + Annotated[str, StringConstraints(strict=True, pattern=r"^[A-Z]{3}$")] | None + ) = Field( + examples=["USD", "EUR", "JPY"], + default=None, + description="Uppercase ISO 4217 alphabetic code. Populate only when snapshot evidence deterministically identifies the currency; do not assume `USD` from an ambiguous `$` symbol alone.", + json_schema_extra=_code_list( + "ISO", "ISO 4217", "https://www.iso.org/iso-4217-currency-codes.html" + ), + ) + + +class Language(_SchemaModel): + """Represent a snapshot language using source text and/or BCP 47. + + Parameters + ---------- + source_text : str | None + Displayed language label, when present. + tag : str | None + Canonical BCP 47 language tag. + """ + + model_config = ConfigDict(json_schema_extra=_content_schema("source_text", "tag")) + + source_text: NonEmptyText | None = Field( + examples=["English", "French", "Arabic"], + default=None, + description="Exact text visible in the snapshot that explicitly names the language. Omit it when the language is inferred from the snapshot content.", + ) + tag: ( + Annotated[ + NonEmptyText, + WithJsonSchema( + {"type": "string", "minLength": 1, "pattern": _BCP47_PATTERN.pattern} + ), + ] + | None + ) = Field( + examples=["en", "fr", "ar"], + default=None, + description="Canonical BCP 47 language tag.", + json_schema_extra=_code_list( + "IETF/IANA", + "BCP 47", + "https://www.iana.org/assignments/language-subtag-registry/", + ), + ) + + @field_validator("tag") + @classmethod + def _validate_tag(cls, value: str | None) -> str | None: + if value is None: + return None + if not _BCP47_PATTERN.fullmatch(value): + raise ValueError("Language tag must use BCP 47 syntax.") + if value in _GRANDFATHERED_TAGS: + return value + parts = value.split("-") + variants: set[str] = set() + singletons: set[str] = set() + in_extension = False + for part in [] if parts[0].lower() == "x" else parts[1:]: + lower = part.lower() + if lower == "x": + break + if len(part) == 1: + if lower in singletons: + raise ValueError( + "Language tag cannot repeat an extension singleton." + ) + singletons.add(lower) + in_extension = True + elif not in_extension and ( + len(part) >= 5 or (len(part) == 4 and part[0].isdigit()) + ): + if lower in variants: + raise ValueError("Language tag cannot repeat a variant.") + variants.add(lower) + if parts[0].lower() == "x": + canonical = [part.lower() for part in parts] + if value != "-".join(canonical): + raise ValueError("Language tag must use canonical BCP 47 casing.") + return value + canonical = [parts[0].lower()] + extension = False + for part in parts[1:]: + if len(part) == 1: + extension = True + canonical.append(part.lower()) + elif extension: + canonical.append(part.lower()) + elif len(part) == 4 and part.isalpha(): + canonical.append(part.title()) + elif (len(part) == 2 and part.isalpha()) or ( + len(part) == 3 and part.isdigit() + ): + canonical.append(part.upper()) + else: + canonical.append(part.lower()) + if value != "-".join(canonical): + raise ValueError("Language tag must use canonical BCP 47 casing.") + return value + + @model_validator(mode="after") + def _validate_language(self) -> Language: + if self.source_text is None and self.tag is None: + raise ValueError("A language requires source_text or tag.") + return self + + +class MultiAxisAssignment(_SchemaModel): + """Bind a variable to one distinct axis in a multi-axis graph. + + Parameters + ---------- + dimension : AxisDimension + Cartesian dimension of the axis. + position : AxisPosition + Side of the plot where the axis appears. + position_index : int + One-based order from the plotting area outward on that side. + """ + + model_config = ConfigDict( + json_schema_extra={ + "allOf": [ + { + "if": { + "properties": {"dimension": {"const": "x"}}, + "required": ["dimension"], + }, + "then": {"properties": {"position": {"enum": ["top", "bottom"]}}}, + }, + { + "if": { + "properties": {"dimension": {"const": "y"}}, + "required": ["dimension"], + }, + "then": {"properties": {"position": {"enum": ["left", "right"]}}}, + }, + ], + "x-validation-rules": [ + "x axes use top or bottom; y axes use left or right." + ], + } + ) + + dimension: AxisDimension = Field( + examples=["x", "y"], + description="Cartesian dimension of the assigned axis.", + ) + position: AxisPosition = Field( + examples=["top", "bottom", "left", "right"], + description="Side of the plot where the assigned axis appears.", + ) + position_index: PositiveInt = Field( + examples=[1, 2], + description="One-based order from the plotting area outward among axes on the same side.", + ) + + @model_validator(mode="after") + def _validate_position(self) -> MultiAxisAssignment: + positions = { + AxisDimension.X: {AxisPosition.TOP, AxisPosition.BOTTOM}, + AxisDimension.Y: {AxisPosition.LEFT, AxisPosition.RIGHT}, + } + if self.position not in positions[self.dimension]: + raise ValueError( + f"{self.dimension.value}-axis position must be one of " + f"{sorted(position.value for position in positions[self.dimension])}." + ) + return self + + +class Variable(_SchemaModel): + """Represent a measured variable and its applicable qualifiers. + + Parameters + ---------- + name : str | None + Variable, indicator, metric, or measured concept, when identifiable. + unit : Unit | None + Applicable unit. + currency : Currency | None + Applicable currency. + analytical_roles : list[AnalyticalRole] | None + Explicit roles in a statistical analysis. + axis_roles : list[AxisRole] | None + Ordinary Cartesian axis roles in a figure. + multi_axis_assignments : list[MultiAxisAssignment] | None + Explicit assignments to distinct axes in a multi-axis graph. + statistical_forms : list[StatisticalFormTerm] | None + Applicable statistical forms. + """ + + model_config = ConfigDict(json_schema_extra=_variable_schema) + + name: NonEmptyText | None = Field( + examples=["GDP Growth", "Inflation", "Literacy Rate", "Refugee Population"], + default=None, + description="The measured quantity, indicator, metric, or outcome represented by the values. Use a variable for what is measured; use a dimension for categories that classify, group, or organize observations. For example, in a table of total loans by client sector, `Total loans` is a variable while `Corporates`, `Households`, and `Other clients` are categories of a client-sector dimension. Do not use `%`, `Value`, `Unknown`, or another unit or placeholder as the name.", + json_schema_extra=_standards(("https://schema.org/variableMeasured", "close")), + ) + unit: Unit | None = Field( + examples=[ + {"source_text": "Percent", "code": "P1"}, + {"source_text": "USD"}, + {"source_text": "People"}, + {"source_text": "Kilometers", "code": "KMT"}, + ], + default=None, + description="The complete displayed unit or scale expression for the variable, such as `%`, people, kilometers, `US$ millions`, or `US$000`. Preserve the complete measurement expression, excluding the variable name and annotations that do not specify the values' scale. When one applicable currency is supported, repeat it in `currency`; when the expression combines multiple currencies or no single currency is supported, leave `currency` unset. Record a supported magnitude in `unit.multiplier_exponent`. The unit does not by itself establish whether values are a percentage, rate, change, coefficient, or another statistical form; the `%` in `95% CI`, for example, is not a percent unit.", + json_schema_extra=_standards( + ("https://schema.org/unitCode", "related_structural") + ), + ) + currency: Currency | None = Field( + examples=[ + {"source_text": "USD", "code": "USD"}, + {"source_text": "EUR", "code": "EUR"}, + {"source_text": "JPY", "code": "JPY"}, + ], + default=None, + description="The currency denomination used for monetary values.", + json_schema_extra=_standards(("https://schema.org/currency", "close")), + ) + analytical_roles: list[AnalyticalRole] | None = Field( + examples=[["outcome"], ["predictor"], ["control"], ["instrumental"]], + default=None, + min_length=1, + description="Explicit roles in a statistical analysis or analytical table. Use `outcome` for the dependent or response variable, `predictor` for an explanatory or independent variable, `control` for an adjustment covariate, and `instrumental` only for a variable explicitly identified as an instrument. Do not derive these roles merely from row, column, or axis position.", + ) + axis_roles: list[AxisRole] | None = Field( + examples=[["x_axis"], ["y_axis"]], + default=None, + min_length=1, + description="Ordinary Cartesian axis roles in a figure. Use `x_axis` for a variable encoded on the horizontal axis and `y_axis` for one encoded on the vertical axis. Do not use this field for analytical-table roles or table row and column organization.", + ) + multi_axis_assignments: list[MultiAxisAssignment] | None = Field( + examples=[ + [ + { + "dimension": "y", + "position": "left", + "position_index": 1, + } + ], + [ + { + "dimension": "y", + "position": "right", + "position_index": 1, + } + ], + ], + default=None, + min_length=1, + description="Assignments to distinct Cartesian axes in a multi-axis graph. Use this field only when variables are assigned to different axes of the same dimension and the side or outward position distinguishes them. Use `axis_roles` for an ordinary single or shared x- or y-axis. `position_index` is 1 for the axis nearest the plotting area on a given side and increases outward.", + ) + statistical_forms: list[StatisticalFormTerm] | None = Field( + examples=[ + [{"source_text": "Count", "normalized_value": "count"}], + [ + { + "source_text": "Percentage", + "normalized_value": "percentage", + } + ], + [{"source_text": "Rate", "normalized_value": "rate"}], + [{"source_text": "Index", "normalized_value": "index"}], + [{"source_text": "Average"}], + ], + default=None, + min_length=1, + description="The mathematical or statistical form of the variable's values. Prefer explicit source terminology. Use `percentage` for a percentage, share, or proportion per hundred; `rate` for a quantity presented as a rate relative to a base, including incidence, growth, interest, or tariff rates; and `change` for a difference, movement, or return. When `%` specifies the scale in which the variable's values are expressed, preserve it in `unit`. Do not create a percent unit when the symbol belongs to a confidence level or significance annotation, or merely appears in a title or note without specifying the variable's measurement scale. Do not select a statistical form from `%` alone.", + json_schema_extra=_standards(("https://schema.org/statType", "close")), + ) + + @model_validator(mode="after") + def _validate_content(self) -> Variable: + if ( + self.name is None + and self.unit is None + and self.currency is None + and self.statistical_forms is None + ): + raise ValueError( + "A variable requires a name, unit, currency, or statistical forms." + ) + if self.name is None and ( + self.analytical_roles is not None + or self.axis_roles is not None + or self.multi_axis_assignments is not None + ): + raise ValueError( + "Analytical roles, axis roles, and multi-axis assignments require " + "a variable name." + ) + return self + + +class CategoryGroup(_SchemaModel): + """Represent one explicit nonrecursive category grouping. + + Parameters + ---------- + name : str + Explicit group heading. + categories : list[CodedTerm] + Categories directly contained by the group. + """ + + name: NonEmptyText = Field( + examples=["Violation of the right to liberty"], + description="Explicit category-group heading.", + ) + categories: list[CodedTerm] = Field( + examples=[ + [ + {"source_text": "Arbitrary arrests"}, + {"source_text": "Abductions"}, + ] + ], + min_length=1, + description="Categories directly contained by the group.", + ) + + +class Dimension(_SchemaModel): + """Represent a classificatory dimension and its visible organization. + + Parameters + ---------- + name : str | None + Dimension name, when identifiable. + categories : list[CodedTerm] | None + Ordered ungrouped categories. + category_groups : list[CategoryGroup] | None + One level of explicit category groups. + presentation_roles : list[PresentationRole] | None + Explicit row and/or column roles. + """ + + model_config = ConfigDict( + json_schema_extra=_content_schema("name", "categories", "category_groups") + ) + + name: NonEmptyText | None = Field( + examples=["Country", "Year", "Education Level", "Industry Sector", "Scenario"], + default=None, + description="The characteristic used to classify, group, or organize observations rather than the quantity being measured. Examples include year, country, sex, sector, and scenario. Repeated series or table headings that apply the same measure to different groups are normally dimension categories, not separate variables.", + ) + categories: list[CodedTerm] | None = Field( + examples=[ + [{"source_text": "Male"}, {"source_text": "Female"}], + [ + {"source_text": "Agriculture"}, + {"source_text": "Manufacturing"}, + {"source_text": "Services"}, + ], + [ + {"source_text": "Kenya"}, + {"source_text": "Uganda"}, + {"source_text": "Tanzania"}, + ], + [ + {"source_text": "Low"}, + {"source_text": "Medium"}, + {"source_text": "High"}, + ], + ], + default=None, + min_length=1, + description="The explicit labels or values belonging to the dimension, such as Male and Female for a sex dimension or Corporates and Households for a client-sector dimension.", + json_schema_extra=_standards( + ("http://www.w3.org/2004/02/skos/core#Concept", "related_structural") + ), + ) + category_groups: list[CategoryGroup] | None = Field( + examples=[ + [ + { + "name": "Violation of the right to liberty", + "categories": [ + {"source_text": "Arbitrary arrests"}, + {"source_text": "Abductions"}, + ], + } + ] + ], + default=None, + min_length=1, + description="One level of explicit category groups.", + json_schema_extra=_standards( + ("http://www.w3.org/2004/02/skos/core#broader", "related_structural") + ), + ) + presentation_roles: list[PresentationRole] | None = Field( + examples=[["row"], ["column"]], + default=None, + min_length=1, + description="How the dimension organizes a table. `row` means its categories vary downward through rows; `column` means they vary horizontally across columns. Assign the role from the direction in which the categories vary, not from the physical cell containing the dimension label. Do not use these roles for figure axes.", + ) + + @model_validator(mode="after") + def _validate_content(self) -> Dimension: + if ( + self.name is None + and self.categories is None + and self.category_groups is None + ): + raise ValueError( + "A dimension requires a name, categories, or category groups." + ) + return self + + +class TemporalExpression(_SchemaModel): + """Represent source-visible time with optional normalized bounds. + + Parameters + ---------- + source_text : str | None + Complete time expression explicitly visible in the snapshot. + start : str | None + Normalized starting value. + end : str | None + Normalized ending value. + relation : TemporalRelation | None + Relationship between the normalized bounds. + precision : TemporalPrecision | None + Precision shared by the normalized bounds. + """ + + model_config = ConfigDict(json_schema_extra=_temporal_schema) + + source_text: NonEmptyText | None = Field( + examples=["2015–2020", "FY2023", "January 2024"], + default=None, + description="Exact complete time expression visible in the snapshot. Omit it when normalized bounds are inferred from separate labels or structure rather than transcribed as one expression.", + ) + start: NonEmptyText | None = Field( + examples=["2015", "2024-01", "2024-01-01T12:00:00.1Z"], + default=None, + description="Normalized start.", + ) + end: NonEmptyText | None = Field( + examples=["2020", "2024-03"], default=None, description="Normalized end." + ) + relation: TemporalRelation | None = Field( + examples=["interval", "point", "open_ended_interval"], + default=None, + description="Relationship between normalized bounds. `point` represents a single represented-data time point and requires `start` only. `interval` represents a bounded period and requires `start` and `end`. `open_ended_interval` represents a period for which the snapshot explicitly supports only one temporal boundary. Use `start` for expressions such as “since 2015” and `end` for expressions such as “through 2020.” Do not use `open_ended_interval` merely because extraction failed to identify the other boundary.", + ) + precision: TemporalPrecision | None = Field( + examples=["year", "month", "day", "datetime"], + default=None, + description="Precision of normalized bounds.", + ) + + @model_validator(mode="after") + def _validate_bounds(self) -> TemporalExpression: + bounds = [bound for bound in (self.start, self.end) if bound is not None] + if self.source_text is None and not bounds: + raise ValueError("A temporal expression requires source text or a bound.") + if not bounds: + if self.relation is not None or self.precision is not None: + raise ValueError("Temporal relation and precision require a bound.") + return self + if self.relation is None or self.precision is None: + raise ValueError( + "Normalized temporal bounds require relation and precision." + ) + if self.relation is TemporalRelation.POINT: + if self.start is None or self.end is not None: + raise ValueError("point requires start only.") + elif self.relation is TemporalRelation.INTERVAL: + if self.start is None or self.end is None: + raise ValueError("interval requires start and end.") + elif len(bounds) != 1: + raise ValueError("open_ended_interval requires exactly one bound.") + for bound in bounds: + _validate_temporal_value(bound, self.precision) + if self.start is not None and self.end is not None: + if _temporal_sort_value(self.start, self.precision) > _temporal_sort_value( + self.end, self.precision + ): + raise ValueError("Temporal start cannot be after end.") + return self + + +def _validate_temporal_value(value: str, precision: TemporalPrecision) -> None: + if precision is TemporalPrecision.YEAR and _YEAR_PATTERN.fullmatch(value): + return + if precision is TemporalPrecision.MONTH and _MONTH_PATTERN.fullmatch(value): + return + if precision is TemporalPrecision.DAY and _DAY_PATTERN.fullmatch(value): + date.fromisoformat(value) + return + if precision is TemporalPrecision.DATETIME and _DATETIME_PATTERN.fullmatch(value): + parsed, _ = _parse_datetime(value) + if parsed.tzinfo is not None: + return + raise ValueError(f"Temporal value {value!r} does not match {precision.value}.") + + +def _parse_datetime(value: str) -> tuple[datetime, Decimal]: + # Parse whole seconds on Python 3.10; compare the original fraction exactly. + fraction = re.search(r"[.,]([0-9]+)(?=Z|[+-][0-9]{2}:[0-9]{2}$)", value) + if fraction: + seconds = Decimal("0." + fraction.group(1)) + value = value[: fraction.start()] + value[fraction.end() :] + else: + seconds = Decimal(0) + return datetime.fromisoformat(value.replace("Z", "+00:00")), seconds + + +def _temporal_sort_value( + value: str, precision: TemporalPrecision +) -> str | tuple[datetime, Decimal]: + if precision is TemporalPrecision.DATETIME: + return _parse_datetime(value) + return value + + +class TemporalCoverage(_SchemaModel): + """Group represented-data time and reporting interval. + + Parameters + ---------- + period : TemporalExpression | None + Represented-data temporal expression. + reporting_interval : ReportingIntervalTerm | None + Interval between successive represented-data time points. + """ + + model_config = ConfigDict( + json_schema_extra=_content_schema("period", "reporting_interval") + ) + + period: TemporalExpression | None = Field( + examples=[ + { + "source_text": "2015–2020", + "start": "2015", + "end": "2020", + "relation": "interval", + "precision": "year", + }, + {"source_text": "FY2023"}, + { + "source_text": "January 2024", + "start": "2024-01", + "relation": "point", + "precision": "month", + }, + ], + default=None, + description="The period or date range represented by the data.\n\nThis field describes **when the represented data apply**. It does not describe when the snapshot artifact or parent document was created, prepared, issued, published, revised, or retrieved. When an explicit artifact date appears only as part of a footer or provenance statement, preserve the complete statement in `interpretive_notes` rather than treating the date as `temporal_coverage.period`.", + json_schema_extra=_standards( + ("https://schema.org/temporalCoverage", "exact"), + ("http://purl.org/dc/terms/temporal", "close"), + ), + ) + reporting_interval: ReportingIntervalTerm | None = Field( + examples=[ + {"source_text": "Annual", "normalized_value": "annual"}, + {"source_text": "Monthly", "normalized_value": "monthly"}, + {"source_text": "Quarterly", "normalized_value": "quarterly"}, + {"source_text": "Daily", "normalized_value": "daily"}, + ], + default=None, + description="The interval between successive time points represented by the data, such as hourly, daily, monthly, quarterly, or annual. This describes the spacing of the represented data, not the document's publication schedule or the overall period covered.", + ) + + @model_validator(mode="after") + def _validate_content(self) -> TemporalCoverage: + if self.period is None and self.reporting_interval is None: + raise ValueError( + "Temporal coverage must contain period or reporting_interval." + ) + return self + + +class Place(_SchemaModel): + """Represent a source-grounded place with optional standard identifiers. + + Parameters + ---------- + source_text : str | None + Displayed place expression. + name : str | None + Preferred place name. + iso3_code : str | None + ISO 3166-1 alpha-3 country code. + subdivision_code : str | None + ISO 3166-2 subdivision code. + m49_code : str | None + UN M49 statistical-area code. + identifiers : list[Identifier] | None + Other authoritative identifiers. + """ + + model_config = ConfigDict( + json_schema_extra=_content_schema( + "source_text", + "name", + "iso3_code", + "subdivision_code", + "m49_code", + "identifiers", + ) + ) + + source_text: NonEmptyText | None = Field( + examples=["Global", "Kenya", "Sub-Saharan Africa", "Latin America"], + default=None, + description="Exact place expression visible in the snapshot. Preserve the displayed wording here; omit it when the place is inferred from structure or context rather than transcribed.", + ) + name: NonEmptyText | None = Field( + examples=["Kenya", "Sub-Saharan Africa", "Philippines"], + default=None, + description="Preferred normalized name for the place represented by `source_text` or by an authoritative identifier. This may standardize spelling or abbreviations and need not appear verbatim in the snapshot.", + ) + iso3_code: ( + Annotated[str, StringConstraints(strict=True, pattern=r"^[A-Z]{3}$")] | None + ) = Field( + examples=["KEN", "PHL"], + default=None, + description="ISO 3166-1 alpha-3 code for a country or area. Do not use World Bank aggregate or region codes.", + json_schema_extra=_code_list( + "ISO", "ISO 3166-1", "https://www.iso.org/iso-3166-country-codes.html" + ), + ) + subdivision_code: ( + Annotated[ + str, StringConstraints(strict=True, pattern=r"^[A-Z]{2}-[A-Z0-9]{1,3}$") + ] + | None + ) = Field( + examples=["US-CA"], + default=None, + description="ISO 3166-2 subdivision code.", + json_schema_extra=_code_list( + "ISO", "ISO 3166-2", "https://www.iso.org/iso-3166-country-codes.html" + ), + ) + m49_code: ( + Annotated[str, StringConstraints(strict=True, pattern=r"^[0-9]{3}$")] | None + ) = Field( + examples=["002", "202"], + default=None, + description="UN M49 statistical-area code.", + json_schema_extra=_code_list( + "United Nations", + "M49", + "https://unstats.un.org/unsd/methodology/m49/", + ), + ) + identifiers: list[Identifier] | None = Field( + examples=[[{"value": "KEN", "scheme": "ISO 3166-1 alpha-3"}]], + default=None, + min_length=1, + description="Other authoritative identifiers.", + ) + + @model_validator(mode="after") + def _validate_name(self) -> Place: + if not any( + ( + self.source_text, + self.name, + self.iso3_code, + self.subdivision_code, + self.m49_code, + self.identifiers, + ) + ): + raise ValueError( + "A place requires a name, source text, code, or identifier." + ) + return self + + +class GeographicLocation(Place): + """Represent an additional named location and its optional role and type. + + Parameters + ---------- + source_text : str | None + Displayed place expression. + name : str | None + Preferred place name. + iso3_code : str | None + ISO 3166-1 alpha-3 country code. + subdivision_code : str | None + ISO 3166-2 subdivision code. + m49_code : str | None + UN M49 statistical-area code. + identifiers : list[Identifier] | None + Other authoritative identifiers. + role : str | None + Explicit source-grounded geographic role. + type : CodedTerm | None + Physical or administrative type of the named location. + """ + + role: NonEmptyText | None = Field( + examples=[ + "Country of origin", + "Host country", + "Destination", + "Reporting location", + ], + default=None, + description="The relation of this location to the represented data, such as country of origin, host country, destination, or reporting location. Record only a relation supported by the snapshot; this field does not describe what kind of place it is or the level at which data are reported.", + ) + type: CodedTerm | None = Field( + examples=[ + {"source_text": "Refugee camp"}, + {"source_text": "Hospital"}, + {"source_text": "School"}, + {"source_text": "District"}, + ], + default=None, + description="The physical or administrative type of a named geographic location represented in the snapshot. This field describes what the location is; use `geographic_coverage.level` for the administrative or spatial level at which the snapshot's data are reported.", + ) + + +class GeographicCoverage(_SchemaModel): + """Group overall geographic scope, locations, and reporting level. + + Parameters + ---------- + scope : Place | None + Overall geographic coverage or focus. + locations : list[GeographicLocation] | None + Additional named locations. + level : GeographicLevelTerm | None + Administrative, geographic, or reporting level. + """ + + model_config = ConfigDict( + json_schema_extra=_content_schema("scope", "locations", "level") + ) + + scope: Place | None = Field( + examples=[ + {"source_text": "Global", "name": "Global", "m49_code": "001"}, + {"source_text": "Kenya", "name": "Kenya", "iso3_code": "KEN"}, + { + "source_text": "Sub-Saharan Africa", + "name": "Sub-Saharan Africa", + "m49_code": "202", + }, + {"source_text": "Latin America", "name": "Latin America"}, + ], + default=None, + description="The overall geographic coverage or focus of the snapshot. Use `locations` for additional places represented within that coverage. A scope may be a country, region, aggregate, or other area and does not by itself state the reporting level.", + json_schema_extra=_standards(("https://schema.org/spatialCoverage", "exact")), + ) + locations: list[GeographicLocation] | None = Field( + examples=[ + [{"name": "Uganda", "iso3_code": "UGA"}], + [{"name": "Nairobi"}], + [{"name": "West Africa", "m49_code": "011"}], + [{"name": "Burkina Faso", "iso3_code": "BFA"}], + ], + default=None, + min_length=1, + description="Additional named places represented within the snapshot. Record the overall coverage in `scope`, and use each location's `role` for its relation to the data. Assign identifiers to the place they identify; do not copy a parent area's code onto a child location.", + json_schema_extra=_standards( + ("https://schema.org/spatialCoverage", "related_structural") + ), + ) + level: GeographicLevelTerm | None = Field( + examples=[ + {"source_text": "Country", "normalized_value": "country"}, + {"source_text": "Province"}, + {"source_text": "District"}, + {"source_text": "Facility", "normalized_value": "site"}, + ], + default=None, + description="The geographic granularity at which observations are reported, such as country, province, district, or site. This describes the level of the data, not the overall coverage area, an individual place, or a location's semantic role.", + ) + + @model_validator(mode="after") + def _validate_content(self) -> GeographicCoverage: + if self.scope is None and self.locations is None and self.level is None: + raise ValueError("Geographic coverage must contain a value.") + return self + + +class Provenance(_SchemaModel): + """Separate derivation sources from credited agents. + + Parameters + ---------- + sources : list[EntityReference] | None + Entities from which represented data derive. + attributions : list[Attribution] | None + Agents explicitly credited for the snapshot artifact. + """ + + model_config = ConfigDict( + json_schema_extra=_content_schema("sources", "attributions") + ) + + sources: list[EntityReference] | None = Field( + examples=[ + [{"name": "World Development Indicators"}], + [{"name": "DHS"}], + [{"name": "UNHCR Registration Data"}], + [{"name": "National Census"}], + ], + default=None, + min_length=1, + description="Represented-data derivation sources.", + json_schema_extra=_standards( + ("http://www.w3.org/ns/prov#wasDerivedFrom", "related_structural"), + ("http://purl.org/dc/terms/source", "close"), + ), + ) + attributions: list[Attribution] | None = Field( + examples=[[{"name": "Map Design Unit", "role": {"source_text": "Map maker"}}]], + default=None, + min_length=1, + description="Agents explicitly credited for the snapshot artifact; include a role when it is explicit.", + json_schema_extra=_standards( + ("http://www.w3.org/ns/prov#wasAttributedTo", "related_structural") + ), + ) + + @model_validator(mode="after") + def _validate_content(self) -> Provenance: + if self.sources is None and self.attributions is None: + raise ValueError("Provenance must contain a source or attribution.") + return self + + +class Project(_SchemaModel): + """Represent project, program, operation, or initiative context. + + Parameters + ---------- + name : str | None + Associated project-context name. + identifiers : list[Identifier] | None + Formal project or operation identifiers. + components : list[EntityReference] | None + Explicitly identified subordinate components. + """ + + model_config = ConfigDict( + json_schema_extra=_content_schema("name", "identifiers", "components") + ) + + name: NonEmptyText | None = Field( + examples=[ + "Niger - COVID-19 Emergency Response Project", + "Jordan Health Sector Reform Project", + "Lebanon - Health Resilience Project", + ], + default=None, + description="The project, program, operation, or initiative associated with the snapshot.", + json_schema_extra=_standards(("https://schema.org/name", "exact")), + ) + identifiers: list[Identifier] | None = Field( + examples=[[{"value": "P171254"}], [{"value": "P178944"}]], + default=None, + min_length=1, + description="The formal identifier assigned to the associated project or operation.", + json_schema_extra=_standards( + ("https://schema.org/identifier", "standard_broader") + ), + ) + components: list[EntityReference] | None = Field( + examples=[ + [{"name": "Component 3: Project management"}], + [{"name": "Results Area 1"}], + ], + default=None, + min_length=1, + description="The project component, workstream, or results area represented by the snapshot.", + json_schema_extra=_standards( + ("https://schema.org/hasPart", "related_structural") + ), + ) + + @model_validator(mode="after") + def _validate_content(self) -> Project: + if self.name is None and self.identifiers is None and self.components is None: + raise ValueError("Project must contain a name, identifier, or component.") + return self + + +class Financing(_SchemaModel): + """Group project-financing measures, funders, and instruments. + + Parameters + ---------- + measures : list[str] | None + Financial quantities or funding-related measures. + funders : list[EntityReference] | None + Named funding sources. + instruments : list[CodedTerm] | None + Financing mechanisms. + """ + + model_config = ConfigDict( + json_schema_extra=_content_schema("measures", "funders", "instruments") + ) + + measures: list[NonEmptyText] | None = Field( + examples=[ + ["Project Cost"], + ["Disbursement"], + ["Financing Gap"], + ["Budget Allocation"], + ], + default=None, + min_length=1, + description="The financial quantity or funding-related measure represented by the snapshot.", + ) + funders: list[EntityReference] | None = Field( + examples=[ + [{"name": "IDA"}], + [{"name": "IBRD"}], + [{"name": "Government"}], + [{"name": "European Union"}], + ], + default=None, + min_length=1, + description="The organization or funding source providing financial support.", + json_schema_extra=_standards(("https://schema.org/funder", "exact")), + ) + instruments: list[CodedTerm] | None = Field( + examples=[ + [{"source_text": "Grant"}], + [{"source_text": "Loan"}], + [{"source_text": "Credit"}], + [{"source_text": "Trust Fund"}], + ], + default=None, + min_length=1, + description="The financing mechanism associated with the represented activity.", + json_schema_extra=_standards( + ( + "https://reference.iatistandard.org/en/iati-standard/203/codelists/financetype/", + "close", + ) + ), + ) + + @model_validator(mode="after") + def _validate_content(self) -> Financing: + if self.measures is None and self.funders is None and self.instruments is None: + raise ValueError("Financing must contain a measure, funder, or instrument.") + return self + + +class DataSnapshotMetadata(_SchemaModel): + """Represent canonical semantic metadata for one data snapshot. + + Parameters + ---------- + title : str | None + Primary title, caption, or heading. + document_label : str | None + Label assigned within the parent source document. + subject_domains : list[str] | None + Broad thematic, policy, or sectoral domains. + subject_summary : str | None + Concise analytical summary. + panel_titles : list[str] | None + Ordered explicit panel titles. + variables : list[Variable] | None + Measured concepts and their qualifiers. + dimensions : list[Dimension] | None + Classificatory dimensions and visible organization. + population_group : str | None + Human population represented by the data. + visualization_types : list[VisualizationTypeTerm] | None + Visible visualization forms. + temporal_coverage : TemporalCoverage | None + Represented-data time and reporting interval. + geographic_coverage : GeographicCoverage | None + Geographic scope, locations, and level. + comparisons : list[str] | None + Explicit comparisons or named comparators. + provenance : Provenance | None + Derivation sources and artifact attributions. + languages : list[Language] | None + Languages used within the snapshot. + interpretive_notes : list[str] | None + Complete source-visible interpretive statements. + project : Project | None + Associated project or operational context. + intervention_types : list[str] | None + Represented interventions or activities. + financing : Financing | None + Project-financing context. + analysis_methods : list[str] | None + Explicit analytical methods. + data_collection_methods : list[CodedTerm] | None + Explicit data-collection methods. + """ + + model_config = ConfigDict( + extra="forbid", + title="Data Snapshot Metadata Schema v1.4", + json_schema_extra={ + "$schema": "https://json-schema.org/draft/2020-12/schema", + "x-schema-version": "1.4", + "x-status": "implementation", + "x-validation-rules": [ + "Python is the canonical validator. JSON Schema enforces exported structural rules; format assertions require a format-aware validator.", + "Python additionally applies NFC normalization, outer whitespace trimming, and stable exact deduplication.", + "Python additionally checks URI syntax before URL normalization, language-tag casing and uniqueness, calendar dates, and chronological ordering.", + "Missing and null values are equivalent. Serialize records with exclude_none=True to omit unavailable values.", + "External registry membership, source-grounding, and semantic correctness are not validated. Unpinned code-list metadata identifies a syntax authority only.", + ], + }, + ) + + title: NonEmptyText | None = Field( + examples=[ + "Inflation Rate by Country", + "Annual Government Expenditure", + "Monthly labor income in Afghanistan and remittances from abroad", + "Table 6: Determinants of illegal land reallocation at village level", + ], + default=None, + description="The primary title, caption, or heading that identifies the data snapshot.", + json_schema_extra=_standards( + ("http://purl.org/dc/terms/title", "exact"), + ("https://schema.org/name", "exact"), + ), + ) + document_label: NonEmptyText | None = Field( + examples=["Figure 3", "Table 4.2", "Annex B", "Exhibit 7"], + default=None, + description="An exact visible label assigned to the snapshot within its source document, such as `Table 3`, `Figure 4.2`, or `Annex B`. Do not use a title, caption, page number, filename, or identifier inferred from document position.", + json_schema_extra=_standards( + ( + "https://jats.nlm.nih.gov/publishing/tag-library/1.3/element/label.html", + "close", + ), + ("https://schema.org/identifier", "standard_broader"), + ), + ) + subject_domains: list[NonEmptyText] | None = Field( + examples=[ + ["Education"], + ["Health"], + ["Macroeconomics"], + ["Agriculture"], + ["Forced Displacement"], + ], + default=None, + min_length=1, + description="The broad thematic, policy, or sectoral domain represented by the snapshot.", + json_schema_extra=_standards( + ("http://purl.org/dc/terms/subject", "standard_broader"), + ("https://schema.org/about", "standard_broader"), + ), + ) + subject_summary: NonEmptyText | None = Field( + examples=[ + "Trends in primary school enrollment", + "Distribution of humanitarian funding", + "Comparison of poverty rates across regions", + ], + default=None, + description="A concise summary describing the primary analytical subject or purpose of the snapshot.", + json_schema_extra=_standards( + ("https://schema.org/abstract", "close"), + ("http://purl.org/dc/terms/description", "standard_broader"), + ), + ) + panel_titles: list[NonEmptyText] | None = Field( + examples=[["(A) Poverty Rate"], ["(B) Literacy Rate"], ["Monthly Returns"]], + default=None, + min_length=1, + description="The title or heading of an individual panel within a multi-panel snapshot.\n\nPopulate only when panel titles are explicitly present.", + json_schema_extra=_standards( + ( + "https://jats.nlm.nih.gov/publishing/tag-library/1.3/element/fig-group.html", + "close", + ), + ("https://schema.org/hasPart", "related_structural"), + ), + ) + variables: list[Variable] | None = Field( + examples=[ + [{"name": "GDP Growth"}], + [{"name": "Inflation"}], + [{"name": "Literacy Rate"}], + [{"name": "Refugee Population"}], + [{"name": None, "unit": {"source_text": "%"}}], + [ + { + "name": "Share of households", + "unit": {"source_text": "%", "code": "P1"}, + "statistical_forms": [ + { + "source_text": "Share", + "normalized_value": "percentage", + } + ], + } + ], + [ + { + "name": "Budgeted cost", + "unit": {"source_text": "US$"}, + "currency": {"source_text": "US$", "code": "USD"}, + } + ], + [ + { + "name": "Monthly returns", + "unit": {"source_text": "%", "code": "P1"}, + "statistical_forms": [ + { + "source_text": "% change", + "normalized_value": "change", + } + ], + } + ], + [ + { + "name": "pce_residual", + "statistical_forms": [ + { + "source_text": "95% CI", + "normalized_value": "confidence_interval", + } + ], + } + ], + [ + { + "name": "Tariff rate", + "unit": {"source_text": "ad-valorem in %"}, + "statistical_forms": [ + { + "source_text": "Tariff rate", + "normalized_value": "rate", + } + ], + } + ], + [ + { + "name": "Amount", + "unit": { + "source_text": "US$m/SDRm", + "multiplier_exponent": 6, + }, + } + ], + [ + { + "name": "Revenue", + "unit": { + "source_text": "US$ millions", + "multiplier_exponent": 6, + }, + "currency": {"source_text": "US$", "code": "USD"}, + }, + { + "name": "Operating cost", + "unit": { + "source_text": "US$ millions", + "multiplier_exponent": 6, + }, + "currency": {"source_text": "US$", "code": "USD"}, + }, + ], + ], + default=None, + min_length=1, + description="Measured quantities, indicators, metrics, or outcomes and their applicable qualifiers. Use variables for what is measured and dimensions for the characteristics that classify or organize observations. Populate `name` when the measured concept can be identified; otherwise omit it rather than using a unit, `%`, `Value`, `Unknown`, or another placeholder. Retain an unnamed variable only when it has a unit, currency, or statistical form. Preserve each complete displayed unit or scale expression. When one applicable currency is supported, repeat it in `currency`; when the expression combines currencies or no single currency is supported, leave `currency` unset. Record a supported magnitude in `unit.multiplier_exponent`. Analytical roles, axis roles, and multi-axis assignments require a named variable. Repeat shared qualifiers on every variable to which they apply.", + json_schema_extra=_standards( + ("https://schema.org/variableMeasured", "close"), + ("https://ddialliance.org/Specification/DDI-Lifecycle/3.3/", "close"), + ), + ) + dimensions: list[Dimension] | None = Field( + examples=[ + [{"name": "Country"}], + [{"name": "Year"}], + [{"name": "Education Level"}], + [{"name": "Industry Sector"}], + [{"name": "Scenario"}], + [{"name": "Country", "presentation_roles": ["row"]}], + [{"name": "Indicator", "presentation_roles": ["row"]}], + [{"name": "Sector", "presentation_roles": ["row"]}], + [{"name": "Year", "presentation_roles": ["column"]}], + [{"name": "Region", "presentation_roles": ["column"]}], + [{"name": "Funding Source", "presentation_roles": ["column"]}], + ], + default=None, + min_length=1, + description="Characteristics that classify, group, or organize observations, together with their categories and table presentation. Use a dimension for year, geography, sex, sector, scenario, and similar groupings; use a variable for the quantity being measured.", + json_schema_extra=_standards( + ("https://sdmx.org/", "close"), + ("https://ddialliance.org/Specification/DDI-Lifecycle/3.3/", "close"), + ), + ) + population_group: NonEmptyText | None = Field( + examples=[ + "Refugees", + "Children under five", + "Female respondents", + "Host communities", + "Technical education graduates", + ], + default=None, + description="The human population, beneficiary group, or demographic group that is the primary subject of the represented data. This field describes who the data are about. Do not use it for companies, institutions, products, locations, or mixed entity classes, and do not repeat a population category used only to disaggregate observations.", + json_schema_extra=_standards( + ( + "https://docs.ddialliance.org/DDI-Lifecycle/3.3/xmlschema/schemas/conceptualcomponent_xsd/elements/Universe.html", + "close", + ), + ("https://schema.org/populationType", "close"), + ), + ) + visualization_types: list[VisualizationTypeTerm] | None = Field( + examples=[ + [{"normalized_value": "bar_chart"}], + [{"source_text": "Bar Graph", "normalized_value": "bar_chart"}], + [{"source_text": "Waffle chart"}], + None, + [{"source_text": "Bar chart", "normalized_value": "bar_chart"}], + [{"source_text": "Line chart", "normalized_value": "line_chart"}], + [{"source_text": "Table", "normalized_value": "table"}], + [{"source_text": "Map", "normalized_value": "map"}], + [{"source_text": "Heatmap", "normalized_value": "heatmap"}], + [ + {"normalized_value": "composite_figure"}, + {"normalized_value": "line_chart"}, + {"normalized_value": "map"}, + ], + [ + { + "source_text": "Composite figure: line charts and map", + "normalized_value": "composite_figure", + } + ], + ], + default=None, + min_length=1, + description="Visualization forms used to encode the represented data. Types are not mutually exclusive: for a composite or multi-panel artifact, record `composite_figure` and each identifiable component type, such as `line_chart` and `map`. Use `normalized_value` when a type is inferred from visual form. Use `source_text` only when wording in the snapshot explicitly names the form; do not manufacture source wording from the design. Preserve an explicitly written unfamiliar type in `source_text` without a normalized value. If neither a listed type nor an explicit unfamiliar label is supported, return `null`; do not force a match to the closest vocabulary value. Use `panel_titles` for explicit panel headings.", + json_schema_extra=_standards( + ("http://purl.org/dc/terms/type", "standard_broader"), + ("https://schema.org/additionalType", "standard_broader"), + ), + ) + temporal_coverage: TemporalCoverage | None = Field( + examples=[ + { + "period": { + "source_text": "2015–2020", + "start": "2015", + "end": "2020", + "relation": "interval", + "precision": "year", + } + }, + {"period": {"source_text": "FY2023"}}, + { + "period": { + "source_text": "January 2024", + "start": "2024-01", + "relation": "point", + "precision": "month", + } + }, + { + "reporting_interval": { + "source_text": "Annual", + "normalized_value": "annual", + } + }, + { + "reporting_interval": { + "source_text": "Monthly", + "normalized_value": "monthly", + } + }, + { + "reporting_interval": { + "source_text": "Quarterly", + "normalized_value": "quarterly", + } + }, + { + "reporting_interval": { + "source_text": "Daily", + "normalized_value": "daily", + } + }, + ], + default=None, + description="When the represented data apply and the interval between successive represented-data time points.", + json_schema_extra=_standards( + ("https://schema.org/temporalCoverage", "exact"), + ("http://purl.org/dc/terms/temporal", "close"), + ), + ) + geographic_coverage: GeographicCoverage | None = Field( + examples=[ + { + "scope": { + "source_text": "Global", + "name": "Global", + "m49_code": "001", + } + }, + { + "scope": { + "source_text": "Kenya", + "name": "Kenya", + "iso3_code": "KEN", + } + }, + { + "scope": { + "source_text": "Sub-Saharan Africa", + "name": "Sub-Saharan Africa", + "m49_code": "202", + } + }, + {"scope": {"source_text": "Latin America", "name": "Latin America"}}, + {"locations": [{"name": "Uganda", "iso3_code": "UGA"}]}, + {"locations": [{"name": "Nairobi"}]}, + {"locations": [{"name": "West Africa", "m49_code": "011"}]}, + {"locations": [{"name": "Burkina Faso", "iso3_code": "BFA"}]}, + {"level": {"source_text": "Country", "normalized_value": "country"}}, + {"level": {"source_text": "Province"}}, + {"level": {"source_text": "District"}}, + {"level": {"source_text": "Facility", "normalized_value": "site"}}, + ], + default=None, + description="The snapshot's overall geographic scope, additional named locations, and reporting level. Keep these meanings separate: `scope` is the coverage area, `locations` are specific places represented within it, and `level` is the granularity at which observations are reported.", + json_schema_extra=_standards( + ("https://schema.org/spatialCoverage", "exact"), + ("http://purl.org/dc/terms/spatial", "exact"), + ), + ) + comparisons: list[NonEmptyText] | None = Field( + examples=[ + ["Male vs Female"], + ["Rural vs Urban"], + ["Baseline vs Endline"], + ["Treatment vs Control"], + ["Before vs After"], + ["Low-income vs Middle-income vs High-income"], + ["Europe & Central Asia benchmark"], + ["Sub-Saharan Africa benchmark"], + ], + default=None, + min_length=1, + description="An explicit comparative relationship or named benchmark, comparator, reference group, cohort, scenario, or entity against which represented data are compared. Record the complete relation when visible, such as `Treatment vs Control`, or a named standalone benchmark. Do not use this field for totals, ordinary cross-tab categories, dimensions, or measures that merely appear together without an explicit comparison.", + ) + provenance: Provenance | None = Field( + examples=[ + {"sources": [{"name": "World Development Indicators"}]}, + {"sources": [{"name": "DHS"}]}, + {"sources": [{"name": "UNHCR Registration Data"}]}, + {"sources": [{"name": "National Census"}]}, + { + "attributions": [ + {"name": "Map Design Unit", "role": {"source_text": "Map maker"}} + ], + }, + ], + default=None, + description="The named dataset, survey, publication, organization, or credited agent from which the represented data originate or which is explicitly credited with producing the snapshot artifact.\n\nUse `sources` for represented-data derivation sources and `attributions` for credited agents; include an attribution role when it is explicit. Do not copy the parent document's authors or publisher into this field solely because they are associated with the document; the source or attribution must be explicitly relevant to the snapshot or its represented data.", + ) + languages: list[Language] | None = Field( + examples=[ + [{"source_text": "English", "tag": "en"}], + [{"source_text": "French", "tag": "fr"}], + [{"source_text": "Arabic", "tag": "ar"}], + ], + default=None, + min_length=1, + description="The language used within the snapshot.", + json_schema_extra=_standards( + ("https://schema.org/inLanguage", "exact"), + ("http://purl.org/dc/terms/language", "exact"), + ), + ) + interpretive_notes: list[NonEmptyText] | None = Field( + examples=[ + ["Values are provisional."], + ["Estimates exclude informal employment."], + ["Data collected using 2022 census boundaries."], + ["Sample: 1,204 respondents."], + ["Shaded areas show 95% confidence intervals."], + ["Prepared by the Map Design Unit, March 2024."], + ], + default=None, + min_length=1, + description="Explanatory, methodological, uncertainty, or provenance statements explicitly provided within the snapshot that aid interpretation or traceability.\n\nThis field may preserve complete notes containing sample-size statements, explanations of confidence intervals, standard errors or uncertainty bands, and footer statements that include an artifact date or production credit. It retains the statement as text; it does not create separate structured fields for sample size, uncertainty representation, or artifact publication date.\n\nPopulate only when such notes are explicitly present.", + json_schema_extra=_standards( + ("http://purl.org/dc/terms/description", "standard_broader"), + ("https://schema.org/description", "standard_broader"), + ), + ) + project: Project | None = Field( + examples=[ + {"name": "Niger - COVID-19 Emergency Response Project"}, + {"name": "Jordan Health Sector Reform Project"}, + {"name": "Lebanon - Health Resilience Project"}, + ], + default=None, + description="Associated project, program, operation, or initiative.", + json_schema_extra=_standards( + ("https://schema.org/Project", "standard_narrower"), + ( + "https://iatistandard.org/en/iati-standard/203/activity-standard/", + "close", + ), + ), + ) + intervention_types: list[NonEmptyText] | None = Field( + examples=[ + ["Cash Transfer"], + ["Vaccination"], + ["School Construction"], + ], + default=None, + min_length=1, + description="The intervention, service, policy, or operational activity represented.", + ) + financing: Financing | None = Field( + examples=[ + {"measures": ["Project Cost"]}, + {"measures": ["Disbursement"]}, + {"measures": ["Financing Gap"]}, + {"measures": ["Budget Allocation"]}, + {"funders": [{"name": "IDA"}]}, + {"funders": [{"name": "IBRD"}]}, + {"funders": [{"name": "Government"}]}, + {"funders": [{"name": "European Union"}]}, + {"instruments": [{"source_text": "Grant"}]}, + {"instruments": [{"source_text": "Loan"}]}, + {"instruments": [{"source_text": "Credit"}]}, + {"instruments": [{"source_text": "Trust Fund"}]}, + ], + default=None, + description="Project-financing measures, funders, and instruments.", + ) + analysis_methods: list[NonEmptyText] | None = Field( + examples=[ + ["Difference-in-Differences"], + ["Regression"], + ["Tobit model"], + ["Cost-Benefit Analysis"], + ], + default=None, + min_length=1, + description="The analytical, statistical, or computational method used to produce the reported results.\n\nPopulate only when explicitly stated.", + json_schema_extra=_standards( + ("https://schema.org/measurementTechnique", "related_structural"), + ("http://www.w3.org/ns/prov#Activity", "related_structural"), + ), + ) + data_collection_methods: list[CodedTerm] | None = Field( + examples=[ + [{"source_text": "Household Survey"}], + [{"source_text": "Administrative Records"}], + [{"source_text": "Key Informant Interviews"}], + [{"source_text": "Census"}], + ], + default=None, + min_length=1, + description="The method or instrument used to collect the underlying data.\n\nPopulate only when explicitly stated.", + json_schema_extra=_standards( + ( + "https://docs.ddialliance.org/DDI-Lifecycle/3.3/model/composite-types/ModeOfCollectionType/", + "exact", + ), + ("https://schema.org/measurementMethod", "close"), + ), + ) diff --git a/src/ai4data/data_snapshot/model_config.py b/src/ai4data/data_snapshot/model_config.py new file mode 100644 index 00000000..6b80e579 --- /dev/null +++ b/src/ai4data/data_snapshot/model_config.py @@ -0,0 +1,37 @@ +"""Resolve model settings from environment overrides and packaged JSON defaults.""" + +import json +import os +from functools import cache +from importlib.resources import files + + +@cache +def _load_model_defaults() -> dict[str, dict[str, str]]: + """Read the bundled model defaults once per process.""" + resource = files("ai4data.data_snapshot").joinpath("config", "models.json") + return json.loads(resource.read_text(encoding="utf-8")) + + +def get_model_setting(model: str, setting: str) -> str: + """Resolve one model setting, preferring a nonempty environment override. + + Parameters + ---------- + model : str + Extractor identifier used in the packaged configuration. + setting : str + Model location field, such as ``repo_id`` or ``cache_dir``. + + Returns + ------- + str + Environment value or the corresponding packaged JSON default. + + Raises + ------ + KeyError + If the model or setting has no packaged default and no override is set. + """ + variable = f"AI4DATA_SNAPSHOT_{model.upper()}_{setting.upper()}" + return os.getenv(variable) or _load_model_defaults()[model][setting] diff --git a/src/ai4data/data_snapshot/naming.py b/src/ai4data/data_snapshot/naming.py new file mode 100644 index 00000000..6c79008a --- /dev/null +++ b/src/ai4data/data_snapshot/naming.py @@ -0,0 +1,33 @@ +"""Build deterministic data snapshot output names.""" + +from pathlib import Path + + +def build_snapshot_stem(source_pdf: str | Path, label: str, index: int) -> str: + """Build a zero-based snapshot stem from its source document and label. + + Parameters + ---------- + source_pdf : str | Path + Source PDF path whose filename stem identifies the document. + label : str + Snapshot label, either ``"figure"`` or ``"table"``. + index : int + Zero-based per-label snapshot index for the source document. + + Returns + ------- + str + Stem in ``{source}_{figure|table}_{index:03d}`` form. + + Raises + ------ + ValueError + If the label is unsupported or the index is negative. + """ + normalized_label = label.lower() + if normalized_label not in {"figure", "table"}: + raise ValueError(f"Unsupported snapshot label: {label}") + if index < 0: + raise ValueError("Snapshot index must be zero or greater.") + return f"{Path(source_pdf).stem}_{normalized_label}_{index:03d}" diff --git a/src/ai4data/data_snapshot/pipeline.py b/src/ai4data/data_snapshot/pipeline.py new file mode 100644 index 00000000..3f82ad70 --- /dev/null +++ b/src/ai4data/data_snapshot/pipeline.py @@ -0,0 +1,166 @@ +"""Run snapshot extraction with one explicitly selected model.""" + +import re +from pathlib import Path + +from tqdm.auto import tqdm + +from ai4data.data_snapshot.adapters import create_adapter +from ai4data.data_snapshot.extraction import extract_snapshots, validate_unique_stems +from ai4data.data_snapshot.inputs import discover_pdfs +from ai4data.data_snapshot.metadata_extraction import extract_metadata_files + + +def run_extraction( + input_path: str | Path, + output_dir: str | Path, + model_name: str = "doclayoutyolo", + *, + dpi: int = 300, + force: bool = False, +) -> list[Path]: + """Extract snapshots beneath a directory named for the selected model. + + Parameters + ---------- + input_path : str | Path + One PDF file or a directory searched recursively for PDFs. + output_dir : str | Path + Root output directory. + model_name : str, optional + Identifier of the single snapshot extractor (default: doclayoutyolo). + dpi : int, optional + PyMuPDF rendering resolution in dots per inch. + force : bool, optional + Delete existing snapshots and metadata for each source PDF before + extracting it. + + Returns + ------- + list[Path] + Available paths beneath ``output_dir/model_name/snapshots``. + + Raises + ------ + ValueError + If the model identifier is not implemented. + """ + normalized_name = model_name.strip().lower() + if normalized_name not in {"doclayoutyolo", "tfid", "yolo11", "yolo26"}: + raise ValueError(f"Unsupported model: {model_name}") + pdf_paths = discover_pdfs(input_path) + validate_unique_stems(pdf_paths) + model_dir = Path(output_dir) / normalized_name + pending_pdf_paths = _prepare_pdf_paths(pdf_paths, model_dir, force=force) + if not pending_pdf_paths: + return [] + + detector = create_adapter(normalized_name) + output_paths: list[Path] = [] + pdf_progress = tqdm(pending_pdf_paths, desc="PDFs", unit="pdf") + for pdf_path in pdf_progress: + pdf_progress.set_postfix_str(pdf_path.name) + output_paths.extend(extract_snapshots(pdf_path, model_dir, detector, dpi=dpi)) + return output_paths + + +def run_pipeline( + input_path: str | Path, + output_dir: str | Path, + model_name: str = "doclayoutyolo", + *, + dpi: int = 300, + force: bool = False, + metadata_provider: str = "openrouter", +) -> list[Path]: + """Run snapshot extraction followed by semantic metadata extraction. + + Parameters + ---------- + input_path : str | Path + One PDF file or a directory searched recursively for PDFs. + output_dir : str | Path + Root output directory. + model_name : str, optional + Identifier of the single snapshot extractor (default: doclayoutyolo). + dpi : int, optional + PyMuPDF rendering resolution in dots per inch. + force : bool, optional + Replace outputs belonging to each source PDF. + metadata_provider : str, optional + API provider used for semantic metadata extraction. + + Returns + ------- + list[Path] + Snapshot paths followed by metadata paths. + + Raises + ------ + RuntimeError + If requested metadata extraction or validation fails. + ValueError + If the model identifier or input PDF set is invalid. + """ + normalized_name = model_name.strip().lower() + snapshot_paths = run_extraction( + input_path, + output_dir, + normalized_name, + dpi=dpi, + force=force, + ) + metadata_paths = extract_metadata_files( + snapshot_paths, + Path(output_dir) / normalized_name / "metadata", + metadata_provider=metadata_provider, + ) + return [*snapshot_paths, *metadata_paths] + + +def _prepare_pdf_paths( + pdf_paths: list[Path], model_dir: Path, *, force: bool +) -> list[Path]: + """Skip or clean existing outputs and return PDFs ready to process.""" + pending_pdf_paths: list[Path] = [] + for pdf_path in pdf_paths: + existing_paths = _existing_output_paths(model_dir, pdf_path.stem) + if not existing_paths: + pending_pdf_paths.append(pdf_path) + continue + if not force: + print( + f"Existing output files for {pdf_path.name} exist. " + "Skipped data snapshot and metadata extraction." + ) + continue + print( + f"Existing output files for {pdf_path.name} exist. " + "Cleaning up before continuing...", + end=" ", + flush=True, + ) + for existing_path in existing_paths: + existing_path.unlink() + print("Done!") + pending_pdf_paths.append(pdf_path) + return pending_pdf_paths + + +def _existing_output_paths(model_dir: Path, pdf_stem: str) -> list[Path]: + """Return exact crop and metadata outputs for one PDF stem.""" + pattern = re.compile(rf"{re.escape(pdf_stem)}_(?:figure|table)_\d{{3,}}") + output_paths: list[Path] = [] + for directory, suffix in ( + (model_dir / "snapshots", ".png"), + (model_dir / "metadata", ".json"), + ): + if directory.is_dir(): + output_paths.extend( + path + for path in directory.iterdir() + if path.is_file() + and path.suffix.lower() == suffix + and pattern.fullmatch(path.stem) + ) + return sorted(output_paths) diff --git a/src/ai4data/data_snapshot/rendering.py b/src/ai4data/data_snapshot/rendering.py new file mode 100644 index 00000000..934a3305 --- /dev/null +++ b/src/ai4data/data_snapshot/rendering.py @@ -0,0 +1,53 @@ +"""Render PDF pages for detection and cropping with PyMuPDF.""" + +from collections.abc import Iterator +from pathlib import Path + +import pymupdf +from PIL import Image +from tqdm.auto import tqdm + + +def iter_pdf_pages( + pdf_path: str | Path, dpi: int = 300 +) -> Iterator[tuple[int, Image.Image]]: + """Yield RGB page images rendered by PyMuPDF in document order. + + Parameters + ---------- + pdf_path : str | Path + PDF file to render. + dpi : int, optional + Rendering resolution in dots per inch. + + Yields + ------ + tuple[int, PIL.Image.Image] + Zero-based page index and rendered RGB page image. + + Raises + ------ + FileNotFoundError + If the PDF path does not identify an existing file. + ValueError + If the path is not a PDF or ``dpi`` is not positive. + """ + path = Path(pdf_path) + if not path.is_file(): + raise FileNotFoundError(f"PDF file does not exist: {path}") + if path.suffix.lower() != ".pdf": + raise ValueError(f"Input file is not a PDF: {path}") + if dpi <= 0: + raise ValueError("DPI must be greater than zero.") + + with pymupdf.open(path) as document: + pages = tqdm( + document, + total=document.page_count, + desc=f"Pages ({path.name})", + unit="page", + leave=False, + ) + for page_index, page in enumerate(pages): + pixmap = page.get_pixmap(dpi=dpi, alpha=False) + yield page_index, pixmap.pil_image().convert("RGB") diff --git a/src/ai4data/data_snapshot/submission.py b/src/ai4data/data_snapshot/submission.py new file mode 100644 index 00000000..bafa0cc5 --- /dev/null +++ b/src/ai4data/data_snapshot/submission.py @@ -0,0 +1,17 @@ +"""Submit extracted data snapshots and metadata.""" + +from pathlib import Path + + +def submit_snapshot(snapshot_path: str | Path, metadata_path: str | Path) -> None: + """Print the data snapshot and metadata paths that would be submitted. + + Parameters + ---------- + snapshot_path : str | Path + Path to the extracted data snapshot image. + metadata_path : str | Path + Path to the corresponding metadata JSON file. + """ + # Submission is not yet implemented; this mock only prints the paths. + print(f"Would submit {snapshot_path} with {metadata_path}") diff --git a/tests/conftest.py b/tests/conftest.py index c9b097bd..1a7e3e31 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -14,6 +14,8 @@ def _auto_mock_classifier(monkeypatch, mock_classifier_pipeline, request): return if "test_cookbook" in str(request.node.fspath): return # cookbook scripts are standalone and do not import ai4data + if "data_snapshot" in request.node.path.parts: + return # snapshot tests inject their own detectors and API clients from ai4data.data_use.models.model_manager import ModelManager monkeypatch.setattr( diff --git a/tests/data_snapshot/test_cli.py b/tests/data_snapshot/test_cli.py new file mode 100644 index 00000000..45eca059 --- /dev/null +++ b/tests/data_snapshot/test_cli.py @@ -0,0 +1,212 @@ +"""Test the pipeline command-line interfaces.""" + +from collections.abc import Callable +from pathlib import Path + +import pytest + +from ai4data.data_snapshot import cli + + +@pytest.mark.parametrize("command", [cli.main, cli.snapshot_main]) +@pytest.mark.parametrize("exported_repo", [None, "exported/repository"]) +def test_pdf_commands_load_env_before_creating_detector( + command: Callable[[list[str]], int], + exported_repo: str | None, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Load working-directory settings for both PDF commands without inference.""" + from ai4data.data_snapshot.adapters.doclayoutyolo import DocLayoutYOLOConfig + + monkeypatch.chdir(tmp_path) + key = "AI4DATA_SNAPSHOT_DOCLAYOUTYOLO_REPO_ID" + monkeypatch.delenv("PYTHON_DOTENV_DISABLED", raising=False) + monkeypatch.setenv(key, exported_repo or "temporary") + if exported_repo is None: + monkeypatch.delenv(key) + (tmp_path / ".env").write_text(f"{key}=file/repository\n", encoding="utf-8") + + def check_config(*args: object, **kwargs: object) -> list[Path]: + """Inspect the configuration before the mocked operation runs.""" + assert DocLayoutYOLOConfig().repo_id == (exported_repo or "file/repository") + return [] + + monkeypatch.setattr(cli, "run_pipeline", check_config) + monkeypatch.setattr(cli, "run_extraction", check_config) + assert command(["--input=example.pdf"]) == 0 + + +@pytest.mark.parametrize( + ("extra_args", "expected_force", "expected_provider"), + [ + ([], False, "openrouter"), + (["-f", "--metadata-provider=openai"], True, "openai"), + ], +) +def test_main_runs_pipeline_with_defaults_and_force( + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], + extra_args: list[str], + expected_force: bool, + expected_provider: str, +) -> None: + """Forward arguments and print paths returned by the pipeline.""" + calls: list[tuple[Path, Path, str, int, bool, str]] = [] + + def fake_run_pipeline( + input_path: Path, + output_dir: Path, + model_name: str, + *, + dpi: int, + force: bool, + metadata_provider: str, + ) -> list[Path]: + """Record the CLI values without running a model.""" + calls.append( + (input_path, output_dir, model_name, dpi, force, metadata_provider) + ) + return [Path("output/yolo11/snapshots/example_table_000.png")] + + monkeypatch.setattr(cli, "run_pipeline", fake_run_pipeline) + + result = cli.main(["--model=yolo11", "--input=data/example.pdf", *extra_args]) + + assert result == 0 + assert calls == [ + ( + Path("data/example.pdf"), + Path("output"), + "yolo11", + 300, + expected_force, + expected_provider, + ) + ] + assert capsys.readouterr().out == ( + "output/yolo11/snapshots/example_table_000.png\n" + ) + + +@pytest.mark.parametrize( + ("model_args", "expected_model"), + [([], "doclayoutyolo"), (["--model=tfid"], "tfid")], +) +def test_snapshot_main_runs_only_snapshot_extraction( + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], + model_args: list[str], + expected_model: str, +) -> None: + """Use the snapshot-only operation and print its output paths.""" + calls: list[str] = [] + + def fake_run_extraction(*args: object, **kwargs: object) -> list[Path]: + """Record that snapshot extraction was selected.""" + calls.append(str(args[2])) + return [Path(f"output/{args[2]}/snapshots/example_table_000.png")] + + monkeypatch.setattr(cli, "run_extraction", fake_run_extraction) + + result = cli.snapshot_main(["--input=data/example.pdf", *model_args]) + + assert result == 0 + assert calls == [expected_model] + assert capsys.readouterr().out.endswith("example_table_000.png\n") + + +def test_metadata_main_infers_output_and_forwards_force( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], +) -> None: + """Infer the metadata sibling and run only metadata extraction.""" + snapshot_path = tmp_path / "tfid" / "snapshots" / "report_table_000.png" + snapshot_path.parent.mkdir(parents=True) + snapshot_path.write_bytes(b"image") + calls: list[tuple[list[Path], Path, bool, str]] = [] + + def fake_extract_metadata_files( + snapshot_paths: list[Path], + output_dir: Path, + *, + force: bool, + metadata_provider: str, + ) -> list[Path]: + """Record metadata command values without making an API request.""" + paths = list(snapshot_paths) + calls.append((paths, output_dir, force, metadata_provider)) + return [output_dir / "report_table_000.json"] + + monkeypatch.setattr(cli, "extract_metadata_files", fake_extract_metadata_files) + + result = cli.metadata_main( + [ + f"--input={snapshot_path}", + "--force", + "--metadata-provider=openai", + ] + ) + + metadata_dir = tmp_path / "tfid" / "metadata" + assert result == 0 + assert calls == [([snapshot_path], metadata_dir, True, "openai")] + assert capsys.readouterr().out == f"{metadata_dir / 'report_table_000.json'}\n" + + +def test_metadata_main_requires_output_for_arbitrary_image( + tmp_path: Path, capsys: pytest.CaptureFixture[str] +) -> None: + """Require an explicit output directory outside a snapshots directory.""" + snapshot_path = tmp_path / "report.png" + snapshot_path.write_bytes(b"image") + + with pytest.raises(SystemExit) as error: + cli.metadata_main([f"--input={snapshot_path}"]) + + assert error.value.code == 2 + assert "--output-dir is required" in capsys.readouterr().err + + +def test_main_reports_pipeline_input_error( + monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str] +) -> None: + """Return a command-line error for an unsupported model.""" + + def reject_model(*args: object, **kwargs: object) -> list[Path]: + """Raise the pipeline's unsupported-model error.""" + raise ValueError("Unsupported model: unknown") + + monkeypatch.setattr(cli, "run_pipeline", reject_model) + + with pytest.raises(SystemExit) as error: + cli.main(["--model=unknown", "--input=data/example.pdf"]) + + assert error.value.code == 2 + assert "Unsupported model: unknown" in capsys.readouterr().err + + +def test_main_reports_runtime_error_without_usage( + monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str] +) -> None: + """Report a valid command's runtime failure without argparse help.""" + + def reject_request(*args: object, **kwargs: object) -> list[Path]: + """Raise a metadata runtime error.""" + raise RuntimeError("OpenAI credits are exhausted.") + + monkeypatch.setattr(cli, "run_pipeline", reject_request) + + result = cli.main(["--model=tfid", "--input=data/example.pdf"]) + + assert result == 1 + assert capsys.readouterr().err == "error: OpenAI credits are exhausted.\n" + + +def test_main_rejects_nonpositive_dpi() -> None: + """Reject a rendering resolution that cannot produce an image.""" + with pytest.raises(SystemExit) as error: + cli.main(["--model=yolo11", "--input=data/example.pdf", "--dpi=0"]) + + assert error.value.code == 2 diff --git a/tests/data_snapshot/test_detection.py b/tests/data_snapshot/test_detection.py new file mode 100644 index 00000000..7df73c83 --- /dev/null +++ b/tests/data_snapshot/test_detection.py @@ -0,0 +1,29 @@ +"""Test the canonical layout detection contract.""" + +import pytest + +from ai4data.data_snapshot import Detection + + +def test_detection_accepts_canonical_values() -> None: + """Retain a valid normalized figure detection.""" + detection = Detection("figure", (0.1, 0.2, 0.8, 0.9), 0.75) + + assert detection.label == "figure" + + +@pytest.mark.parametrize( + ("label", "bbox", "score"), + [ + ("caption", (0.1, 0.2, 0.8, 0.9), 0.75), + ("figure", (-0.1, 0.2, 0.8, 0.9), 0.75), + ("table", (0.8, 0.2, 0.1, 0.9), 0.75), + ("table", (0.1, 0.2, 0.8, 0.9), 1.1), + ], +) +def test_detection_rejects_invalid_values( + label: str, bbox: tuple[float, float, float, float], score: float +) -> None: + """Reject noncanonical labels, boxes, and scores.""" + with pytest.raises(ValueError): + Detection(label, bbox, score) # type: ignore[arg-type] diff --git a/tests/data_snapshot/test_doclayoutyolo.py b/tests/data_snapshot/test_doclayoutyolo.py new file mode 100644 index 00000000..54fdf479 --- /dev/null +++ b/tests/data_snapshot/test_doclayoutyolo.py @@ -0,0 +1,21 @@ +"""Test the DocLayout-YOLO adapter without loading model weights.""" + +from pathlib import Path + +from ai4data.data_snapshot.adapters.doclayoutyolo import ( + DocLayoutYOLOAdapter, + DocLayoutYOLOConfig, +) + + +def test_doclayoutyolo_adapter_uses_official_defaults_in_local_cache() -> None: + """Keep the official checkpoint and inference defaults in local cache.""" + model = object() + adapter = DocLayoutYOLOAdapter(model=model) + + assert adapter._model is model + assert adapter.config == DocLayoutYOLOConfig() + assert adapter.config.filename == "doclayout_yolo_docstructbench_imgsz1024.pt" + assert Path(adapter.config.cache_dir) == Path(".cache/models") + assert adapter.config.conf == 0.2 + assert adapter.config.imgsz == 1024 diff --git a/tests/data_snapshot/test_extraction.py b/tests/data_snapshot/test_extraction.py new file mode 100644 index 00000000..2c0cc9e1 --- /dev/null +++ b/tests/data_snapshot/test_extraction.py @@ -0,0 +1,107 @@ +"""Test snapshot detection and cropping orchestration.""" + +from pathlib import Path + +import pymupdf +import pytest +from PIL import Image + +from ai4data.data_snapshot import Detection, extract_snapshots + + +class _FakeDetector: + """Return fixed detections for extraction tests.""" + + def detect(self, image: Image.Image) -> list[Detection]: + """Return intentionally unordered figure and table detections.""" + return [ + Detection("figure", (0.5, 0.5, 1.0, 1.0), 0.8), + Detection("table", (0.0, 0.0, 0.5, 0.5), 0.9), + Detection("figure", (0.0, 0.0, 0.5, 0.5), 0.7), + ] + + +def _write_pdf(path: Path) -> None: + """Write one 72-point square PDF page for a test.""" + with pymupdf.open() as document: + document.new_page(width=72, height=72) + document.save(path) + + +def test_extract_snapshots_saves_spatially_ordered_zero_based_names( + tmp_path: Path, +) -> None: + """Crop from one render and assign independent label counters.""" + pdf_path = tmp_path / "report.pdf" + _write_pdf(pdf_path) + + paths = extract_snapshots(pdf_path, tmp_path / "output", _FakeDetector(), dpi=72) + + assert [path.name for path in paths] == [ + "report_table_000.png", + "report_figure_000.png", + "report_figure_001.png", + ] + assert [Image.open(path).size for path in paths] == [ + (36, 36), + (36, 36), + (36, 36), + ] + + +def test_extract_snapshots_skips_existing_pdf( + tmp_path: Path, capsys: pytest.CaptureFixture[str] +) -> None: + """Preserve existing snapshots and report a skipped source PDF.""" + pdf_path = tmp_path / "report.pdf" + _write_pdf(pdf_path) + existing = tmp_path / "output" / "snapshots" / "report_table_000.png" + existing.parent.mkdir(parents=True) + existing.write_bytes(b"keep") + + paths = extract_snapshots(pdf_path, tmp_path / "output", _FakeDetector(), dpi=72) + + assert paths == [existing] + assert existing.read_bytes() == b"keep" + assert f"Skipping {pdf_path}: existing snapshots found" in capsys.readouterr().out + + +def test_extract_snapshots_force_replaces_existing_pdf(tmp_path: Path) -> None: + """Remove all prior crops for one PDF before extracting it again.""" + pdf_path = tmp_path / "report.pdf" + _write_pdf(pdf_path) + snapshots_dir = tmp_path / "output" / "snapshots" + snapshots_dir.mkdir(parents=True) + existing = snapshots_dir / "report_table_000.png" + stale = snapshots_dir / "report_figure_999.png" + existing.write_bytes(b"old") + stale.write_bytes(b"stale") + + paths = extract_snapshots( + pdf_path, + tmp_path / "output", + _FakeDetector(), + dpi=72, + force=True, + ) + + assert [path.name for path in paths] == [ + "report_table_000.png", + "report_figure_000.png", + "report_figure_001.png", + ] + assert existing.read_bytes() != b"old" + assert not stale.exists() + + +def test_extract_snapshots_rejects_duplicate_document_stems(tmp_path: Path) -> None: + """Reject recursive inputs that would produce colliding output names.""" + first_dir = tmp_path / "first" + second_dir = tmp_path / "second" + first_dir.mkdir() + second_dir.mkdir() + _write_pdf(first_dir / "report.pdf") + _write_pdf(second_dir / "report.pdf") + + with pytest.raises(ValueError, match="filename stems must be unique"): + extract_snapshots(tmp_path, tmp_path / "output", _FakeDetector(), dpi=72) diff --git a/tests/data_snapshot/test_foundation.py b/tests/data_snapshot/test_foundation.py new file mode 100644 index 00000000..2f394063 --- /dev/null +++ b/tests/data_snapshot/test_foundation.py @@ -0,0 +1,110 @@ +"""Test PDF discovery, rendering, and snapshot naming contracts.""" + +from pathlib import Path + +import pymupdf +import pytest + +from ai4data.data_snapshot import ( + build_snapshot_stem, + discover_pdfs, + discover_snapshot_images, + iter_pdf_pages, +) + + +def test_discover_pdfs_accepts_file_or_directory(tmp_path: Path) -> None: + """Discover one file directly or PDFs recursively in path order.""" + first_pdf = tmp_path / "a.pdf" + second_pdf = tmp_path / "nested" / "b.PDF" + second_pdf.parent.mkdir() + first_pdf.touch() + second_pdf.touch() + (tmp_path / "notes.txt").touch() + + assert discover_pdfs(first_pdf) == [first_pdf] + assert discover_pdfs(tmp_path) == [first_pdf, second_pdf] + + +def test_discover_pdfs_rejects_invalid_inputs(tmp_path: Path) -> None: + """Reject missing, non-PDF, and empty-directory inputs.""" + text_path = tmp_path / "notes.txt" + text_path.touch() + empty_dir = tmp_path / "empty" + empty_dir.mkdir() + + with pytest.raises(FileNotFoundError, match="Input does not exist"): + discover_pdfs(tmp_path / "missing") + with pytest.raises(ValueError, match="not a PDF"): + discover_pdfs(text_path) + with pytest.raises(FileNotFoundError, match="No PDF files"): + discover_pdfs(empty_dir) + + +def test_discover_snapshot_images_accepts_file_or_directory(tmp_path: Path) -> None: + """Discover one image directly or supported images recursively in order.""" + first_image = tmp_path / "a.png" + second_image = tmp_path / "nested" / "b.JPG" + second_image.parent.mkdir() + first_image.touch() + second_image.touch() + (tmp_path / "notes.txt").touch() + + assert discover_snapshot_images(first_image) == [first_image] + assert discover_snapshot_images(tmp_path) == [first_image, second_image] + + +def test_discover_snapshot_images_rejects_invalid_inputs(tmp_path: Path) -> None: + """Reject missing, unsupported, and empty-directory image inputs.""" + text_path = tmp_path / "notes.txt" + text_path.touch() + empty_dir = tmp_path / "empty" + empty_dir.mkdir() + + with pytest.raises(FileNotFoundError, match="Input does not exist"): + discover_snapshot_images(tmp_path / "missing") + with pytest.raises(ValueError, match="not a supported image"): + discover_snapshot_images(text_path) + with pytest.raises(FileNotFoundError, match="No snapshot images"): + discover_snapshot_images(empty_dir) + + +def test_iter_pdf_pages_yields_rgb_pages(tmp_path: Path) -> None: + """Render pages once in document order using PyMuPDF.""" + pdf_path = tmp_path / "document.pdf" + with pymupdf.open() as document: + document.new_page(width=72, height=36) + document.new_page(width=144, height=72) + document.save(pdf_path) + + pages = list(iter_pdf_pages(pdf_path, dpi=72)) + + assert [(index, image.mode, image.size) for index, image in pages] == [ + (0, "RGB", (72, 36)), + (1, "RGB", (144, 72)), + ] + + +def test_iter_pdf_pages_rejects_invalid_dpi(tmp_path: Path) -> None: + """Reject a non-positive rendering resolution before opening the PDF.""" + pdf_path = tmp_path / "document.pdf" + pdf_path.touch() + + with pytest.raises(ValueError, match="DPI"): + list(iter_pdf_pages(pdf_path, dpi=0)) + + +def test_build_snapshot_stem_uses_zero_based_per_label_names() -> None: + """Normalize supported labels and preserve zero-based indices.""" + assert build_snapshot_stem("reports/source.pdf", "Figure", 0) == ( + "source_figure_000" + ) + assert build_snapshot_stem("reports/source.pdf", "table", 7) == ("source_table_007") + + +def test_build_snapshot_stem_rejects_invalid_values() -> None: + """Reject unsupported labels and negative indices.""" + with pytest.raises(ValueError, match="Unsupported snapshot label"): + build_snapshot_stem("source.pdf", "caption", 0) + with pytest.raises(ValueError, match="zero or greater"): + build_snapshot_stem("source.pdf", "figure", -1) diff --git a/tests/data_snapshot/test_installation.py b/tests/data_snapshot/test_installation.py new file mode 100644 index 00000000..60c326a3 --- /dev/null +++ b/tests/data_snapshot/test_installation.py @@ -0,0 +1,144 @@ +"""Check namespace imports and installed command entry points without models.""" + +import json +import subprocess +import sys +from importlib import import_module +from importlib.resources import files +from pathlib import Path + +import pytest + + +def test_package_import_does_not_load_snapshot_dependencies(tmp_path: Path) -> None: + """Keep base imports independent of snapshot and model runtimes.""" + script = """ +import importlib.abc +import sys + +class BlockOptionalImports(importlib.abc.MetaPathFinder): + def find_spec(self, fullname, path=None, target=None): + if fullname.split('.')[0] in { + 'PIL', 'pymupdf', 'openai', 'dotenv', 'torch', 'transformers', + 'ultralytics', 'doclayout_yolo', 'huggingface_hub', + }: + raise AssertionError(f'Unexpected optional import: {fullname}') + +sys.meta_path.insert(0, BlockOptionalImports()) +from ai4data import data_snapshot +assert data_snapshot.discover_pdfs.__module__ == 'ai4data.data_snapshot.inputs' +try: + data_snapshot.unknown +except AttributeError: + pass +else: + raise AssertionError('Unknown attribute was accepted') +""" + subprocess.run([sys.executable, "-c", script], cwd=tmp_path, check=True) + + +@pytest.mark.parametrize( + "module", + [ + "ai4data.data_snapshot", + "ai4data.data_snapshot.extract_data_snapshot", + "ai4data.data_snapshot.extract_metadata", + ], +) +def test_module_help_runs_outside_checkout(module: str, tmp_path: Path) -> None: + """Exercise each module entry point without inference or API calls.""" + result = subprocess.run( + [sys.executable, "-m", module, "--help"], + cwd=tmp_path, + capture_output=True, + text=True, + check=True, + ) + assert module in result.stdout + assert "--input" in result.stdout + + +def test_packaged_resources_are_available() -> None: + """Read both provider presets and prompt files through the installed package.""" + from importlib.resources import files + + from ai4data.data_snapshot.metadata_extraction import load_extraction_config + from ai4data.data_snapshot.metadata_extraction.extraction import ( + get_extraction_config_path, + ) + + for provider in ("openai", "openrouter"): + config = load_extraction_config(get_extraction_config_path(provider)) + assert config["provider"] == provider + resources = files("ai4data.data_snapshot.metadata_extraction") + for prompt in ("system.md", "user.md"): + assert resources.joinpath("prompts", prompt).read_text(encoding="utf-8") + model_defaults = json.loads( + files("ai4data.data_snapshot") + .joinpath("config", "models.json") + .read_text(encoding="utf-8") + ) + assert set(model_defaults) == {"doclayoutyolo", "yolo11", "yolo26", "tfid"} + assert all( + isinstance(value, str) and value.strip() + for settings in model_defaults.values() + for value in settings.values() + ) + + +@pytest.mark.parametrize( + ("module", "config_name"), + [ + ("doclayoutyolo", "DocLayoutYOLOConfig"), + ("yolo11", "YOLO11Config"), + ("yolo26", "YOLO26Config"), + ("tfid", "TFIDConfig"), + ], +) +def test_model_configuration_environment_precedence( + module: str, config_name: str, monkeypatch: pytest.MonkeyPatch +) -> None: + """Resolve per-model defaults at construction time, below explicit settings.""" + adapter_module = import_module(f"ai4data.data_snapshot.adapters.{module}") + config_type = getattr(adapter_module, config_name) + prefix = f"AI4DATA_SNAPSHOT_{module.upper()}_" + defaults = json.loads( + files("ai4data.data_snapshot") + .joinpath("config", "models.json") + .read_text(encoding="utf-8") + )[module] + defaults["cache_dir"] = Path(defaults["cache_dir"]) + for other in ("DOCLAYOUTYOLO", "YOLO11", "YOLO26", "TFID"): + for field in defaults: + monkeypatch.setenv(f"AI4DATA_SNAPSHOT_{other}_{field.upper()}", "other") + for field in defaults: + monkeypatch.delenv(prefix + field.upper()) + config = config_type() + assert {field: getattr(config, field) for field in defaults} == defaults + + for field in defaults: + monkeypatch.setenv(prefix + field.upper(), "") + assert config_type() == config + + from ai4data.data_snapshot import model_config + + alternate_defaults = {field: f"preset-{field}" for field in defaults} + monkeypatch.setattr( + model_config, "_load_model_defaults", lambda: {module: alternate_defaults} + ) + config = config_type() + assert { + field: str(getattr(config, field)) for field in defaults + } == alternate_defaults + + for value in ("first", "second"): + for field in defaults: + monkeypatch.setenv(prefix + field.upper(), f"{value}-{field}") + config = config_type() + assert {field: str(getattr(config, field)) for field in defaults} == { + field: f"{value}-{field}" for field in defaults + } + + explicit = {field: f"explicit-{field}" for field in defaults} + config = config_type(**explicit) + assert {field: getattr(config, field) for field in defaults} == explicit diff --git a/tests/data_snapshot/test_metadata_batch.py b/tests/data_snapshot/test_metadata_batch.py new file mode 100644 index 00000000..d1bf7010 --- /dev/null +++ b/tests/data_snapshot/test_metadata_batch.py @@ -0,0 +1,127 @@ +"""Test writing validated snapshot metadata files.""" + +import json +from pathlib import Path + +import pytest + +from ai4data.data_snapshot.metadata_extraction import ExtractionResult +from ai4data.data_snapshot.metadata_extraction import batch +from ai4data.data_snapshot.metadata_schema import DataSnapshotMetadata + + +def _result(title: str) -> ExtractionResult: + """Build a successful extraction result for a test.""" + return ExtractionResult( + metadata=DataSnapshotMetadata(title=title), + model="test-model", + response_id="resp_test", + api_status="completed", + incomplete_details=None, + elapsed_seconds=0.1, + usage=None, + raw_output=None, + error_type=None, + error=None, + ) + + +def test_extract_metadata_files_writes_validated_json( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """Write one canonical JSON file beside the model's snapshots directory.""" + snapshot_path = tmp_path / "snapshots" / "report_table_000.png" + snapshot_path.parent.mkdir() + snapshot_path.write_bytes(b"image") + monkeypatch.setattr( + batch, "extract_metadata", lambda path, **kwargs: _result("A title") + ) + + paths = batch.extract_metadata_files([snapshot_path], tmp_path / "metadata") + + assert paths == [tmp_path / "metadata" / "report_table_000.json"] + assert json.loads(paths[0].read_text(encoding="utf-8"))["title"] == "A title" + + +def test_extract_metadata_files_skips_existing_output( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], +) -> None: + """Preserve an existing JSON file without making an API request.""" + snapshot_path = tmp_path / "snapshots" / "report_table_000.png" + metadata_path = tmp_path / "metadata" / "report_table_000.json" + snapshot_path.parent.mkdir() + metadata_path.parent.mkdir() + snapshot_path.write_bytes(b"image") + metadata_path.write_text('{"keep": true}\n', encoding="utf-8") + + def fail_if_called(path: Path) -> ExtractionResult: + """Fail if the skipped snapshot reaches extraction.""" + raise AssertionError(f"Unexpected extraction call: {path}") + + monkeypatch.setattr(batch, "extract_metadata", fail_if_called) + + paths = batch.extract_metadata_files([snapshot_path], tmp_path / "metadata") + + assert paths == [metadata_path] + assert metadata_path.read_text(encoding="utf-8") == '{"keep": true}\n' + assert capsys.readouterr().out == ( + "Existing metadata output for report_table_000.png exists. " + "Skipped metadata extraction.\n" + ) + + +def test_force_replaces_metadata_without_modifying_snapshot( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], +) -> None: + """Replace current JSON while preserving the source snapshot.""" + snapshot_path = tmp_path / "snapshots" / "report_table_000.png" + metadata_dir = tmp_path / "metadata" + metadata_path = metadata_dir / "report_table_000.json" + snapshot_path.parent.mkdir() + metadata_dir.mkdir() + snapshot_path.write_bytes(b"image") + metadata_path.write_text('{"old": true}\n', encoding="utf-8") + monkeypatch.setattr( + batch, "extract_metadata", lambda path, **kwargs: _result("New title") + ) + + batch.extract_metadata_files( + [snapshot_path], + metadata_dir, + force=True, + ) + + assert json.loads(metadata_path.read_text(encoding="utf-8"))["title"] == "New title" + assert snapshot_path.read_bytes() == b"image" + assert "Cleaning up before continuing... Done!" in capsys.readouterr().out + + +def test_extraction_failure_writes_nothing( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """Raise a useful error without creating invalid metadata JSON.""" + snapshot_path = tmp_path / "snapshots" / "report_table_000.png" + snapshot_path.parent.mkdir() + snapshot_path.write_bytes(b"image") + failure = ExtractionResult( + metadata=None, + model="test-model", + response_id="resp_failed", + api_status="incomplete", + incomplete_details={"reason": "max_output_tokens"}, + elapsed_seconds=0.1, + usage=None, + raw_output="{", + error_type="ValueError", + error="Response was incomplete: max_output_tokens.", + ) + monkeypatch.setattr(batch, "extract_metadata", lambda path, **kwargs: failure) + + with pytest.raises(RuntimeError, match="max_output_tokens"): + batch.extract_metadata_files([snapshot_path], tmp_path / "metadata") + + assert not (tmp_path / "metadata").exists() diff --git a/tests/data_snapshot/test_metadata_extraction.py b/tests/data_snapshot/test_metadata_extraction.py new file mode 100644 index 00000000..0bdc42f0 --- /dev/null +++ b/tests/data_snapshot/test_metadata_extraction.py @@ -0,0 +1,256 @@ +"""Test Schema v1.4 single-snapshot metadata extraction offline.""" + +import json +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from ai4data.data_snapshot import DataSnapshotMetadata, extract_metadata +from ai4data.data_snapshot.metadata_extraction import load_extraction_config +from ai4data.data_snapshot.metadata_extraction.extraction import ( + _create_api_client, + _response_format, + get_extraction_config_path, +) + + +class _FakeResponses: + """Record Responses API requests and return a configured result.""" + + def __init__(self, response: object | Exception) -> None: + self.response = response + self.request: dict[str, object] | None = None + + def create(self, **arguments: object) -> object: + """Record a request and return or raise the configured result.""" + self.request = arguments + if isinstance(self.response, Exception): + raise self.response + return self.response + + +def _write_image(path: Path) -> None: + """Write a minimal file accepted as a PNG input.""" + path.write_bytes(b"png-data") + + +def _write_config(path: Path, **overrides: object) -> None: + """Write a minimal extraction configuration.""" + config: dict[str, object] = {"model": "test-model", "max_output_tokens": 25_000} + config.update(overrides) + path.write_text(json.dumps(config), encoding="utf-8") + + +def test_default_config_uses_openrouter_luna() -> None: + """Use GPT-5.6 Luna through OpenRouter by default.""" + config = load_extraction_config() + + assert config == { + "provider": "openrouter", + "model": "openai/gpt-5.6-luna", + "max_output_tokens": 25_000, + "timeout": 900.0, + "service_tier": "flex", + "reasoning": {"effort": "medium"}, + } + + +def test_openai_config_preserves_runtime_settings() -> None: + """Keep the direct OpenAI Luna request settings available.""" + config = load_extraction_config(get_extraction_config_path("openai")) + + assert config["provider"] == "openai" + assert config["model"] == "gpt-5.6-luna" + assert config["max_output_tokens"] == 25_000 + assert config["service_tier"] == "flex" + assert config["reasoning"] == {"effort": "medium"} + + +def test_openrouter_client_uses_compatible_endpoint( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Point the existing OpenAI SDK at OpenRouter with its own API key.""" + import openai + + arguments: dict[str, object] = {} + + def fake_openai(**kwargs: object) -> object: + """Record SDK client configuration.""" + arguments.update(kwargs) + return object() + + monkeypatch.setenv("OPENROUTER_API_KEY", "test-key") + monkeypatch.setattr(openai, "OpenAI", fake_openai) + + _create_api_client("openrouter") + + assert arguments == { + "api_key": "test-key", + "base_url": "https://openrouter.ai/api/v1", + "max_retries": 0, + } + + +@pytest.mark.parametrize("exported_key", [None, "exported-key"]) +def test_client_loads_working_directory_env_without_overriding_exports( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, exported_key: str | None +) -> None: + """Read local credentials after installation and preserve shell precedence.""" + import openai + + monkeypatch.chdir(tmp_path) + monkeypatch.delenv("OPENROUTER_API_KEY", raising=False) + monkeypatch.delenv("PYTHON_DOTENV_DISABLED", raising=False) + if exported_key is not None: + monkeypatch.setenv("OPENROUTER_API_KEY", exported_key) + (tmp_path / ".env").write_text("OPENROUTER_API_KEY=file-key\n", encoding="utf-8") + arguments: dict[str, object] = {} + + def fake_openai(**kwargs: object) -> object: + """Record constructor arguments without making a request.""" + arguments.update(kwargs) + return object() + + monkeypatch.setattr(openai, "OpenAI", fake_openai) + _create_api_client("openrouter") + assert arguments["api_key"] == (exported_key or "file-key") + + +def test_extract_metadata_uses_image_and_validates_output(tmp_path: Path) -> None: + """Send one image with a strict schema and validate the returned metadata.""" + image_path = tmp_path / "secret_source_table_000.png" + config_path = tmp_path / "config.json" + _write_image(image_path) + _write_config(config_path) + metadata = DataSnapshotMetadata(title="Displacement by year") + responses = _FakeResponses( + SimpleNamespace( + id="resp_test", + status="completed", + incomplete_details=None, + output_text=metadata.model_dump_json(), + usage=SimpleNamespace(input_tokens=12, output_tokens=4, total_tokens=16), + ) + ) + + result = extract_metadata( + image_path, + config_path=config_path, + client=SimpleNamespace(responses=responses), + ) + + assert result.metadata == metadata + assert result.error is None + assert responses.request is not None + assert responses.request["text"]["format"] == _response_format() + request_input = responses.request["input"] + prompt_text = " ".join( + item["text"] + for message in request_input + for item in message["content"] + if item["type"] == "input_text" + ) + assert image_path.name not in prompt_text + assert request_input[1]["content"][1]["image_url"].startswith( + "data:image/png;base64," + ) + + +def test_response_format_is_strict_and_openai_compatible() -> None: + """Generate Structured Outputs from the canonical Pydantic model.""" + response_format = _response_format() + schema = response_format["schema"] + serialized = json.dumps(schema) + + assert response_format["type"] == "json_schema" + assert response_format["strict"] is True + assert schema["required"] == list(schema["properties"]) + for unsupported in ("allOf", "else", "if", "not", "then"): + assert f'"{unsupported}"' not in serialized + assert "Validation rules:" in serialized + + +def test_incomplete_response_preserves_details_without_parsing(tmp_path: Path) -> None: + """Report output-budget truncation before attempting Pydantic parsing.""" + image_path = tmp_path / "snapshot.png" + config_path = tmp_path / "config.json" + _write_image(image_path) + _write_config(config_path) + response = SimpleNamespace( + id="resp_incomplete", + status="incomplete", + incomplete_details=SimpleNamespace(reason="max_output_tokens"), + output_text='{"title": "truncated', + usage=None, + ) + + result = extract_metadata( + image_path, + config_path=config_path, + client=SimpleNamespace(responses=_FakeResponses(response)), + ) + + assert result.metadata is None + assert result.api_status == "incomplete" + assert result.incomplete_details == {"reason": "max_output_tokens"} + assert result.error == "Response was incomplete: max_output_tokens." + assert result.raw_output == '{"title": "truncated' + + +def test_failed_response_reports_api_error(tmp_path: Path) -> None: + """Include the provider's failure code and message in the result.""" + image_path = tmp_path / "snapshot.png" + config_path = tmp_path / "config.json" + _write_image(image_path) + _write_config(config_path) + response = SimpleNamespace( + id="resp_failed", + status="failed", + incomplete_details=None, + output_text="", + usage=None, + error=SimpleNamespace( + code="provider_error", + message="Structured output was rejected.", + ), + ) + + result = extract_metadata( + image_path, + config_path=config_path, + client=SimpleNamespace(responses=_FakeResponses(response)), + ) + + assert result.metadata is None + assert result.api_status == "failed" + assert result.error == ( + "Response ended with status: failed " + "(provider_error: Structured output was rejected.)." + ) + + +def test_config_rejects_pipeline_owned_fields(tmp_path: Path) -> None: + """Prevent config files from replacing prompts, input, or output schema.""" + config_path = tmp_path / "config.json" + _write_config(config_path, input=[]) + + try: + load_extraction_config(config_path) + except ValueError as error: + assert "'input'" in str(error) + else: + raise AssertionError("Pipeline-owned input config was accepted.") + + +def test_config_rejects_unknown_provider(tmp_path: Path) -> None: + """Reject provider names without a configured client endpoint.""" + config_path = tmp_path / "config.json" + _write_config(config_path, provider="unknown") + + try: + load_extraction_config(config_path) + except ValueError as error: + assert "provider" in str(error) + else: + raise AssertionError("Unknown metadata provider was accepted.") diff --git a/tests/data_snapshot/test_metadata_schema.py b/tests/data_snapshot/test_metadata_schema.py new file mode 100644 index 00000000..f219c9c7 --- /dev/null +++ b/tests/data_snapshot/test_metadata_schema.py @@ -0,0 +1,38 @@ +"""Test the canonical Data Snapshot Metadata Schema v1.4 port.""" + +import pytest +from pydantic import BaseModel, TypeAdapter, ValidationError + +import ai4data.data_snapshot.metadata_schema as metadata_models +from ai4data.data_snapshot import DataSnapshotMetadata + + +def test_minimal_record_is_empty_and_forbids_unknown_fields() -> None: + """Accept an empty record while rejecting fields outside Schema v1.4.""" + assert DataSnapshotMetadata().model_dump(exclude_none=True) == {} + with pytest.raises(ValidationError, match="Extra inputs are not permitted"): + DataSnapshotMetadata.model_validate({"source_document_title": "Removed"}) + + +def test_every_public_field_example_validates() -> None: + """Validate every published example against its canonical field type.""" + for name in metadata_models.__all__: + model = getattr(metadata_models, name) + if not isinstance(model, type) or not issubclass(model, BaseModel): + continue + for field_name, field in model.model_fields.items(): + assert field.examples, f"{name}.{field_name} has no examples" + adapter = TypeAdapter(field.rebuild_annotation()) + for example in field.examples: + adapter.validate_python(example) + if model is DataSnapshotMetadata: + model.model_validate({field_name: example}) + + +def test_export_identifies_schema_v1_4() -> None: + """Preserve the canonical schema title, dialect, and version marker.""" + schema = DataSnapshotMetadata.model_json_schema() + + assert schema["title"] == "Data Snapshot Metadata Schema v1.4" + assert schema["$schema"] == "https://json-schema.org/draft/2020-12/schema" + assert schema["x-schema-version"] == "1.4" diff --git a/tests/data_snapshot/test_pipeline.py b/tests/data_snapshot/test_pipeline.py new file mode 100644 index 00000000..735a930b --- /dev/null +++ b/tests/data_snapshot/test_pipeline.py @@ -0,0 +1,188 @@ +"""Test model selection and model-scoped output orchestration.""" + +from pathlib import Path + +import pymupdf +import pytest +from PIL import Image + +from ai4data.data_snapshot import Detection, run_extraction, run_pipeline +from ai4data.data_snapshot import pipeline +from ai4data.data_snapshot.adapters import create_adapter + + +class _FakeDetector: + """Return one table detection for pipeline tests.""" + + def detect(self, image: Image.Image) -> list[Detection]: + """Return a full-page table detection.""" + return [Detection("table", (0.0, 0.0, 1.0, 1.0), 0.9)] + + +@pytest.mark.parametrize("model_args", [[], ["YOLO11"]]) +def test_run_extraction_selects_one_model_and_scopes_output( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, model_args: list[str] +) -> None: + """Use one selected adapter and write beneath its model directory.""" + pdf_path = tmp_path / "report.pdf" + with pymupdf.open() as document: + document.new_page(width=72, height=72) + document.save(pdf_path) + selected_models: list[str] = [] + + def fake_create_adapter(model_name: str) -> _FakeDetector: + """Record and return the selected test detector.""" + selected_models.append(model_name) + return _FakeDetector() + + monkeypatch.setattr(pipeline, "create_adapter", fake_create_adapter) + + paths = run_extraction(pdf_path, tmp_path / "output", *model_args, dpi=72) + + expected_model = model_args[0].lower() if model_args else "doclayoutyolo" + assert selected_models == [expected_model] + assert paths == [ + tmp_path / "output" / expected_model / "snapshots" / "report_table_000.png" + ] + + +def test_create_adapter_rejects_unknown_model() -> None: + """Reject an unknown model identifier.""" + with pytest.raises(ValueError, match="Unsupported model: unknown"): + create_adapter("unknown") + + +@pytest.mark.parametrize("entry_point", ["run_extraction", "run_pipeline"]) +@pytest.mark.parametrize("model_name", ["unknown", "../other", "absolute"]) +def test_invalid_model_preserves_outputs_before_forced_cleanup( + tmp_path: Path, entry_point: str, model_name: str +) -> None: + """Reject unsupported names and paths without deleting existing outputs.""" + pdf_path = tmp_path / "report.pdf" + pdf_path.write_bytes(b"No rendering should occur for an invalid model.") + output_dir = tmp_path / "output" + output_dir.mkdir() + if model_name == "absolute": + model_name = str(tmp_path / "other") + model_dir = output_dir / model_name + snapshot_path = model_dir / "snapshots" / "report_table_000.png" + metadata_path = model_dir / "metadata" / "report_table_000.json" + snapshot_path.parent.mkdir(parents=True) + metadata_path.parent.mkdir(parents=True) + snapshot_path.write_bytes(b"keep snapshot") + metadata_path.write_bytes(b"keep metadata") + + with pytest.raises(ValueError, match="Unsupported model:"): + getattr(pipeline, entry_point)(pdf_path, output_dir, model_name, force=True) + + assert snapshot_path.read_bytes() == b"keep snapshot" + assert metadata_path.read_bytes() == b"keep metadata" + + +def test_run_pipeline_skips_all_stages_for_existing_outputs( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], +) -> None: + """Skip both stages when an exact output already exists.""" + pdf_path = tmp_path / "report.pdf" + with pymupdf.open() as document: + document.new_page(width=72, height=72) + document.save(pdf_path) + snapshot_path = tmp_path / "output" / "tfid" / "snapshots" / "report_table_000.png" + snapshot_path.parent.mkdir(parents=True) + snapshot_path.write_bytes(b"existing") + monkeypatch.setattr( + pipeline, + "create_adapter", + lambda model_name: (_ for _ in ()).throw( + AssertionError("Skipped extraction loaded the snapshot detector.") + ), + ) + + paths = run_pipeline( + pdf_path, + tmp_path / "output", + "tfid", + dpi=72, + ) + + assert paths == [] + assert capsys.readouterr().out == ( + "Existing output files for report.pdf exist. " + "Skipped data snapshot and metadata extraction.\n" + ) + + +def test_run_extraction_force_cleans_exact_outputs_before_extraction( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], +) -> None: + """Remove crops and metadata for only the forced PDF stem.""" + pdf_path = tmp_path / "report.pdf" + with pymupdf.open() as document: + document.new_page(width=72, height=72) + document.save(pdf_path) + model_dir = tmp_path / "output" / "tfid" + snapshot_path = model_dir / "snapshots" / "report_table_000.png" + metadata_path = model_dir / "metadata" / "report_table_000.json" + prefix_neighbor = model_dir / "snapshots" / "report_figure_2024_table_000.png" + snapshot_path.parent.mkdir(parents=True) + metadata_path.parent.mkdir(parents=True) + snapshot_path.write_bytes(b"old") + metadata_path.write_text("{}\n", encoding="utf-8") + prefix_neighbor.write_bytes(b"keep") + monkeypatch.setattr(pipeline, "create_adapter", lambda model_name: _FakeDetector()) + + paths = run_extraction( + pdf_path, + tmp_path / "output", + "tfid", + dpi=72, + force=True, + ) + + assert metadata_path.exists() is False + assert prefix_neighbor.read_bytes() == b"keep" + assert len(paths) == 1 + assert paths[0] == snapshot_path + assert "Cleaning up before continuing... Done!" in capsys.readouterr().out + + +def test_run_pipeline_extracts_metadata_for_new_snapshots( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """Pass newly extracted crops directly to metadata extraction.""" + pdf_path = tmp_path / "report.pdf" + with pymupdf.open() as document: + document.new_page(width=72, height=72) + document.save(pdf_path) + model_dir = tmp_path / "output" / "tfid" + snapshot_path = model_dir / "snapshots" / "report_table_000.png" + metadata_path = model_dir / "metadata" / "report_table_000.json" + calls: list[tuple[list[Path], Path, str]] = [] + monkeypatch.setattr(pipeline, "create_adapter", lambda model_name: _FakeDetector()) + + def fake_extract_metadata_files( + snapshot_paths: list[Path], + output_dir: Path, + *, + metadata_provider: str, + ) -> list[Path]: + """Record the crops passed to the metadata stage.""" + calls.append((snapshot_paths, output_dir, metadata_provider)) + return [metadata_path] + + monkeypatch.setattr(pipeline, "extract_metadata_files", fake_extract_metadata_files) + + paths = run_pipeline( + pdf_path, + tmp_path / "output", + "tfid", + dpi=72, + metadata_provider="openai", + ) + + assert calls == [([snapshot_path], model_dir / "metadata", "openai")] + assert paths == [snapshot_path, metadata_path] diff --git a/tests/data_snapshot/test_submission.py b/tests/data_snapshot/test_submission.py new file mode 100644 index 00000000..ccb49db4 --- /dev/null +++ b/tests/data_snapshot/test_submission.py @@ -0,0 +1,19 @@ +"""Test the placeholder data snapshot submission function.""" + +from pathlib import Path + +from pytest import CaptureFixture + +from ai4data.data_snapshot import submit_snapshot + + +def test_submit_snapshot_prints_paths(capsys: CaptureFixture[str]) -> None: + """Print both paths without attempting an external submission.""" + submit_snapshot( + Path("snapshots/report_figure_000.png"), "metadata/report_figure_000.json" + ) + + assert capsys.readouterr().out == ( + "Would submit snapshots/report_figure_000.png with " + "metadata/report_figure_000.json\n" + ) diff --git a/tests/data_snapshot/test_tfid.py b/tests/data_snapshot/test_tfid.py new file mode 100644 index 00000000..e4b645dc --- /dev/null +++ b/tests/data_snapshot/test_tfid.py @@ -0,0 +1,96 @@ +"""Test the TF-ID adapter without loading model weights.""" + +from pathlib import Path +from typing import Any + +import pytest +from PIL import Image + +from ai4data.data_snapshot.adapters.tfid import TFIDAdapter, TFIDConfig + + +class _FakeInputs(dict[str, object]): + """Provide the device transfer used by Transformers batch inputs.""" + + def to(self, device: str) -> "_FakeInputs": + """Record no device state for the CPU-only fake inputs.""" + return self + + +class _FakeModel: + """Return a fixed generated token sequence.""" + + device = "cpu" + + def generate(self, **arguments: Any) -> list[list[int]]: + """Return one fake generated sequence.""" + assert arguments == { + "input_ids": "input_ids", + "pixel_values": "pixel_values", + "max_new_tokens": 1024, + "do_sample": False, + "num_beams": 3, + } + return [[1, 2, 3]] + + +class _FakeProcessor: + """Return fixed TF-ID inputs and parsed object detections.""" + + def __call__(self, **arguments: Any) -> _FakeInputs: + """Return the two tensors consumed by generation.""" + assert arguments["text"] == "" + assert arguments["return_tensors"] == "pt" + return _FakeInputs(input_ids="input_ids", pixel_values="pixel_values") + + def batch_decode(self, *args: object, **kwargs: object) -> list[str]: + """Return one generated TF-ID response.""" + return ["generated"] + + def post_process_generation( + self, *args: object, **kwargs: object + ) -> dict[str, Any]: + """Return figure, table, ignored, and undersized predictions.""" + return { + "": { + "bboxes": [ + [10.0, 20.0, 90.0, 80.0], + [0.0, 0.0, 20.0, 10.0], + [10.0, 10.0, 30.0, 30.0], + [0.0, 0.0, 1.0, 1.0], + ], + "labels": ["Figure", "Table", "Text", "Figure"], + "scores": [0.9, 1.2, 0.8, 0.7], + } + } + + +def test_tfid_adapter_normalizes_and_filters_detections() -> None: + """Map TF-ID output to canonical labels, boxes, and scores.""" + adapter = TFIDAdapter( + TFIDConfig(min_area=0.008), + model=_FakeModel(), + processor=_FakeProcessor(), + ) + + detections = adapter.detect(Image.new("RGB", (100, 100))) + + assert [(item.label, item.bbox, item.score) for item in detections] == [ + ("figure", (0.1, 0.2, 0.9, 0.8), 0.9), + ("table", (0.0, 0.0, 0.2, 0.1), 1.0), + ] + + +def test_tfid_config_uses_pinned_large_model_in_local_cache() -> None: + """Keep TF-ID Large and its remote code pinned in the local cache.""" + config = TFIDConfig() + + assert config.model_id == "yifeihu/TF-ID-large" + assert config.revision == "a91ffef47039768b2deff82681b75d77252703f2" + assert Path(config.cache_dir) == Path(".cache/models/tfid") + + +def test_tfid_adapter_requires_model_and_processor_together() -> None: + """Reject incomplete dependency injection before inference.""" + with pytest.raises(ValueError, match="supplied together"): + TFIDAdapter(model=_FakeModel()) diff --git a/tests/data_snapshot/test_yolo11.py b/tests/data_snapshot/test_yolo11.py new file mode 100644 index 00000000..a5bdf014 --- /dev/null +++ b/tests/data_snapshot/test_yolo11.py @@ -0,0 +1,16 @@ +"""Test the YOLO11 adapter without loading model weights.""" + +from pathlib import Path + +from ai4data.data_snapshot.adapters.yolo11 import YOLO11Adapter, YOLO11Config + + +def test_yolo11_adapter_uses_medium_weights_in_local_cache() -> None: + """Keep medium YOLO11 weights in the repository-local ignored cache.""" + model = object() + adapter = YOLO11Adapter(model=model) + + assert adapter._model is model + assert adapter.config == YOLO11Config() + assert adapter.config.filename == "yolo11m_doc_layout.pt" + assert Path(adapter.config.cache_dir) == Path(".cache/models") diff --git a/tests/data_snapshot/test_yolo26.py b/tests/data_snapshot/test_yolo26.py new file mode 100644 index 00000000..0be8c481 --- /dev/null +++ b/tests/data_snapshot/test_yolo26.py @@ -0,0 +1,112 @@ +"""Test the YOLO26 adapter without loading model weights.""" + +import sys +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import pytest +from PIL import Image + +from ai4data.data_snapshot.adapters.yolo26 import YOLO26Adapter, YOLO26Config + + +class _FakeTensor: + """Provide the tensor calls used by the adapter.""" + + def __init__(self, values: list[Any]) -> None: + self._values = values + + def cpu(self) -> "_FakeTensor": + """Return this CPU-backed fake tensor.""" + return self + + def tolist(self) -> list[Any]: + """Return the stored JSON-compatible values.""" + return self._values + + +class _FakeBoxes: + """Provide fake YOLO box tensors.""" + + xyxy = _FakeTensor( + [ + [10.0, 20.0, 90.0, 80.0], + [0.0, 0.0, 20.0, 10.0], + [0.0, 0.0, 1.0, 1.0], + [10.0, 10.0, 30.0, 30.0], + ] + ) + conf = _FakeTensor([0.9, 1.2, 0.5, 0.8]) + cls = _FakeTensor([0.0, 1.0, 0.0, 2.0]) + + +class _FakeResult: + """Provide one fake Ultralytics result.""" + + boxes = _FakeBoxes() + names = {0: "Picture", 1: "Table", 2: "Text"} + + +class _FakeModel: + """Record prediction arguments and return a fake result.""" + + def __init__(self) -> None: + self.arguments: dict[str, Any] | None = None + + def predict(self, image: Image.Image, **arguments: Any) -> list[_FakeResult]: + """Return fixed detections for the supplied image.""" + self.arguments = arguments + return [_FakeResult()] + + +def test_yolo26_adapter_normalizes_and_filters_detections() -> None: + """Map model output to canonical labels, boxes, and scores.""" + model = _FakeModel() + config = YOLO26Config(device="cpu", imgsz=640, min_area=0.008) + adapter = YOLO26Adapter(config, model=model) + + detections = adapter.detect(Image.new("RGB", (100, 100))) + + assert [(item.label, item.bbox, item.score) for item in detections] == [ + ("figure", (0.1, 0.2, 0.9, 0.8), 0.9), + ("table", (0.0, 0.0, 0.2, 0.1), 1.0), + ] + assert model.arguments == { + "imgsz": 640, + "conf": 0.25, + "iou": 0.7, + "device": "cpu", + "verbose": False, + } + + +def test_yolo26_config_rejects_invalid_threshold() -> None: + """Reject invalid model configuration before inference.""" + with pytest.raises(ValueError, match="Confidence"): + YOLO26Config(conf=1.1) + + +def test_yolo26_config_uses_medium_weights_in_local_cache() -> None: + """Keep medium weights in the repository-local ignored cache.""" + config = YOLO26Config() + + assert config.filename == "yolo26m_doc_layout.pt" + assert Path(config.cache_dir) == Path(".cache/models") + + +def test_yolo26_adapter_loads_cached_weights_without_hub( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """Load existing medium weights without contacting Hugging Face.""" + model_path = tmp_path / "yolo26m_doc_layout.pt" + model_path.touch() + monkeypatch.setitem( + sys.modules, + "ultralytics", + SimpleNamespace(YOLO=lambda path: path), + ) + + adapter = YOLO26Adapter(YOLO26Config(cache_dir=tmp_path)) + + assert adapter._model == str(model_path)