diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py new file mode 100644 index 0000000000..d2038ea280 --- /dev/null +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py @@ -0,0 +1,402 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The stored contract for the dataset profiler. + +A ``DatasetProfile`` is machine-owned metadata the profiler computes once, at the Files layer, so +every consumer reads one typed description of a dataset instead of downloading and re-inspecting it. + +The profile has two stored layers: + +* **structure** — predominantly facts: file layout, splits, the derived row schema (``features``), + and per-column ``stats``. The one detected attribute living here is the ``FeatureSchema.semantic_role`` + marker stacked on the feature node it describes. +* **classification** — an objective description of what the data *is*: ``dataset_type``, the + ``format`` / ``prompt_form`` axes, and ``verifiability``. + +Vocabularies (``dataset_type``, ``semantic_role``, ``modality``, ...) are open ``str`` values with +documented canonical sets, not closed enums: only known values are emitted, but consumers must +tolerate unknown ones so the vocabulary can grow without a breaking change. Pydantic's default +``extra="ignore"`` gives the same forward-compatibility for unknown *fields*. + +This module is pydantic-only — no platform dependencies — so the profiler can import it as a +standalone contract and ``DatasetMetadataContent`` can later carry it as a typed field. +""" + +from __future__ import annotations + +from datetime import datetime + +from pydantic import BaseModel, Field, model_validator + +# Semver of THIS contract. Gates consumer compatibility: new detectors or vocabulary values are a +# minor bump; a change to the fields below is a major bump. +PROFILE_SCHEMA_VERSION = "1.0" + + +# ---- classification: an objective description of what the data is (computed, stored) ----------- + + +class Evidence(BaseModel): + """Why the profiler believes what it detected. + + Captured at profile time — the only moment it is cheap and guaranteed to match the stored + result; once the data or the profiler version moves, a re-run explains the *new* snapshot, not + the stored one. + """ + + kind: str = Field( + description="column_name | column_dtype | content_probe | split_name | file_name | card_metadata", + ) + detail: str = Field( + description="Self-describing evidence, e.g. \"answer matches '#### ' in 100% of 1024 sampled rows\".", + ) + + +class Verifiability(BaseModel): + """A found verification target. Present only when one exists; absence *is* the claim (not verifiable).""" + + method: str = Field( + description="extractable_final_answer | ground_truth_column | constraint | test_cases", + ) + coverage: float | None = Field( + default=None, + ge=0.0, + le=1.0, + description="Fraction of sampled rows with a usable verification target.", + ) + evidence: list[Evidence] = Field(default_factory=list) + + +class PartitionClassification(BaseModel): + """What the data *is*, described objectively — the stored basis a downstream reader uses to decide + which tasks the dataset can train. + + Holds the partition-level findings only: column-level semantics are the ``semantic_role`` markers + on the feature nodes they describe, but the evidence for *why* they were assigned is recorded here. + """ + + modality: str = Field(default="text", description="text | image_text | audio_text | ...") + dataset_type: str = Field(description="Dataset-type vocabulary (prompt_completion, preference_pair, ...).") + format: str | None = Field(default=None, description="standard | conversational | mixed") + prompt_form: str | None = Field(default=None, description="explicit | implicit | n/a") + verifiability: Verifiability | None = Field( + default=None, + description="Present only when a verification target was found; keeps its own coverage-scoped evidence.", + ) + evidence: list[Evidence] = Field( + default_factory=list, + description=( + "Why the type / roles / axes were assigned: one flat list, detail strings self-describe " + "what they support. A profile-time snapshot; unrecoverable once the data or profiler version move." + ), + ) + + +# ---- structure: facts + the stacked semantic_role detection (computed) ------------------------- + + +class Quantiles(BaseModel): + """A per-row distribution summary. p99 = long-tail sequence-length signal; max = hard cap.""" + + p50: int + p95: int + p99: int + max: int + + +class TextStats(BaseModel): + """Measurements for a ``string`` column.""" + + chars: Quantiles = Field(description="Per-row character-length distribution.") + + +class MessageStats(BaseModel): + """Measurements for a ``messages`` column (a list of ``{role, content}``).""" + + turns: Quantiles = Field(description="Per-row turn count (p99 -> packing long chats).") + content_chars: Quantiles = Field(description="Per-row total content length -> chat sequence length.") + roles_seen: list[str] = Field( + default_factory=list, + description=( + 'The distinct role strings actually present in the sampled rows, verbatim — e.g. ["system", ' + '"user", "assistant", "tool"], but equally ShareGPT\'s ["human", "gpt"] or a house convention. ' + "A measurement of row content, not a vocabulary the profiler picks from, so it is deliberately " + "not an enum: an unexpected role is the finding worth reporting, and normalizing or dropping it " + "would hide exactly what a consumer needs to see before choosing a chat template." + ), + ) + ends_with_assistant_rate: float = Field( + ge=0.0, + le=1.0, + description="Key signal separating an SFT target (conversation ends on an assistant turn) from a prompt-only row.", + ) + valid_alternation_rate: float = Field(ge=0.0, le=1.0) + has_tool_calls: bool = False + + +class NumericStats(BaseModel): + """Measurements for a numeric column.""" + + min: float + max: float + mean: float + + +class TextQuality(BaseModel): + """Cheap, single-pass corruption signals for a text column. Flags training-wrecking data, not + toxicity / PII. + """ + + whitespace_ratio: float = Field(ge=0.0, le=1.0, description="Padding / bad scraping.") + non_ascii_ratio: float = Field(ge=0.0, le=1.0, description="Encoding / non-Latin signal.") + repetition_score: float = Field(ge=0.0, le=1.0, description="Degenerate repeated-substring loops.") + + +class FeatureSchema(BaseModel): + """One node of the row schema, derived de novo from the data (there is no external JSON-Schema + store to reference). Carries the measured layout (name, dtype, children) plus at most one + detected ``semantic_role`` marker stacked on the same node. + + Recursive and fully expanded: a ``struct`` node has child ``fields``; a ``list`` / ``messages`` + node has an element ``items`` — for ``messages`` the per-message ``{role, content}`` struct is + spelled out, so a vision message whose content is a list of typed parts shows up structurally. + The column-level chat summary lives in ``MessageStats`` on the stats side. This tree is the + clean, bridgeable schema artifact (e.g. to a JSON Schema or a UI columns view). + """ + + name: str = Field(default="", description='Column / struct-field name; "" for a list element.') + dtype: str = Field( + description=( + "string | bool | int8..int64 / uint8..uint64 | float16/32/64 | struct | list | messages | " + "image | audio | video | json | ... — fixed-width numeric widths as the source file reports them." + ), + ) + semantic_role: str | None = Field( + default=None, + description=( + "Detected role (from the role vocabulary), valid at any depth of the tree; omitted when nothing " + "was detected. The only detected attribute in the structure layer — its evidence lands in " + "PartitionClassification.evidence. Named `semantic_role`, not `role`, so it never collides with a " + "message struct's `role` key." + ), + ) + fixed_length: int | None = Field( + default=None, + description=( + "dtype == list: constant observed element count (e.g. an embedding vector's 768), None when " + "variable. Multi-dimensional shapes compose via nesting." + ), + ) + fields: list[FeatureSchema] | None = Field(default=None, description="dtype == struct: named child fields.") + items: FeatureSchema | None = Field(default=None, description="dtype in {list, messages}: element schema.") + + @model_validator(mode="after") + def _fields_and_items_are_exclusive(self) -> FeatureSchema: + """A node is either a named-field container or has a single element schema, never both. + + Deliberately the only structural check here: it holds for *any* dtype, so it costs no + forward compatibility. Tying `fields` / `items` / `fixed_length` to specific dtype values + would instead reject a profile written by a newer profiler that added a container dtype, + which is exactly what the open vocabulary exists to prevent. + """ + if self.fields is not None and self.items is not None: + raise ValueError(f"feature {self.name!r}: `fields` and `items` are mutually exclusive") + return self + + +class CategoricalStats(BaseModel): + """Cardinality signals for string / int columns. + + ``distinct_count`` is always safe to store; the values themselves ARE row data, so they appear + only when proven to be a small enumeration by an exhaustive scan — the same + assert-only-what-was-proven rule applied everywhere the profiler would otherwise leak row content. + """ + + distinct_count: int = Field( + description=( + "Distinct values among scanned rows: ~=rows_scanned -> id-like; a small bounded set " + "corroborates score / category roles." + ), + ) + values: list[str] | None = Field( + default=None, + description="The proven enumeration; only when the scan was exhaustive and distinct_count <= 32.", + ) + + +class ColumnStats(BaseModel): + """Measurements for one top-level column (keyed by name in ``PartitionProfile.stats``). + + The kind-specific block is populated by dtype; deep measurements fold into it (e.g. + ``MessageStats.content_chars``) so stats stay flat — no path addressing to drift against the + schema tree. Never row values — profiles stay safe to display / export without leaking data — + with one gated exception: ``categorical.values``, a proven small enumeration. + """ + + null_rate: float = Field(default=0.0, ge=0.0, le=1.0) + text: TextStats | None = Field(default=None, description="dtype == string") + numeric: NumericStats | None = Field(default=None, description="dtype in {int*, uint*, float*}") + messages: MessageStats | None = Field(default=None, description="dtype == messages (list of {role, content})") + categorical: CategoricalStats | None = Field(default=None, description="low observed cardinality only") + quality: TextQuality | None = Field(default=None, description="dtype == string: corruption signals") + + +class FileRecord(BaseModel): + """One physical file, measured. + + Stores the exact digest inputs (so a profile self-describes its ``content_digest`` and per-file + staleness is computable) plus what the reader learned cheaply. + """ + + path: str = Field(description="Relative path within the fileset.") + size_bytes: int + checksum: str | None = Field( + default=None, + description=( + 'As the Files service reports it (e.g. "sha256:..."). None falls back to a (path, size) digest, ' + "which cannot detect a same-size in-place edit." + ), + ) + num_rows: int | None = Field( + default=None, + description="Exact only (parquet footer / exhaustive scan), else None.", + ) + + +class SplitProfile(BaseModel): + """A split within a partition. + + Resolution precedence (declared structure beats detection): + + 1. HF card front-matter when the fileset ships a README — ``configs[].data_files`` maps splits to + file globs explicitly; + 2. best-effort detection from file paths (train/test/validation markers, sharded layouts like + ``data/train-00000-of-00003.parquet``); path markers are matched against canonical names and + common aliases (val/valid/dev -> validation), the normalized concept lands in ``canonical``, and + the split keeps its on-disk ``name``; + 3. otherwise leave it alone: a single "default" split holding all files. + + A split encoded as a *data column* (a value inside each row rather than a file grouping) is not + resolved here; such files profile as a single split. + """ + + name: str = Field(description="The on-disk name: train | test | train_prefs | ...") + canonical: str | None = Field( + default=None, + description=( + "Normalized concept: train | validation | test; None when nothing matches. E.g. train_prefs -> " + "train, with the variant's intent kept in `name`." + ), + ) + files: list[FileRecord] = Field( + description=( + "Every file resolved into this split, measured. Partitioning is exhaustive and disjoint: each " + "file of the partition lands in exactly one split, so concatenating `files` across splits " + "reconstructs the partition's file list with no gaps or repeats." + ), + ) + num_examples: int | None = Field( + default=None, + description=( + "Rows in this split, counting every file in `files` whether or not it was scanned. Exact when " + "read from parquet footers or an exhaustive scan, otherwise extrapolated from the rows sampled — " + "check `SamplingInfo.exhaustive` before treating it as a fact. None when nothing usable was found." + ), + ) + + +class PartitionProfile(BaseModel): + """A file-group sharing one row schema, one top-level directory, and one split-name variant + (roughly an HF config); named after the directory / variant, else "default". + + File membership and row counts live on ``splits`` — every file lands in exactly one split, so + partition-level files / num_examples would be derivable duplication. + """ + + name: str = "default" + file_format: str = Field(description="jsonl | parquet | csv | arrow") + splits: list[SplitProfile] = Field(description="card-declared > path-detected > single 'default' split.") + features: list[FeatureSchema] = Field( + description="The row schema: measured layout plus detected role markers, derived de novo (nested).", + ) + stats: dict[str, ColumnStats] = Field( + default_factory=dict, + description=( + "Top-level column name -> measurements; sparse (a column with nothing worth measuring is " + "omitted); keys are a subset of the top-level `features` names." + ), + ) + classification: PartitionClassification + + @model_validator(mode="after") + def _stats_keys_subset_of_features(self) -> PartitionProfile: + """``stats`` is keyed by top-level column name, so every key must name a top-level feature + (the producer keys stats by ``feature.name``); a stray key is a malformed profile.""" + unknown = set(self.stats) - {feature.name for feature in self.features} + if unknown: + raise ValueError(f"stats keys must name top-level features; unknown columns: {sorted(unknown)}") + return self + + +# ---- envelope ---------------------------------------------------------------------------------- + + +class SamplingInfo(BaseModel): + """How much of the data the profile is based on. + + Consumers read ``exhaustive`` to decide whether stats are proven facts or estimates (e.g. only an + exhaustive profile can assert enum / required in a bridged JSON Schema, or that verifiability + coverage is truly 1.0). + """ + + exhaustive: bool = Field(description="True => every row of every file was parsed.") + strategy: str = Field( + description=( + "full | stratified_probes | random. Kept explicit alongside `exhaustive` because, with an open " + "strategy vocabulary, consumers can't derive exhaustiveness from the strategy name alone." + ), + ) + rows_scanned: int = Field(description="Total rows actually parsed across all files.") + rows_total: int | None = Field( + default=None, + description=( + "How many rows the whole fileset holds, scanned or not — the denominator `rows_scanned` is a " + "fraction of, so a consumer can judge how representative the stats are. Populated only when the " + "count is exact and cheap (summed parquet footers, or an exhaustive scan); None means unknown, " + "never zero and never an estimate." + ), + ) + files_scanned: int = Field( + description=( + "How many files were opened and read from (a count, not a list — the files themselves are " + "`SplitProfile.files`). Every file should be probed, since head-sampling a subset hides columns " + "that appear only in later shards; expect this to equal the fileset's file count, and be lower " + "only when scale forces file-level sampling." + ), + ) + per_file_row_cap: int | None = Field(default=None, description="Cap that bounded per-file reads, if any.") + seed: int | None = Field(default=None, description="RNG seed used for row selection, for reproducibility.") + + +class DatasetProfile(BaseModel): + """The machine-owned dataset profile — the root of the stored contract.""" + + profile_schema_version: str = Field( + default=PROFILE_SCHEMA_VERSION, + description='Semver of THIS contract (e.g. "1.0") — gates consumer compatibility.', + ) + content_digest: str = Field(description="Digest over the stored FileRecords; staleness = mismatch.") + created_at: datetime + profiler_info: dict = Field( + default_factory=dict, + description="Free-form profiler metadata (name, version, git sha, timings).", + ) + sampling: SamplingInfo = Field(description="How much data the profile is based on.") + partitions: list[PartitionProfile] = Field( + description="Single partition in the common homogeneous case; there is no fileset-level rollup.", + ) + + +# Resolve the recursive FeatureSchema self-reference (deferred by `from __future__ import annotations`). +FeatureSchema.model_rebuild() diff --git a/packages/nemo_platform_plugin/tests/files/test_dataset_profile.py b/packages/nemo_platform_plugin/tests/files/test_dataset_profile.py new file mode 100644 index 0000000000..ca7ba48484 --- /dev/null +++ b/packages/nemo_platform_plugin/tests/files/test_dataset_profile.py @@ -0,0 +1,356 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the dataset-profiler stored contract. + +Three representative datasets — a conversational prompt/completion set, a conversational preference +pair, and a standard scored-response set — are embedded as YAML fixtures and exercised as executable +expectations: if a field name, alias, or vocabulary value drifts, one of these deserializations +breaks. +""" + +from datetime import datetime + +import pytest +import yaml +from nemo_platform_plugin.files.dataset_profile import ( + PROFILE_SCHEMA_VERSION, + ColumnStats, + DatasetProfile, + Evidence, + FeatureSchema, + FileRecord, + MessageStats, + PartitionClassification, + PartitionProfile, + Quantiles, + SamplingInfo, + SplitProfile, + Verifiability, +) +from pydantic import ValidationError + +# --- Fixture: trl-lib/OpenMathReasoning (conversational prompt_completion, verifiable) --- +OPENMATHREASONING = """ +profile_schema_version: "1.0" +content_digest: sha256:7be1... +created_at: 2026-07-08T22:05:12Z +profiler_info: {name: nemo-dataset-profiler, version: 0.1.0} +sampling: {exhaustive: false, strategy: stratified_probes, rows_scanned: 2112, + rows_total: 3201061, + files_scanned: 33, per_file_row_cap: 64} +partitions: + - name: default + file_format: parquet + splits: + - {name: train, canonical: train, num_examples: 3200861, + files: [{path: train-00000-of-00032.parquet, size_bytes: 193777041, + checksum: sha256:9c1e..., num_rows: 100027}]} + - {name: test, canonical: test, num_examples: 200, + files: [{path: test-00000-of-00001.parquet, size_bytes: 411552, + checksum: sha256:02af..., num_rows: 200}]} + features: + - {name: prompt, dtype: messages, semantic_role: prompt, + items: {dtype: struct, fields: [{name: role, dtype: string}, {name: content, dtype: string}]}} + - {name: completion, dtype: messages, semantic_role: completion, + items: {dtype: struct, fields: [{name: role, dtype: string}, {name: content, dtype: string}]}} + stats: + prompt: {messages: {turns: {p50: 1, p95: 1, p99: 1, max: 1}, content_chars: {p50: 180, p95: 620, p99: 1100, max: 4800}, + roles_seen: [user], ends_with_assistant_rate: 0.0, valid_alternation_rate: 1.0}} + completion: {messages: {turns: {p50: 1, p95: 1, p99: 1, max: 1}, content_chars: {p50: 2400, p95: 7800, p99: 12000, max: 32000}, + roles_seen: [assistant], ends_with_assistant_rate: 1.0, valid_alternation_rate: 1.0}} + classification: + modality: text + dataset_type: prompt_completion + format: conversational + prompt_form: explicit + verifiability: + method: extractable_final_answer + coverage: 0.81 + evidence: [{kind: content_probe, detail: 'completion ends with \\boxed{...} in 81% of 2112 sampled rows'}] + evidence: + - {kind: column_name, detail: "prompt + completion column pair"} + - {kind: content_probe, detail: "prompt ends on a user turn, completion is a single assistant turn"} +""" + +# --- Fixture: trl-lib/hh-rlhf-helpful-base (conversational preference_pair, explicit) ----- +HH_RLHF_HELPFUL_BASE = """ +profile_schema_version: "1.0" +content_digest: sha256:5d20... +created_at: 2026-07-08T22:41:37Z +profiler_info: {name: nemo-dataset-profiler, version: 0.1.0} +sampling: {exhaustive: false, strategy: stratified_probes, rows_scanned: 1024, + rows_total: 46189, files_scanned: 2, per_file_row_cap: 512} +partitions: + - name: default + file_format: parquet + splits: + - {name: train, canonical: train, num_examples: 43835, + files: [{path: train-00000-of-00001.parquet, size_bytes: 22105331, + checksum: sha256:77b0..., num_rows: 43835}]} + - {name: test, canonical: test, num_examples: 2354, + files: [{path: test-00000-of-00001.parquet, size_bytes: 1198422, + checksum: sha256:5c1d..., num_rows: 2354}]} + features: + - {name: prompt, dtype: messages, semantic_role: prompt, + items: {dtype: struct, fields: [{name: role, dtype: string}, {name: content, dtype: string}]}} + - {name: chosen, dtype: messages, semantic_role: chosen, + items: {dtype: struct, fields: [{name: role, dtype: string}, {name: content, dtype: string}]}} + - {name: rejected, dtype: messages, semantic_role: rejected, + items: {dtype: struct, fields: [{name: role, dtype: string}, {name: content, dtype: string}]}} + stats: + prompt: {messages: {turns: {p50: 3, p95: 8, p99: 9, max: 9}, content_chars: {p50: 640, p95: 3200, p99: 5400, max: 9800}, + roles_seen: [user, assistant], ends_with_assistant_rate: 0.0, valid_alternation_rate: 1.0}} + chosen: {messages: {turns: {p50: 1, p95: 1, p99: 1, max: 1}, content_chars: {p50: 420, p95: 1400, p99: 2100, max: 3600}, + roles_seen: [assistant], ends_with_assistant_rate: 1.0, valid_alternation_rate: 1.0}} + rejected: {messages: {turns: {p50: 1, p95: 1, p99: 1, max: 1}, content_chars: {p50: 410, p95: 1380, p99: 2050, max: 3500}, + roles_seen: [assistant], ends_with_assistant_rate: 1.0, valid_alternation_rate: 1.0}} + classification: + modality: text + dataset_type: preference_pair + format: conversational + prompt_form: explicit + evidence: + - {kind: column_name, detail: "chosen + rejected column pair"} + - {kind: content_probe, detail: "prompt carries the multi-turn history ending on a user turn"} +""" + +# --- Fixture: nvidia/HelpSteer2 (standard scored_response, no verifiability) -------------- +HELPSTEER2 = """ +profile_schema_version: "1.0" +content_digest: sha256:c41f... +created_at: 2026-07-09T10:12:45Z +profiler_info: {name: nemo-dataset-profiler, version: 0.1.0} +sampling: {exhaustive: false, strategy: stratified_probes, rows_scanned: 1024, + rows_total: 21362, files_scanned: 2, per_file_row_cap: 512} +partitions: + - name: default + file_format: parquet + splits: + - {name: train, canonical: train, num_examples: 20324, + files: [{path: train-00000-of-00001.parquet, size_bytes: 44201991, + checksum: sha256:e410..., num_rows: 20324}]} + - {name: validation, canonical: validation, num_examples: 1038, + files: [{path: validation-00000-of-00001.parquet, size_bytes: 2311008, + checksum: sha256:8bd2..., num_rows: 1038}]} + features: + - {name: prompt, dtype: string, semantic_role: prompt} + - {name: response, dtype: string, semantic_role: completion} + - {name: helpfulness, dtype: int64, semantic_role: score} + - {name: correctness, dtype: int64, semantic_role: score} + - {name: coherence, dtype: int64, semantic_role: score} + - {name: complexity, dtype: int64, semantic_role: score} + - {name: verbosity, dtype: int64, semantic_role: score} + stats: + prompt: {text: {chars: {p50: 320, p95: 2200, p99: 5600, max: 12000}}, + quality: {whitespace_ratio: 0.16, non_ascii_ratio: 0.004, repetition_score: 0.02}} + response: {text: {chars: {p50: 1350, p95: 3900, p99: 6200, max: 10500}}, + quality: {whitespace_ratio: 0.15, non_ascii_ratio: 0.003, repetition_score: 0.04}} + helpfulness: {numeric: {min: 0, max: 4, mean: 2.8}, categorical: {distinct_count: 5}} + correctness: {numeric: {min: 0, max: 4, mean: 2.9}, categorical: {distinct_count: 5}} + coherence: {numeric: {min: 0, max: 4, mean: 3.5}, categorical: {distinct_count: 5}} + complexity: {numeric: {min: 0, max: 4, mean: 1.6}, categorical: {distinct_count: 5}} + verbosity: {numeric: {min: 0, max: 4, mean: 1.5}, categorical: {distinct_count: 5}} + classification: + modality: text + dataset_type: scored_response + format: standard + prompt_form: explicit + evidence: + - {kind: column_name, detail: "prompt + response pair; five rating columns match score aliases"} + - {kind: content_probe, detail: "all five ratings bounded 0-4 with 5 distinct values"} + - {kind: card_metadata, detail: "README tag 'human-feedback' corroborates scored human ratings"} +""" + +FIXTURES = { + "OpenMathReasoning": OPENMATHREASONING, + "hh-rlhf-helpful-base": HH_RLHF_HELPFUL_BASE, + "HelpSteer2": HELPSTEER2, +} + + +def _build_profile() -> DatasetProfile: + """A hand-built profile exercising every model in the contract.""" + return DatasetProfile( + content_digest="sha256:deadbeef", + created_at=datetime(2026, 7, 13, 12, 0, 0), + profiler_info={"name": "nemo-dataset-profiler", "version": "0.1.0"}, + sampling=SamplingInfo( + exhaustive=False, + strategy="stratified_probes", + rows_scanned=1024, + rows_total=2048, + files_scanned=2, + per_file_row_cap=512, + seed=7, + ), + partitions=[ + PartitionProfile( + file_format="parquet", + splits=[ + SplitProfile( + name="train", + canonical="train", + num_examples=2048, + files=[ + FileRecord(path="train-00000.parquet", size_bytes=123, checksum="sha256:ab", num_rows=2048) + ], + ) + ], + features=[ + FeatureSchema(name="prompt", dtype="string", semantic_role="prompt"), + FeatureSchema(name="response", dtype="string", semantic_role="completion"), + ], + stats={ + "prompt": ColumnStats(text=None), + "response": ColumnStats(numeric=None), + }, + classification=PartitionClassification( + dataset_type="prompt_completion", + format="standard", + prompt_form="explicit", + verifiability=Verifiability( + method="extractable_final_answer", + coverage=0.9, + evidence=[Evidence(kind="content_probe", detail="ends with #### in 90% of rows")], + ), + evidence=[Evidence(kind="column_name", detail="prompt + response pair")], + ), + ) + ], + ) + + +def test_schema_version_defaults_to_constant(): + profile = _build_profile() + assert profile.profile_schema_version == PROFILE_SCHEMA_VERSION == "1.0" + + +def test_round_trip_json_is_lossless(): + profile = _build_profile() + restored = DatasetProfile.model_validate_json(profile.model_dump_json()) + assert restored == profile + + +@pytest.mark.parametrize("name", list(FIXTURES)) +def test_fixture_deserializes(name): + """Every fixture loads into the contract and round-trips.""" + profile = DatasetProfile.model_validate(yaml.safe_load(FIXTURES[name])) + assert profile.profile_schema_version == "1.0" + assert profile.partitions[0].name == "default" + # Round-trip through JSON is lossless. + assert DatasetProfile.model_validate_json(profile.model_dump_json()) == profile + + +def test_openmathreasoning_locks_contract_shape(): + profile = DatasetProfile.model_validate(yaml.safe_load(OPENMATHREASONING)) + part = profile.partitions[0] + assert part.classification.dataset_type == "prompt_completion" + assert part.classification.format == "conversational" + # semantic_role is stacked on the feature node; message struct spelled out under items. + prompt_feature = part.features[0] + assert prompt_feature.name == "prompt" + assert prompt_feature.dtype == "messages" + assert prompt_feature.semantic_role == "prompt" + assert prompt_feature.items.fields[0].name == "role" + # Verifiability carries its own coverage + scoped evidence. + verify = part.classification.verifiability + assert verify.method == "extractable_final_answer" + assert verify.coverage == pytest.approx(0.81) + # Message stats fold into the messages block. + assert part.stats["prompt"].messages.ends_with_assistant_rate == 0.0 + assert part.stats["completion"].messages.roles_seen == ["assistant"] + + +def test_helpsteer2_flat_schema_and_no_verifiability(): + profile = DatasetProfile.model_validate(yaml.safe_load(HELPSTEER2)) + part = profile.partitions[0] + assert part.classification.dataset_type == "scored_response" + assert part.classification.format == "standard" + # Absence of a verifiability object *is* the "not verifiable" claim. + assert part.classification.verifiability is None + # Physical column name != role (response -> completion). + response_feature = next(f for f in part.features if f.name == "response") + assert response_feature.semantic_role == "completion" + # Bounded rating scale corroborated by cardinality. + assert part.stats["helpfulness"].categorical.distinct_count == 5 + assert part.stats["helpfulness"].numeric.max == 4.0 + # card_metadata evidence survives (declared card tags corroborate, never override). + kinds = {e.kind for e in part.classification.evidence} + assert "card_metadata" in kinds + + +def test_semantic_role_reachable_at_any_depth(): + """A role marker nested inside a response list (e.g. a rank) is reachable; a flat column->role + dict could not address it.""" + answers = FeatureSchema( + name="answers", + dtype="list", + items=FeatureSchema( + dtype="struct", + fields=[ + FeatureSchema(name="answer", dtype="string", semantic_role="completion"), + FeatureSchema(name="model", dtype="string", semantic_role="provenance"), + FeatureSchema(name="rank", dtype="int64", semantic_role="rank"), + ], + ), + ) + assert answers.items.fields[2].semantic_role == "rank" + # Round-trips through JSON without losing the nested marker. + restored = FeatureSchema.model_validate_json(answers.model_dump_json()) + assert restored.items.fields[2].semantic_role == "rank" + + +def test_vocabularies_are_open(): + """Unknown vocabulary values must be accepted so the vocabulary can grow.""" + classification = PartitionClassification( + modality="video_text", + dataset_type="some_future_type", + format="mixed", + ) + assert classification.dataset_type == "some_future_type" + feature = FeatureSchema(dtype="tensor", semantic_role="a_role_added_next_year") + assert feature.semantic_role == "a_role_added_next_year" + + +def test_fields_and_items_are_mutually_exclusive(): + """A node cannot be both a named-field container and a single-element container.""" + with pytest.raises(ValidationError, match="mutually exclusive"): + FeatureSchema( + name="broken", + dtype="struct", + fields=[FeatureSchema(name="a", dtype="string")], + items=FeatureSchema(dtype="string"), + ) + + +def test_container_shape_is_not_pinned_to_known_dtypes(): + """The exclusivity check must not become a dtype whitelist: a container dtype added by a newer + profiler still loads on an older reader, which is what the open vocabulary buys.""" + future_map = FeatureSchema(name="attrs", dtype="map", fields=[FeatureSchema(name="k", dtype="string")]) + assert [field.name for field in future_map.fields or []] == ["k"] + future_tensor = FeatureSchema(name="embedding", dtype="tensor", fixed_length=768) + assert future_tensor.fixed_length == 768 + + +def test_unknown_fields_are_ignored_for_forward_compat(): + """A profile written by a newer minor version (extra fields) still loads on an older reader.""" + doc = yaml.safe_load(HELPSTEER2) + doc["some_future_top_level_field"] = {"anything": 1} + doc["partitions"][0]["classification"]["future_axis"] = "value" + profile = DatasetProfile.model_validate(doc) + assert profile.partitions[0].classification.dataset_type == "scored_response" + + +def test_quantiles_and_message_stats_construct(): + """Smoke-check the leaf stat models are wired as documented.""" + stats = MessageStats( + turns=Quantiles(p50=1, p95=3, p99=5, max=9), + content_chars=Quantiles(p50=100, p95=500, p99=900, max=2000), + roles_seen=["user", "assistant", "tool"], + ends_with_assistant_rate=1.0, + valid_alternation_rate=0.98, + has_tool_calls=True, + ) + assert stats.turns.max == 9 + assert stats.has_tool_calls is True