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 index d2038ea280..5a2e247c7b 100644 --- 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 @@ -19,8 +19,20 @@ 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. +**Why this lives in a shared package rather than with the profiler that writes it.** The Files +service is a first-class consumer, not a bystander: it stores a profile as its own entity and serves +it, so its entities, endpoints and schemas all need this type. Were the contract to live in the +datasets plugin, a core service would depend on an optional plugin to deserialize rows in its own +database — a deployment that installs no profiler still holds stored profiles and still answers +``GET .../filesets/{name}/profile``. Keeping it here means neither side depends on the other: the +module is pydantic-only, with no platform dependencies, so the profiler imports it standalone while +Files imports it as the type it persists. + +It sits under ``files/`` because Files is what stores and serves it, alongside the rest of that +service's shared contract — including ``metadata.py``, which houses the equally dataset-shaped +``DatasetMetadataContent``. Note that a profile is *not* carried inside fileset metadata: it is a +separate entity, so writing one cannot clobber an unrelated metadata edit that lands between a read +and a write. """ from __future__ import annotations @@ -30,7 +42,9 @@ 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. +# minor bump; a change to the fields below is a major bump. Still 1.0 because nothing consumes it +# yet — the fields have moved a great deal, but pre-release churn is not a break for anyone, and the +# first number that means something is the one shipped alongside the first consumer. PROFILE_SCHEMA_VERSION = "1.0" @@ -46,7 +60,12 @@ class Evidence(BaseModel): """ kind: str = Field( - description="column_name | column_dtype | content_probe | split_name | file_name | card_metadata", + description=( + "column_name | column_dtype | content_probe | split_name | file_name | card_metadata | " + "user_hint | error — `user_hint` for a caller-supplied column role the data could not " + "support, and `error` for when a detector could not run at all, so an absent finding is " + "distinguishable from a finding of absence." + ), ) detail: str = Field( description="Self-describing evidence, e.g. \"answer matches '#### ' in 100% of 1024 sampled rows\".", @@ -77,7 +96,25 @@ class PartitionClassification(BaseModel): """ modality: str = Field(default="text", description="text | image_text | audio_text | ...") - dataset_type: str = Field(description="Dataset-type vocabulary (prompt_completion, preference_pair, ...).") + dataset_type: str = Field( + description=( + "Dataset-type vocabulary (prompt_completion, preference_pair, ...). A SUMMARY, not the " + "basis for a decision — it is the most specific single structure the roles satisfy, and a " + "dataset routinely satisfies several. The `semantic_role` markers are what a consumer " + "should match on; `candidates` lists everything this one is a projection of." + ), + ) + candidates: list[str] = Field( + default_factory=list, + description=( + "Every dataset type the assigned roles satisfy, most specific first, so " + "`candidates[0] == dataset_type`. prompt + completion + score + label is genuinely both " + "scored_response and unpaired_preference; reporting only the first made rule order an " + "invisible tie-break and hid that the data supports more than one use. Deliberately not a " + 'capability list ("supports DPO") — trainer requirements shift and differ per framework, ' + "so that mapping belongs in the consumer, computed from the roles." + ), + ) 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( @@ -97,7 +134,21 @@ class PartitionClassification(BaseModel): class Quantiles(BaseModel): - """A per-row distribution summary. p99 = long-tail sequence-length signal; max = hard cap.""" + """A per-row distribution summary. p99 = long-tail sequence-length signal; max = hard cap. + + The shape is the point, not the precision. Mean and max cannot tell "uniformly medium-length" + apart from "mostly short with a long tail", and those call for opposite sequence budgets — set + one from `max` and most of the memory is wasted, set it from the mean and the tail is silently + truncated. Reading p50 against p99 is what answers it. + + **p50 / p95 / p99 are estimates, within a couple of percent.** They are read off counters bucketed + by magnitude rather than from the lengths themselves, which is what keeps the profiler's memory + flat in rows. Every row is counted, so the *rank* is exact; only the value is rounded, and it is + rounded to a bound that does not grow with the dataset. That is the cheap error to accept here, + because whoever reads these rounds to a power of two anyway. + + **`max` is exact**, always, and is the only number here safe to treat as a hard bound. + """ p50: int p95: int @@ -123,7 +174,9 @@ class MessageStats(BaseModel): '"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." + "would hide exactly what a consumer needs to see before choosing a chat template. " + "Bounded: this is fed straight from row content, and a column with more distinct roles " + "than fit here is not a chat column, which the first few dozen already say." ), ) ends_with_assistant_rate: float = Field( @@ -144,8 +197,21 @@ class NumericStats(BaseModel): class TextQuality(BaseModel): - """Cheap, single-pass corruption signals for a text column. Flags training-wrecking data, not - toxicity / PII. + """Corruption signals for a text column. Flags training-wrecking data, not toxicity / PII. + + **Estimates, not counts**, and the only measurements in the profile that are. These three are + all the per-character work there is — every other statistic is O(1) per row, and the content + probes are literal searches costing a fraction of these — so scanning every row of a large + column costs more than the entire rest of the profile. They are also ratios, which a sample of + tens of thousands of rows pins down far past the precision anyone reads them to. Bounding them + is what makes reading every row of a dataset affordable. + + The sample is contiguous blocks, spaced evenly across the column: deterministic, so two runs over + the same bytes agree, and spread rather than taken from the head, so a sorted shard does not + decide the answer. Blocks rather than every n-th row because an even step aliases against + periodic data — a set that round-robins over sources, or carries k responses per prompt, is + periodic by construction, and a step sharing a factor with that period samples one phase and + only that phase. A column smaller than the bound is measured in full. """ whitespace_ratio: float = Field(ge=0.0, le=1.0, description="Padding / bad scraping.") @@ -181,11 +247,14 @@ class FeatureSchema(BaseModel): "message struct's `role` key." ), ) - fixed_length: int | None = Field( + semantic_role_source: str | 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." + "Where `semantic_role` came from: detected | declared. A declared role was supplied by the " + "caller and only accepted because the dtype could carry it; a detected one was inferred from " + "the column name. Kept as a field rather than left to evidence prose because the distinction " + "is per-column and actionable — a UI renders a declared role as confirmed and a detected one " + "as a suggestion to correct." ), ) fields: list[FeatureSchema] | None = Field(default=None, description="dtype == struct: named child fields.") @@ -196,9 +265,9 @@ 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. + forward compatibility. Tying `fields` / `items` 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") @@ -206,22 +275,41 @@ def _fields_and_items_are_exclusive(self) -> FeatureSchema: class CategoricalStats(BaseModel): - """Cardinality signals for string / int columns. + """The vocabulary of a column that has one. + + Present only when the column really is a bounded controlled vocabulary. Absent otherwise, and + the absence *is* the claim: this column is not a vocabulary. - ``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. + It used to be a general cardinality count on every string and numeric column. Counting distinct + values exactly means *retaining* them, and for a column of prompts the set of distinct values is + the column. What that bought was a reading of "9,954 distinct in 10,000 rows", which says free + text, which ``semantic_role`` and the length quantiles already said for free. Nothing read it + either: the only consumers of the number are a ``<= 2`` test that confirms a binary label and + the ``<= 32`` gate on ``values`` below. + + The values themselves ARE row content, so they appear only for a column whose detected role makes + it a controlled vocabulary — the assert-only-what-was-proven rule applied to the one place the + profiler would otherwise leak the data it is describing. """ distinct_count: int = Field( description=( - "Distinct values among scanned rows: ~=rows_scanned -> id-like; a small bounded set " - "corroborates score / category roles." + "How many distinct values the vocabulary holds. Exact, with no cap to have silently hit: " + "this model is built only for a column that stayed inside the vocabulary bounds all the " + "way through, so there is nothing to caveat. A small bounded set corroborates score / " + "category roles, and `<= 2` is what confirms a binary preference label." ), ) values: list[str] | None = Field( default=None, - description="The proven enumeration; only when the scan was exhaustive and distinct_count <= 32.", + description=( + "The observed values, present only when this column's `semantic_role` marks it a controlled " + "vocabulary (label | provenance | meta | rank) and it holds at most 32 of them. Cardinality " + "alone cannot be the gate: it inverts on small data, where every column holds few distinct " + "values — free text included — so a three-row dataset had its prompts stored verbatim. A role " + "says what a column *is*, at any size. Read `PartitionProfile.rows_complete` to know whether " + "this is the whole vocabulary or only what the sampled rows showed." + ), ) @@ -231,37 +319,38 @@ class ColumnStats(BaseModel): 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. + with one role-gated exception: ``categorical.values``, and only for a column whose role makes it + a controlled vocabulary rather than free text that happens to repeat. """ 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") + categorical: CategoricalStats | None = Field( + default=None, + description="Present only when the column is a bounded controlled vocabulary; absence means it is not one.", + ) quality: TextQuality | None = Field(default=None, description="dtype == string: corruption signals") -class FileRecord(BaseModel): - """One physical file, measured. +class FileError(BaseModel): + """A file the profiler could not fully use, and why. - 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. + Only failures are enumerated. Healthy files are counted (``SplitProfile.num_files``), because a + per-file record for each of them scaled the profile with shard count while telling a reader + nothing: at 512 shards those records were 95% of the payload and every one of them said "this + file was fine". Problems are the part worth naming, and there are few. """ path: str = Field(description="Relative path within the fileset.") - size_bytes: int - checksum: str | None = Field( - default=None, + error: str = Field( 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." + "Why this file was not fully read: unreadable, corrupt, partially parsed, or in a format " + "with no reader. A file that was read cleanly never appears here, so the absence of a path " + "is itself the claim that it was fine." ), ) - num_rows: int | None = Field( - default=None, - description="Exact only (parquet footer / exhaustive scan), else None.", - ) class SplitProfile(BaseModel): @@ -289,33 +378,90 @@ class SplitProfile(BaseModel): "train, with the variant's intent kept in `name`." ), ) - files: list[FileRecord] = Field( + data_files: str | None = Field( + default=None, + description=( + "A glob selecting exactly this split's files, relative to the fileset root: \"helpsteer2/" + 'train*.parquet". Gives the files back their addressability without giving back the ' + "per-file manifest — one pattern per split, whatever the shard count — so a consumer can " + "hand a reader the files of one split without listing the fileset and re-deriving which " + "shards belong where. Named for HF card front-matter's `configs[].data_files`, which is " + "the declared form of this same claim and, once cards are parsed, the thing that will " + "replace this inference rather than sit beside it in a second vocabulary.\n\n" + "`*` spans any run of characters except `/` — the one reading shared by shell globs, " + "Python's glob, fsspec and HF — so the pattern means the same thing wherever it is pasted. " + "`**` is never emitted, because its meaning is not shared.\n\n" + "None when no single pattern selects these files and nothing else (shards spread across " + "subdirectories, say). Never approximate: a pattern is emitted only after being matched " + "back against every file in the fileset and found to select this split exactly. A glob is " + "an instruction to go read files, so a near miss is not a rougher answer — it silently " + "pulls a README, or a neighbouring split's shards, into a training set." + ), + ) + num_files: int = Field( + default=0, 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." + "How many files resolved into this split. Partitioning is exhaustive and disjoint — each file " + "of the partition lands in exactly one split — so these sum to the partition's file count. " + "A count rather than a list: the paths of healthy shards are the one part of a profile that " + "grows without bound and informs no decision." + ), + ) + size_bytes: int = Field( + default=0, + description=( + "On-disk bytes of this split's files, summed. Answers whether the data fits wherever the " + "reader means to put it — the first question asked of an unfamiliar dataset, and one a row " + "count cannot answer, since a row ranges from an integer score to a reasoning trace. " + "Unlike `num_examples` this is never None: it comes from the file listing rather than from " + "reading, so a file that failed mid-read still contributes its size. Bytes as stored — " + "compressed, and several times this once decoded into memory. Covers only files a " + "partition grouped; a format with no reader never reaches a split, so weigh the whole " + "fileset with `SamplingInfo.bytes_present`." ), ) 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." + "Rows in this split, counting every one of its files whether or not that file's rows were " + "read. Always " + "exact — summed from parquet footers or from files read to their end — and None the moment any " + "one file's count is unknown. Never an estimate, so it carries no accuracy caveat: a capped run " + "still reports the true total whenever the footers knew it." ), ) 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". + """A file-group sharing one row schema and one source directory (roughly an HF config). 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") + name: str = Field( + default="", + description=( + "Identifies this partition, and unique within a profile. It is the path prefix its files " + 'share within the fileset: a top-level directory, or "" when they sit at the fileset ' + "root. Empty is a safe sentinel precisely because no directory can be named it, so " + "root-level files stay distinct from a directory literally called 'default'. Once card " + "front-matter is parsed, a declared config name populates this field instead — the same " + 'claim from a better source. For display, read it as `name or "default"`: storing that ' + "default was a lossy habit, because a lone partition under `data/` then reported " + '"default" and threw away the only thing identifying it.' + ), + ) + file_formats: list[str] = Field( + default_factory=list, + description=( + "The distinct formats this partition's files are in, sorted — normally exactly one, and " + "more than one when a stray .jsonl sits beside .parquet shards. That is noise, not a " + "second dataset, so it stays in this partition and shows up here rather than splitting it. " + "jsonl | parquet are read today; csv | arrow are reserved vocabulary the profiler cannot " + "read yet and reports on `DatasetProfile.file_errors` instead." + ), + ) 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).", @@ -327,6 +473,21 @@ class PartitionProfile(BaseModel): "omitted); keys are a subset of the top-level `features` names." ), ) + rows_complete: bool = Field( + description=( + "True => every row of every file in THIS partition was read. Only then can a consumer " + "assert enum / required in a bridged JSON Schema, or read a verifiability coverage of " + "1.0 as literal.\n\n" + "Named for what it measures. It was `stats_complete`, which promised more than it " + "delivered: `Quantiles` and `TextQuality` are estimates by construction however much was " + "read, each bounded for the cost reasons its own docstring gives. Whether a number is " + "exact is a property of that number, and every one of them says so; this says only " + "whether anything was missed on the way in.\n\n" + "Scoped to the partition because that is where it is decided — a corrupt shard in one " + "partition says nothing about the measurements in another, and a fileset-wide flag " + "quietly downgraded every partition to the worst one." + ), + ) classification: PartitionClassification @model_validator(mode="after") @@ -343,50 +504,87 @@ def _stats_keys_subset_of_features(self) -> PartitionProfile: class SamplingInfo(BaseModel): - """How much of the data the profile is based on. + """How much of the data the profile is based on — coverage, stated as numbers. + + Deliberately carries no ``exhaustive`` flag. That bit was answering two questions at once: "are + these measurements facts or estimates?", which is a property of each measurement and is now + stated by each of them, and "did I see all the data?", which is this block's job and needs + numerators and denominators rather than a boolean. It also folded together causes that call for + different people to act — a short read is the caller's choice, a corrupt shard is the data + owner's problem, and a missing reader is ours. - 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). + Nor does it record the caller's row limit. Reading everything is now the default and costs what + reading some of it costs, so a short read is unusual — and when it happens ``rows_scanned`` + against ``rows_present`` already says so. *Why* is not the profile's business: a limit is an + input, and the only other cause is a file that failed, which is named on ``file_errors``. + + The dataset-wide question is still one expression away, and still says which half failed:: + + all(p.rows_complete for p in profile.partitions) and not profile.file_errors """ - exhaustive: bool = Field(description="True => every row of every file was parsed.") - strategy: str = Field( + rows_scanned: int = Field(description="Total rows actually parsed across all files.") + rows_present: int | None = Field( + default=None, 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." + "How many rows the fileset holds, scanned or not — the denominator `rows_scanned` is a " + "fraction of. Populated whenever every file's count is *known*, regardless of how much was " + "read: a row-capped run over parquet still knows its totals from the footers, and that is " + "exactly when the ratio carries information. None means at least one file's count is " + "unknown — never zero, never an estimate." ), ) - rows_scanned: int = Field(description="Total rows actually parsed across all files.") - rows_total: int | None = Field( - default=None, + files_read: int = Field( 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 actually opened and read from. A count, not a list -- the paths of healthy " + "shards are the one part of a profile that grows without bound and informs no " + "decision; `SplitProfile.num_files` counts them per split, and only the ones that " + "went wrong are named, on `DatasetProfile.file_errors`." + ) + ) + files_present: int = Field( + description=( + "Data files the fileset holds, whether or not this run could read them — the denominator " + "`files_read` is a fraction of. Includes files in formats with no reader, since those are " + "data that went unprofiled (they are named on `DatasetProfile.file_errors`). A README " + "is not data and is counted nowhere. Every readable file should be opened, since " + "head-sampling a *subset of files* hides columns that appear only in later shards, so expect " + "these two to match until scale forces file-level sampling." ), ) - files_scanned: int = Field( + bytes_present: int = Field( + default=0, 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." + "On-disk bytes of every data file the fileset holds, whether or not this run could read it " + "— the size of the dataset as it sits, independent of how much was profiled. Redundant " + "with the sum over `SplitProfile.size_bytes` exactly when nothing failed, and load-bearing " + "when something did: a file in a format with no reader never reaches a partition, so a " + "directory of .csv shards beside one .parquet would otherwise weigh in at the parquet " + "alone. Same reason `files_present` is kept alongside the per-split counts — a denominator " + "stops being derivable the moment coverage is partial, which is the only time it is read." ), ) - 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.""" + """The machine-owned dataset profile — the root of the stored contract. + + Deliberately carries no staleness marker, and no per-file manifest to reconstruct one from. A + stored digest would freeze "which files count as inputs" into the data at write time, and that + judgment moves: once card front-matter drives split declaration, ``README.md`` becomes an input. + Changing the rule would then invalidate every stored profile at once, with no way to tell a real + change from a definition change. + + So a profile says when it was made and nothing about whether it still holds. ``created_at`` is + the whole of it. That is deliberate while profiling is user-triggered and nothing consumes + freshness; when something does, the cheap primitive is a fileset version token from the storage + backend, which costs no listing and freezes no policy — not a manifest reconstructed here. + """ 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, @@ -396,6 +594,17 @@ class DatasetProfile(BaseModel): partitions: list[PartitionProfile] = Field( description="Single partition in the common homogeneous case; there is no fileset-level rollup.", ) + file_errors: list[FileError] = Field( + default_factory=list, + description=( + "Every file the profiler could not fully use, from anywhere in the fileset: a format with " + "no reader, a corrupt shard, a partially parsed one. Reporting them is what keeps a " + "directory of .csv shards from profiling as an exhaustively scanned *empty* dataset, " + "indistinguishable from one that really is empty. One list rather than two, because " + '"a file I could not use" is the same finding whether or not a partition managed to group ' + 'it first, and a reader asking "did anything go wrong?" should not have to look twice.' + ), + ) # Resolve the recursive FeatureSchema self-reference (deferred by `from __future__ import annotations`). diff --git a/packages/nemo_platform_plugin/tests/files/test_dataset_profile.py b/packages/nemo_platform_plugin/tests/files/test_dataset_profile.py index ca7ba48484..ed2e8359c0 100644 --- a/packages/nemo_platform_plugin/tests/files/test_dataset_profile.py +++ b/packages/nemo_platform_plugin/tests/files/test_dataset_profile.py @@ -19,7 +19,7 @@ DatasetProfile, Evidence, FeatureSchema, - FileRecord, + FileError, MessageStats, PartitionClassification, PartitionProfile, @@ -33,26 +33,23 @@ # --- 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} +sampling: {rows_scanned: 2112, rows_present: 3201061, + files_read: 33, files_present: 33, bytes_present: 31821490182} partitions: - - name: default - file_format: parquet + - name: "" + file_formats: [parquet] + rows_complete: false 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}]} + - {name: train, canonical: train, num_examples: 3200861, num_files: 32, + size_bytes: 31819412254, data_files: 'train*.parquet'} + - {name: test, canonical: test, num_examples: 200, num_files: 1, + size_bytes: 2077928, data_files: 'test*.parquet'} features: - - {name: prompt, dtype: messages, semantic_role: prompt, + - {name: prompt, dtype: messages, semantic_role: prompt, semantic_role_source: detected, items: {dtype: struct, fields: [{name: role, dtype: string}, {name: content, dtype: string}]}} - - {name: completion, dtype: messages, semantic_role: completion, + - {name: completion, dtype: messages, semantic_role: completion, semantic_role_source: detected, 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}, @@ -62,6 +59,7 @@ classification: modality: text dataset_type: prompt_completion + candidates: [prompt_completion] format: conversational prompt_form: explicit verifiability: @@ -76,27 +74,25 @@ # --- 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} +sampling: {rows_scanned: 1024, rows_present: 46189, + files_read: 2, files_present: 2, bytes_present: 27055195} partitions: - - name: default - file_format: parquet + - name: "" + file_formats: [parquet] + rows_complete: false 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}]} + - {name: train, canonical: train, num_examples: 43835, num_files: 1, + size_bytes: 25670988, data_files: 'train*.parquet'} + - {name: test, canonical: test, num_examples: 2354, num_files: 1, + size_bytes: 1384207, data_files: 'test*.parquet'} features: - - {name: prompt, dtype: messages, semantic_role: prompt, + - {name: prompt, dtype: messages, semantic_role: prompt, semantic_role_source: detected, items: {dtype: struct, fields: [{name: role, dtype: string}, {name: content, dtype: string}]}} - - {name: chosen, dtype: messages, semantic_role: chosen, + - {name: chosen, dtype: messages, semantic_role: chosen, semantic_role_source: detected, items: {dtype: struct, fields: [{name: role, dtype: string}, {name: content, dtype: string}]}} - - {name: rejected, dtype: messages, semantic_role: rejected, + - {name: rejected, dtype: messages, semantic_role: rejected, semantic_role_source: detected, 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}, @@ -108,6 +104,7 @@ classification: modality: text dataset_type: preference_pair + candidates: [preference_pair] format: conversational prompt_form: explicit evidence: @@ -118,29 +115,27 @@ # --- 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} +sampling: {rows_scanned: 1024, rows_present: 21362, + files_read: 2, files_present: 2, bytes_present: 19459677} partitions: - - name: default - file_format: parquet + - name: "" + file_formats: [parquet] + rows_complete: false 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}]} + - {name: train, canonical: train, num_examples: 20324, num_files: 1, + size_bytes: 18495985, data_files: 'train*.parquet'} + - {name: validation, canonical: validation, num_examples: 1038, num_files: 1, + size_bytes: 963692, data_files: 'validation*.parquet'} 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} + - {name: prompt, dtype: string, semantic_role: prompt, semantic_role_source: detected} + - {name: response, dtype: string, semantic_role: completion, semantic_role_source: detected} + - {name: helpfulness, dtype: int64, semantic_role: score, semantic_role_source: detected} + - {name: correctness, dtype: int64, semantic_role: score, semantic_role_source: detected} + - {name: coherence, dtype: int64, semantic_role: score, semantic_role_source: detected} + - {name: complexity, dtype: int64, semantic_role: score, semantic_role_source: detected} + - {name: verbosity, dtype: int64, semantic_role: score, semantic_role_source: detected} 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}} @@ -154,6 +149,7 @@ classification: modality: text dataset_type: scored_response + candidates: [scored_response, prompt_completion] format: standard prompt_form: explicit evidence: @@ -172,29 +168,25 @@ 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, + rows_present=2048, + files_read=2, + files_present=2, + row_budget=1024, ), partitions=[ PartitionProfile( - file_format="parquet", + file_formats=["parquet"], + rows_complete=False, splits=[ SplitProfile( name="train", canonical="train", num_examples=2048, - files=[ - FileRecord(path="train-00000.parquet", size_bytes=123, checksum="sha256:ab", num_rows=2048) - ], + num_files=1, ) ], features=[ @@ -207,6 +199,7 @@ def _build_profile() -> DatasetProfile: }, classification=PartitionClassification( dataset_type="prompt_completion", + candidates=["prompt_completion"], format="standard", prompt_form="explicit", verifiability=Verifiability( @@ -237,7 +230,8 @@ 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" + # All three ship their shards at the fileset root, so the shared path prefix is empty. + assert profile.partitions[0].name == "" # Round-trip through JSON is lossless. assert DatasetProfile.model_validate_json(profile.model_dump_json()) == profile @@ -262,10 +256,52 @@ def test_openmathreasoning_locks_contract_shape(): assert part.stats["completion"].messages.roles_seen == ["assistant"] +@pytest.mark.parametrize("name", sorted(FIXTURES)) +def test_split_sizes_account_for_the_whole_fileset(name): + """On a clean profile the splits weigh the whole fileset, so `bytes_present` is the same number + reached without going through partitions. That redundancy is the point: it is what lets the + figure survive a file no partition could group.""" + profile = DatasetProfile.model_validate(yaml.safe_load(FIXTURES[name])) + assert not profile.file_errors + from_splits = sum(split.size_bytes for part in profile.partitions for split in part.splits) + assert from_splits == profile.sampling.bytes_present + + +@pytest.mark.parametrize("name", sorted(FIXTURES)) +def test_split_globs_are_one_pattern_each_and_never_cross_a_directory(name): + """`data_files` is a single pattern, not a manifest, so it cannot reintroduce the per-file growth + the split-level counts exist to avoid. `**` is never emitted, because its meaning is not shared + across glob implementations.""" + profile = DatasetProfile.model_validate(yaml.safe_load(FIXTURES[name])) + for part in profile.partitions: + for split in part.splits: + assert isinstance(split.data_files, str) + assert "**" not in split.data_files + + +def test_a_split_with_no_expressible_pattern_says_so(): + """None is a first-class answer: shards spread across subdirectories need `**` to cover, and a + pattern that resolves differently in the reader than in the profiler is worse than none.""" + split = SplitProfile(name="train", num_files=2) + assert split.data_files is None + + +def test_a_split_weighs_something_even_when_its_row_count_does_not(): + """Size is read off the file listing and a row count off the data, so they go unknown + independently — `num_examples` is None-able and `size_bytes` is not.""" + split = SplitProfile(name="train", num_files=3, size_bytes=4096) + assert split.num_examples is None + assert split.size_bytes == 4096 + + 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" + # A scored prompt/completion set is also a plain prompt_completion set. `dataset_type` is the + # most specific reading; `candidates` is what the same columns otherwise support. + assert part.classification.candidates == ["scored_response", "prompt_completion"] + assert part.classification.candidates[0] == part.classification.dataset_type assert part.classification.format == "standard" # Absence of a verifiability object *is* the "not verifiable" claim. assert part.classification.verifiability is None @@ -329,8 +365,8 @@ def test_container_shape_is_not_pinned_to_known_dtypes(): 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 + future_tensor = FeatureSchema(name="embedding", dtype="tensor", items=FeatureSchema(dtype="float32")) + assert future_tensor.items is not None and future_tensor.items.dtype == "float32" def test_unknown_fields_are_ignored_for_forward_compat(): @@ -342,6 +378,42 @@ def test_unknown_fields_are_ignored_for_forward_compat(): assert profile.partitions[0].classification.dataset_type == "scored_response" +def test_file_errors_are_the_only_channel_for_trouble(): + # Healthy files are counted, never listed, so a reader asking "did anything go wrong?" reads one + # list whose length is the number of problems -- not one that grows with the shard count and is + # 95% success records at scale. + doc = yaml.safe_load(HELPSTEER2) + doc["file_errors"] = [ + {"path": "train-00007-of-00032.parquet", "error": "ArrowInvalid: not a parquet file"}, + {"path": "notes.csv", "error": "no reader for '.csv' files"}, + ] + profile = DatasetProfile.model_validate(doc) + + assert [e.path for e in profile.file_errors] == ["train-00007-of-00032.parquet", "notes.csv"] + # A shard the profiler could not read and a format it has no reader for are the same finding, + # and land in the same place whether or not a partition managed to group the file first. + assert all(isinstance(e, FileError) and e.error for e in profile.file_errors) + assert DatasetProfile.model_validate_json(profile.model_dump_json()) == profile + + +def test_a_clean_profile_names_no_files_at_all(): + profile = DatasetProfile.model_validate(yaml.safe_load(HELPSTEER2)) + assert profile.file_errors == [] + assert [s.num_files for s in profile.partitions[0].splits] == [1, 1] + + +def test_a_profile_written_before_the_digest_was_dropped_still_loads(): + # `content_digest` was removed rather than repaired: it froze "which files count as inputs" into + # stored data at write time, and that judgment moves. Profiles already written with it have to + # keep loading, or removing it would break every one of them at once — the very failure mode the + # removal exists to avoid. + doc = yaml.safe_load(HELPSTEER2) + doc["content_digest"] = "sha256:7be1c0ffee" + profile = DatasetProfile.model_validate(doc) + assert not hasattr(profile, "content_digest") + 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( diff --git a/plugins/nemo-datasets/pyproject.toml b/plugins/nemo-datasets/pyproject.toml new file mode 100644 index 0000000000..0ba323e177 --- /dev/null +++ b/plugins/nemo-datasets/pyproject.toml @@ -0,0 +1,30 @@ +[project] +name = "nemo-datasets-plugin" +description = "Dataset profiler for NeMo Platform filesets." +requires-python = ">=3.11,<3.15" +dependencies = [ + "nemo-platform-plugin", + "nemo-platform-sdk", + "pyarrow>=19.0.1", + "pydantic>=2.10.3", +] +version = "0.1.0" + +[tool.uv.sources] +nemo-platform-plugin = { workspace = true } +nemo-platform-sdk = { workspace = true } + +# Deliberately contributes no `nemo.cli` entry point. The profiler runs as a job task +# (`python -m nemo_datasets_plugin.tasks.profile`), which is invoked by the platform rather than +# typed by a user, so its inputs can keep moving while the feature is new. + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build.targets.wheel] +packages = ["src/nemo_datasets_plugin"] + +[tool.pytest.ini_options] +testpaths = ["tests"] +pythonpath = ["src"] diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/classify.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/classify.py new file mode 100644 index 0000000000..146b7fcd61 --- /dev/null +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/classify.py @@ -0,0 +1,446 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Classification: assign column roles, resolve the format/prompt-form axes, and pick a dataset type. + +Roles are inferred from column names, gated by dtype, and stacked onto the feature nodes as +``semantic_role`` markers. The dataset type is the most specific structure the assigned roles +satisfy. + +Content probes are *measured* in :mod:`stats` over every column; this module only interprets the +counts. Roles still order that interpretation — a column known to be the ground truth is a better +answer than one that merely looks like it — but they no longer gate it, so a dataset whose columns +carry unrecognized names keeps whatever its content proves. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +from nemo_datasets_plugin.profiler.stats import ColumnProbes +from nemo_platform_plugin.files.dataset_profile import ( + ColumnStats, + Evidence, + FeatureSchema, + PartitionClassification, + Verifiability, +) + +# Column-name aliases -> role. Score is handled separately (name alias + numeric dtype gate). +_ALIAS_ROLES = { + "prompt": "prompt", + "question": "prompt", + "instruction": "prompt", + "problem": "prompt", + "query": "prompt", + "context": "context", + "input": "context", + "passage": "context", + "document": "context", + "system": "system", + "system_prompt": "system", + "response": "completion", + "output": "completion", + "answer": "completion", + "completion": "completion", + "solution": "completion", + "messages": "messages", + "conversation": "messages", + "conversations": "messages", + "chosen": "chosen", + "rejected": "rejected", + "label": "label", + "rank": "rank", + "ground_truth": "ground_truth", + "reference_answer": "ground_truth", + "verification_info": "ground_truth", + "test_cases": "ground_truth", + "completions": "stepwise_completions", + "labels": "stepwise_labels", + "tools": "tools", + "image": "image", + "images": "image", + "id": "id", + "prompt_id": "id", + "source": "provenance", + "dataset": "provenance", + "model": "provenance", + "category": "meta", +} + +_SCORE_ALIASES = { + "score", + "score_chosen", + "score_rejected", + "helpfulness", + "correctness", + "coherence", + "complexity", + "verbosity", + "quality", + "rating", + "reward", +} + +_TEXT_DTYPES = {"string", "messages"} +# A verification target is naturally a container: test_cases (list), verification_info (struct), or a +# plain string answer — but never a bare scalar/number, which is far more likely a label or score. +_GROUND_TRUTH_DTYPES = {"string", "messages", "list", "struct", "json"} +# Roles whose string-vs-messages dtype decides the format axis. +_SHAPE_ROLES = {"prompt", "completion", "chosen", "rejected", "messages"} + + +def _is_numeric(dtype: str) -> bool: + return dtype.startswith(("int", "uint", "float")) + + +def _is_binary(column: ColumnStats | None) -> bool: + """Whether a column was observed to hold at most two distinct values.""" + return column is not None and column.categorical is not None and column.categorical.distinct_count <= 2 + + +def _is_label_column(feature: FeatureSchema, stats: dict[str, ColumnStats]) -> bool: + """Whether a column named ``label`` really carries a binary preference label. + + A bool says so outright. An integer is the more common on-disk encoding (0/1, as KTO-style sets + ship it), but only when the observed values really are binary — a wider integer range is a class + index or a rating, which is a different claim, so it stays unroled. + """ + if feature.dtype == "bool": + return True + return _is_numeric(feature.dtype) and _is_binary(stats.get(feature.name)) + + +def _dtype_allows(feature: FeatureSchema, role: str, stats: dict[str, ColumnStats]) -> bool: + """Whether this column's dtype can carry ``role`` at all. + + Applied to detected *and* declared roles alike. A hint says which column, not what the data is: + without this, one typo (``{"score": "prompt"}`` on an int column) would silently produce a + nonsense classification and the profile would become a place to store mistakes. + """ + dtype = feature.dtype + if role == "score" or role == "rank": + return _is_numeric(dtype) + if role == "label": + return _is_label_column(feature, stats) + if role == "messages": + return dtype == "messages" + if role in {"prompt", "completion", "chosen", "rejected", "context", "system"}: + return dtype in _TEXT_DTYPES + if role == "ground_truth": + return dtype in _GROUND_TRUTH_DTYPES + if role in {"stepwise_completions", "stepwise_labels"}: + return dtype == "list" + return True # id / provenance / meta / tools / image carry no dtype constraint + + +def _role_for(feature: FeatureSchema, stats: dict[str, ColumnStats]) -> str | None: + """The role this column's *name* implies, if the dtype does not contradict it.""" + name = feature.name.lower() + if name in _SCORE_ALIASES and _is_numeric(feature.dtype): + return "score" + role = _ALIAS_ROLES.get(name) + if role is None: + return None + return role if _dtype_allows(feature, role, stats) else None + + +def _assign_roles( + features: list[FeatureSchema], stats: dict[str, ColumnStats], column_roles: dict[str, str] +) -> list[Evidence]: + """Stack roles onto ``features`` in place; return evidence for any hint the data could not support. + + A declared role wins over the name-alias table — the caller knows their schema and the table is + ~35 English names — but it still has to pass the dtype gate, and a rejected hint is reported + rather than dropped. Silence is what made the alias table's misses so expensive in the first place. + """ + rejected: list[Evidence] = [] + for feature in features: + declared = column_roles.get(feature.name) + if declared is not None: + if _dtype_allows(feature, declared, stats): + feature.semantic_role = declared + feature.semantic_role_source = "declared" + continue + rejected.append( + Evidence( + kind="user_hint", + detail=( + f"hint '{feature.name} -> {declared}' rejected: a {feature.dtype} column " + f"cannot carry that role; falling back to detection" + ), + ) + ) + role = _role_for(feature, stats) + if role is not None: + feature.semantic_role = role + feature.semantic_role_source = "detected" + return rejected + + +def _detect_modality(features: list[FeatureSchema]) -> str: + if any(feature.semantic_role == "image" or feature.dtype == "image" for feature in features): + return "image_text" + return "text" + + +def _detect_format(features: list[FeatureSchema]) -> str | None: + dtypes = {feature.dtype for feature in features if feature.semantic_role in _SHAPE_ROLES} + has_messages = "messages" in dtypes + has_string = "string" in dtypes + if has_messages and has_string: + return "mixed" + if has_messages: + return "conversational" + if has_string: + return "standard" + return None + + +def _detect_prompt_form(roles: set[str]) -> str | None: + if "prompt" in roles: + return "explicit" + if roles & {"chosen", "rejected", "completion"}: + return "implicit" # a prompt exists but is embedded in the completions + return "n/a" + + +def _messages_stats(features: list[FeatureSchema], stats: dict[str, ColumnStats]): + for feature in features: + if feature.semantic_role == "messages": + column = stats.get(feature.name) + if column is not None: + return column.messages + return None + + +def _detect_types(features: list[FeatureSchema], stats: dict[str, ColumnStats]) -> list[str]: + """Every dataset type the assigned roles satisfy, most specific first. + + The chain is ordered by specificity, so the head is the best single answer and the tail is + structures the same columns *also* satisfy. Returning only the head made rule order an invisible + tie-break: prompt + completion + score + label is genuinely both scored_response and + unpaired_preference, and which one a consumer saw depended on line numbers. + """ + roles = {feature.semantic_role for feature in features if feature.semantic_role} + targets = roles & {"completion", "chosen", "rejected", "stepwise_completions"} + candidates: list[str] = [] + + def has(*required: str) -> bool: + return all(role in roles for role in required) + + if has("prompt", "stepwise_completions", "stepwise_labels"): + candidates.append("stepwise_supervision") + if has("chosen", "rejected"): + candidates.append("preference_pair") + if has("prompt", "completion", "score"): + candidates.append("scored_response") + if has("prompt", "completion", "label"): + candidates.append("unpaired_preference") + # `rank` is only a dataset type alongside something to rank. On its own it short-circuited every + # more specific structure above, so a stray numeric column named `rank` — or a ranked variant of + # a preference set — was enough to mislabel the dataset. + if has("rank") and targets: + candidates.append("ranked_responses") + if has("prompt", "completion"): + candidates.append("prompt_completion") + if "messages" in roles: + message_stats = _messages_stats(features, stats) + if message_stats is not None and message_stats.ends_with_assistant_rate < 0.5: + candidates.append("prompt_only") # a chat that ends on a user turn has no training target + else: + candidates.append("messages") + # A prompt with nothing to predict. Guarded on `targets` because with candidates collected rather + # than returned early, a prompt+completion set would otherwise claim prompt_only as well. + if "prompt" in roles and not targets and "prompt_only" not in candidates: + candidates.append("prompt_only") + if len(features) == 1 and features[0].dtype == "string" and features[0].semantic_role is None: + candidates.append("text") + return candidates or ["unknown"] + + +# --- interpreting the content probes -------------------------------------------------------------- + +# A verification target must cover at least this fraction of sampled rows to be asserted. Below it, +# a "hit" is noise -- e.g. one completion in thousands coincidentally ending in `#### ` does +# not make a dataset verifiable. Tune here; the coverage itself is still reported on the Verifiability. +_MIN_VERIFIABILITY_COVERAGE = 0.05 + + +def _pct(fraction: float) -> str: + return f"{round(fraction * 100)}%" + + +def _detect_verifiability(features: list[FeatureSchema], probes: dict[str, ColumnProbes]) -> Verifiability | None: + """The strongest verification target the probes found, if any clears the coverage floor. + + Each method wins only if it clears the floor; otherwise fall through to the next, so a sparse + ground_truth column yields to an extractable-answer signal instead of masking it. + """ + ground_truth = next((feature for feature in features if feature.semantic_role == "ground_truth"), None) + if ground_truth is not None: + probe = probes.get(ground_truth.name) + if probe is not None and probe.rows: + coverage = probe.non_empty / probe.rows + if coverage >= _MIN_VERIFIABILITY_COVERAGE: + detail = f"'{ground_truth.name}' present in {_pct(coverage)} of {probe.rows} sampled rows" + return Verifiability( + method="ground_truth_column", + coverage=coverage, + evidence=[Evidence(kind="content_probe", detail=detail)], + ) + + # A named completion is the authoritative place to look. Without one, take whichever column the + # probes found the strongest signal in and name it — the markers are a fact about that column + # whether or not its name happened to be in the alias table. + completion = next((feature for feature in features if feature.semantic_role == "completion"), None) + searched = [completion] if completion is not None else features + best_name: str | None = None + best_coverage = 0.0 + for feature in searched: + probe = probes.get(feature.name) + if probe is None or not probe.texts: + continue + coverage = probe.extractable_answer / probe.texts + if coverage > best_coverage: + best_name, best_coverage = feature.name, coverage + + if best_name is not None and best_coverage >= _MIN_VERIFIABILITY_COVERAGE: + sampled = probes[best_name].texts + detail = ( + f"'{best_name}' ends with an extractable answer (#### or \\boxed) in " + f"{_pct(best_coverage)} of {sampled} sampled rows" + ) + return Verifiability( + method="extractable_final_answer", + coverage=best_coverage, + evidence=[Evidence(kind="content_probe", detail=detail)], + ) + return None + + +# How much text two answers must open with in common before it reads as a shared prompt rather than +# a shared turn of phrase. Short enough that a one-line question counts, long enough that "I think +# that" does not. +_EMBEDDED_PROMPT_PREFIX_CHARS = 16 + + +def _common_prefix_len(left: str, right: str) -> int: + limit = min(len(left), len(right)) + index = 0 + while index < limit and left[index] == right[index]: + index += 1 + return index + + +@dataclass(frozen=True) +class PrefixPair: + """How often two columns of the same row opened with the same long run of text.""" + + pairs: int = 0 + shared: int = 0 + + +# The column names the alias table maps to the two sides of a preference pair. Looked up directly +# rather than resolved through roles, because this fold runs before classification has assigned any. +_CHOSEN_NAMES = tuple(name for name, role in _ALIAS_ROLES.items() if role == "chosen") +_REJECTED_NAMES = tuple(name for name, role in _ALIAS_ROLES.items() if role == "rejected") + + +class PrefixPairFold: + """The one probe that reads two columns against each other rather than each on its own. + + A preference set whose prompt is embedded in both answers shows up as a long shared prefix + between them, and no per-column measurement can see that. Being relational, it also cannot live + on a column accumulator, so it folds separately over the same batches. + + It takes no schema. Resolving by name straight off each row is what lets it run over a partition + whose columns are not known yet, which is every partition with no declared schema. Values that + are not text contribute nothing, so a `chosen` column that turns out to hold chat rather than + strings simply never counts a pair -- the same answer the dtype check used to give up front. + """ + + def __init__(self) -> None: + self._pairs = 0 + self._shared = 0 + + def update(self, rows: list[dict]) -> None: + for row in rows: + left = next((row.get(name) for name in _CHOSEN_NAMES if isinstance(row.get(name), str)), None) + right = next((row.get(name) for name in _REJECTED_NAMES if isinstance(row.get(name), str)), None) + if left is None or right is None: + continue + self._pairs += 1 + if _common_prefix_len(left, right) >= _EMBEDDED_PROMPT_PREFIX_CHARS: + self._shared += 1 + + def result(self) -> PrefixPair: + return PrefixPair(pairs=self._pairs, shared=self._shared) + + +def _implicit_prompt_evidence( + features: list[FeatureSchema], probes: dict[str, ColumnProbes], prefix_pair: PrefixPair +) -> Evidence | None: + targets = [f for f in features if f.semantic_role in {"chosen", "rejected", "completion"} and f.dtype == "string"] + counted = [probes[f.name] for f in targets if f.name in probes] + sampled = sum(probe.texts for probe in counted) + marked = sum(probe.transcript_marker for probe in counted) + if sampled and marked: + detail = f"embedded transcript markers in {_pct(marked / sampled)} of sampled completions - prompt is embedded" + return Evidence(kind="content_probe", detail=detail) + + if prefix_pair.pairs and prefix_pair.shared / prefix_pair.pairs >= 0.5: + rate = _pct(prefix_pair.shared / prefix_pair.pairs) + return Evidence( + kind="content_probe", + detail=f"chosen/rejected share a common prefix in {rate} of pairs - prompt is embedded", + ) + return None + + +def classify( + features: list[FeatureSchema], + stats: dict[str, ColumnStats], + *, + probes: dict[str, ColumnProbes] | None = None, + prefix_pair: PrefixPair | None = None, + column_roles: dict[str, str] | None = None, +) -> PartitionClassification: + """Assign roles onto ``features`` in place and return the partition's classification. + + ``probes`` are the per-column content measurements, and ``prefix_pair`` the one relational one. + Both are folded over the rows before this runs; nothing here reads a row, which is what lets a + partition be classified without ever having been materialised. Absent, each reads as "nothing + was measured", and role/axis/type inference is unaffected — that needs only schema and stats. + + ``column_roles`` maps a column name to a role the caller is asserting, taking precedence over + the name-alias table but still subject to the dtype gates. It exists because that table is ~35 + English names with no way to say "my `q` column is the prompt", and its misses are silent. + """ + probes = probes or {} + evidence = _assign_roles(features, stats, column_roles or {}) + roles = {feature.semantic_role for feature in features if feature.semantic_role} + candidates = _detect_types(features, stats) + dataset_type = candidates[0] + fmt = _detect_format(features) + prompt_form = _detect_prompt_form(roles) if dataset_type != "unknown" else None + + role_columns = [f"{feature.name} -> {feature.semantic_role}" for feature in features if feature.semantic_role] + if role_columns: + evidence.append(Evidence(kind="column_name", detail=f"columns matched roles: {', '.join(role_columns)}")) + if fmt is not None: + evidence.append(Evidence(kind="column_dtype", detail=f"{fmt} format from role column dtypes")) + if prompt_form == "implicit": + embedded = _implicit_prompt_evidence(features, probes, prefix_pair or PrefixPair()) + if embedded is not None: + evidence.append(embedded) + + return PartitionClassification( + modality=_detect_modality(features), + dataset_type=dataset_type, + candidates=candidates, + format=fmt, + prompt_form=prompt_form, + verifiability=_detect_verifiability(features, probes), + evidence=evidence, + ) diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/file_source.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/file_source.py new file mode 100644 index 0000000000..4500307e50 --- /dev/null +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/file_source.py @@ -0,0 +1,56 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The file-source seam. + +The profiler core reads dataset files only through a :class:`FileSource`, so it never touches a +storage API directly. :class:`LocalFileSource` covers a directory on disk; a ranged-read source over +the Files storage API is a later drop-in behind the same two methods. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import BinaryIO, Protocol + + +@dataclass(frozen=True) +class FileEntry: + """A file's identity — everything the profiler needs before it reads the contents.""" + + path: str # POSIX-style path relative to the source root + size_bytes: int + checksum: str | None = None # "sha256:..." when the source reports one; None otherwise + + +class FileSource(Protocol): + """Read-only access to a set of dataset files.""" + + def list_files(self) -> list[FileEntry]: + """Every file in the source, in a stable order.""" + ... + + def open(self, path: str) -> BinaryIO: + """A binary, seekable stream for one file (``path`` as returned by :meth:`list_files`).""" + ... + + +class LocalFileSource: + """A directory of dataset files on the local filesystem.""" + + def __init__(self, root: str | Path) -> None: + self._root = Path(root) + if not self._root.is_dir(): + raise NotADirectoryError(f"{self._root} is not a directory") + + def list_files(self) -> list[FileEntry]: + entries = [ + FileEntry(path=path.relative_to(self._root).as_posix(), size_bytes=path.stat().st_size) + for path in sorted(self._root.rglob("*")) + if path.is_file() + ] + return entries + + def open(self, path: str) -> BinaryIO: + return open(self._root / path, "rb") diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/partition.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/partition.py new file mode 100644 index 0000000000..3216ee4c13 --- /dev/null +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/partition.py @@ -0,0 +1,53 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Partition grouping. + +A partition is a group of files profiled as a unit. This stage groups by top-level directory; a +later stage refines partitions whose files turn out to disagree on column schema. +""" + +from __future__ import annotations + +from pathlib import PurePosixPath + +from nemo_datasets_plugin.profiler.file_source import FileEntry +from nemo_datasets_plugin.profiler.splits import is_split_directory + + +def _top_dir(path: str) -> str: + """The partition a file belongs to: its first path segment, or ``""`` for a root-level file. + + Empty is a usable name precisely because no directory can be called it, so root-level files never + collide with a directory literally named ``default``. + + A split-named top-level directory (``train/``, ``test/``) is deliberately *not* a partition + dimension. Grouping on it would split one dataset's train and test into unrelated partitions, + each deriving its own schema and classification — the exact structure `splits` exists to model. + Those files fall through to the same partition and are separated by :mod:`splits` instead. + """ + parts = PurePosixPath(path).parts + if len(parts) <= 1 or is_split_directory(parts[0]): + return "" + return parts[0] + + +def group_partitions(entries: list[FileEntry]) -> list[tuple[str, list[FileEntry]]]: + """Group files into (name, files) partitions by top-level directory, sorted by name. + + The name *is* the identity — the shared path prefix, not a display string derived from it. A lone + group under ``data/`` is named ``"data"``, not ``"default"``: reporting the latter discarded the + only thing identifying the partition, and left two partitions that could share a name. + + Files whose top-level directory is a split name (``train/``, ``test/``) group under ``""`` + alongside root-level files: those are one dataset's splits, not separate partitions. + """ + by_dir: dict[str, list[FileEntry]] = {} + for entry in entries: + by_dir.setdefault(_top_dir(entry.path), []).append(entry) + + if len(by_dir) == 1: + # A single group is one partition holding everything, keeping whatever directory it came from. + return [(next(iter(by_dir)), list(entries))] + + return sorted(by_dir.items()) diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py new file mode 100644 index 0000000000..1a37e8e319 --- /dev/null +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py @@ -0,0 +1,497 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The top-level profiling pipeline. + +``profile(source)`` lists the files behind a :class:`FileSource`, groups them into partitions and +splits, reads them, and assembles a ``DatasetProfile``. It produces the structural envelope — +partitions, splits with their counts and sizes, the files it could not use, and the sampling figures +— along with the derived row schema (``features``), per-column ``stats``, and the full +``classification`` (roles, format, prompt form, dataset type, and verifiability). + +Every file is opened — sampling a *subset of files* would hide columns that appear only in later +shards. + +Every partition is **folded**: batches are measured and let go, and nothing kept grows with the file. +An exhaustive read therefore costs what a short one costs, which is why reading everything is the +default. ``row_budget`` survives only as a way to ask for a shorter run. + +What a declared schema buys is not the fold but its sharpness. Parquet footers are read first, so +the columns are known before a row is parsed and each accumulator is chosen up front, and the exact +row count is known too, which is what lets the quality sample be spread across a column not yet +seen. + +Without one — line-delimited data — both wait for the data. Columns are created on first sight and +back-filled with the rows they were absent for, and each carries every shape at once until the last +row has gone by and the dtype resolves. That costs a deferred type per column and nothing else; it +does not cost a second pass, and it does not decide from a prefix. + +A caller who does ask for one gets a target rather than a ceiling. :data:`MIN_ROWS_PER_FILE` is the +floor every file is read to however thin its share gets, since one sampled below it cannot +contribute the columns it alone witnesses. That division outlived the memory problem it was invented +for: reading files in order until a total ran out would leave the later ones unopened, which is the +same coverage hole by another route. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from datetime import datetime, timezone +from pathlib import PurePosixPath + +import pyarrow as pa +from nemo_datasets_plugin.profiler.classify import PrefixPairFold, classify +from nemo_datasets_plugin.profiler.file_source import FileEntry, FileSource +from nemo_datasets_plugin.profiler.partition import group_partitions +from nemo_datasets_plugin.profiler.readers.base import ( + FilePreview, + detect_format, + get_reader, + is_unsupported_data, +) +from nemo_datasets_plugin.profiler.schema import MAX_COLUMNS, columns_were_capped, derive_features +from nemo_datasets_plugin.profiler.splits import infer_data_files, resolve_splits +from nemo_datasets_plugin.profiler.stats import ( + ColumnFold, + InferredColumnFold, + quote_enumerations, +) +from nemo_platform_plugin.files.dataset_profile import ( + ColumnStats, + DatasetProfile, + Evidence, + FeatureSchema, + FileError, + PartitionClassification, + PartitionProfile, + SamplingInfo, + SplitProfile, +) + +PROFILER_NAME = "nemo-dataset-profiler" +PROFILER_VERSION = "0.1.0" + +# Read everything. The budget existed to keep a materialised partition off the heap, and nothing is +# materialised any longer -- a fold's memory is flat in rows, so an exhaustive read costs what a +# short one costs. What it bounded was never really rows, it was risk. +DEFAULT_ROW_BUDGET = None + +# Rows read from a file however thin a caller-supplied budget gets. Below this a file cannot +# contribute the columns it alone witnesses, which is the whole reason every file is opened rather +# than a subset sampled. It is what makes a budget a target rather than a ceiling: 10,000 shards read +# this many each. It survives the default going unbounded because it never had anything to do with +# memory -- dividing a budget across files is about *coverage*, and reading files in order until a +# total ran out would leave the later ones unopened. +MIN_ROWS_PER_FILE = 10 + + +def _format_of(path: str) -> str: + """The registered format of a data file. Callers pass only pre-filtered ``data_entries``, so the + format is always known; a None here would mean that invariant was broken.""" + file_format = detect_format(path) + if file_format is None: + raise ValueError(f"no registered format for {path!r}") + return file_format + + +def profile( + source: FileSource, + *, + created_at: datetime | None = None, + row_budget: int | None = DEFAULT_ROW_BUDGET, + column_roles: dict[str, str] | None = None, +) -> DatasetProfile: + """Profile the dataset behind ``source`` into a ``DatasetProfile``. + + ``row_budget`` bounds how many rows each *partition* reads in total, divided across its files. + It defaults to ``None``, which reads every row: memory is flat in rows either way, so the only + thing a budget buys now is a shorter run. Files smaller than their share are read to the end, so + a budgeted profile of a small dataset is still complete. + + ``column_roles`` maps a column name to a role the caller is asserting, for datasets whose column + names the role table does not recognize. Hints take precedence over name detection but still have + to pass the dtype gates, and a rejected one is reported as evidence rather than dropped. + + ``created_at`` is injectable so a profile can be made reproducible byte-for-byte in tests; it + defaults to the current UTC time. + """ + created_at = created_at or datetime.now(timezone.utc) + all_entries = source.list_files() + data_entries = [entry for entry in all_entries if detect_format(entry.path) is not None] + # Files that plainly hold records but have no reader yet. They are not profiled, but they must be + # reported: silently dropping them let a directory of .csv shards profile as an exhaustively + # scanned, empty dataset — indistinguishable from a dataset that really is empty. They get real + # FileErrors like any other file the profiler could not read, just at the envelope, since no + # partition ever grouped them. Kept as entries, not just paths, because their bytes still count + # toward the size of the fileset even though no partition will ever weigh them. + unreadable_entries = [ + entry + for entry in sorted(all_entries, key=lambda entry: entry.path) + if detect_format(entry.path) is None and is_unsupported_data(entry.path) + ] + file_errors = [ + FileError( + path=entry.path, + error=f"no reader for '{PurePosixPath(entry.path).suffix.lower()}' files", + ) + for entry in unreadable_entries + ] + + partitions: list[PartitionProfile] = [] + rows_scanned = 0 + files_read = 0 + # None once any file's row count is unknown: the fileset's total is then unknowable, not zero. + rows_present: int | None = 0 + + # Every path the source listed, data or not. A split's glob is verified against this rather than + # against the partition's own files, so a pattern can never be emitted that would also pull in a + # README sitting beside the shards. + all_paths = [entry.path for entry in all_entries] + + for name, partition_entries in group_partitions(data_entries): + outcome = _profile_partition( + source, name, partition_entries, row_budget, column_roles or {}, all_paths=all_paths + ) + partitions.append(outcome.partition) + rows_scanned += outcome.rows_scanned + files_read += outcome.files_read + rows_present = _add_known(rows_present, outcome.rows_present) + file_errors.extend(outcome.file_errors) + + # A file the profiler could not use holds an unknown number of rows, so it makes the fileset + # total unknown — whether it was skipped for want of a reader or failed mid-read. + if file_errors: + rows_present = None + + sampling = SamplingInfo( + rows_scanned=rows_scanned, + rows_present=rows_present, + files_read=files_read, # files actually opened and read, not files merely listed + # Every data file, readable or not: the denominator that makes `files_read` a fraction rather + # than a bare count. Non-data files (a README, a LICENSE) are not data and are counted nowhere. + files_present=len(data_entries) + len(unreadable_entries), + # Weighed over the same set, so a fileset the profiler could not read still reports its size. + # Summing the splits would miss the unreadable files, which never reach a partition. + bytes_present=sum(entry.size_bytes for entry in data_entries) + + sum(entry.size_bytes for entry in unreadable_entries), + ) + return DatasetProfile( + created_at=created_at, + profiler_info={"name": PROFILER_NAME, "version": PROFILER_VERSION}, + sampling=sampling, + partitions=partitions, + # Sorted so a reader scanning for trouble sees it in a stable order, whatever partition it + # came from; partitions contribute theirs as they are profiled. + file_errors=sorted(file_errors, key=lambda error: error.path), + ) + + +def _per_file_cap(row_budget: int | None, file_count: int) -> int | None: + """Split a partition's row budget across its files. + + Bounded below by :data:`MIN_ROWS_PER_FILE`, which is what makes the budget a target rather than a + ceiling: at a thousand shards the arithmetic share is ten rows, and at ten thousand it would be + one, which is too thin to witness a column. Overshooting the budget there is the right trade -- + the alternative is sampling a *subset of files*, which hides columns that appear only in later + shards, and file-level sampling is the tier of this problem still to solve. + """ + if row_budget is None: + return None + if file_count <= 1: + return row_budget + return max(MIN_ROWS_PER_FILE, row_budget // file_count) + + +def _add_known(total: int | None, addend: int | None) -> int | None: + """Sum two counts, where ``None`` means unknown and poisons the total. + + A fileset whose row count is unknown for even one file has an unknown total — reporting the sum + of the rest would look like a fact and read low. + """ + if total is None or addend is None: + return None + return total + addend + + +def _unify_schemas(schemas: list[pa.Schema]) -> pa.Schema | None: + """One schema describing every file of the partition, or None when they cannot be reconciled. + + Taking the first file's schema and ignoring the rest makes the profile depend on which shard + happens to sort first: a column that appears only in a later shard would vanish from ``features`` + (and so from ``stats``), and the same data would classify differently under a different file + order. Unifying is order-independent for the common case — later shards adding columns. + + A genuine type conflict for the same column name has no correct answer here, so we return None + and let the caller fall back to inferring from the rows themselves, which widens the conflicting + column to ``json`` rather than asserting one shard's type over the other's. + """ + if not schemas: + return None + if len(schemas) == 1: + return schemas[0] + try: + return pa.unify_schemas(schemas) + except pa.ArrowException: + return None + + +@dataclass(frozen=True) +class _PartitionOutcome: + """One partition plus what it contributes to the dataset-level sampling envelope.""" + + partition: PartitionProfile + rows_scanned: int + files_read: int # files actually opened and read, so `files_read` can exclude failures + rows_present: int | None # rows known to exist here, or None once any file's count is unknown + file_errors: list[FileError] # files this partition grouped but could not fully read + + +def _capped_columns_evidence(features: list[FeatureSchema]) -> list[Evidence]: + """Say so when the schema stopped at the cap rather than at the end of the data. + + A profile that quietly described 4,096 of a file's columns as though they were all of them would + be worse than one that failed: the reader has no way to tell a wide table from a broken one. + """ + if not columns_were_capped(features): + return [] + return [ + Evidence( + kind="error", + detail=( + f"stopped at {MAX_COLUMNS} columns; the rest of this partition's schema was not " + f"described. A file whose rows carry unique keys will do this." + ), + ) + ] + + +def _peek_files(source: FileSource, entries: list[FileEntry]) -> dict[str, FilePreview]: + """What each file declares about itself, before any of them is read. + + A failure here is not reported: it will surface as a :class:`FileError` when the file is actually + read, with a reason, and reporting it twice would double-count. All this decides is whether the + partition can be folded, and a file that cannot be peeked cannot. + """ + previews: dict[str, FilePreview] = {} + for entry in entries: + try: + previews[entry.path] = get_reader(_format_of(entry.path)).peek(source, entry) + except Exception: + previews[entry.path] = FilePreview() + return previews + + +def _expected_rows(previews: dict[str, FilePreview], row_cap: int | None) -> int | None: + """How many rows the fold is about to see, if every file said. + + Capped per file the same way the read will be, so a budgeted run spreads its quality sample over + what it will actually scan rather than over what the dataset holds. + """ + total = 0 + for preview in previews.values(): + if preview.num_rows is None: + return None + total += min(preview.num_rows, row_cap) if row_cap is not None else preview.num_rows + return total + + +class _PartitionFolds: + """The two folds a partition needs, driven together over the same batches. + + One is per column; the other compares two columns of a row against each other and so belongs to + neither. Keeping them side by side is what lets the file loop hand over a batch and forget it. + """ + + def __init__(self, features: list[FeatureSchema] | None, expected_rows: int | None) -> None: + # Declared: the columns are known, so the accumulators are chosen now. Inferred: they are + # discovered as they appear and typed once every row has gone by. + self.features = features or [] + self._columns: ColumnFold | InferredColumnFold = ( + ColumnFold(features, expected_rows) if features is not None else InferredColumnFold(expected_rows) + ) + self._prefix = PrefixPairFold() + self._prefix_error: Evidence | None = None + + def update(self, rows: list[dict]) -> None: + self._columns.update(rows) + # Guarded like the columns are. Unguarded, the only thing that would catch this is the + # per-file handler, which would report odd *data* as a bad *file* -- collapsing the one + # distinction the two failure domains exist to keep. + if self._prefix_error is None: + try: + self._prefix.update(rows) + except Exception as exc: + self._prefix_error = Evidence( + kind="error", + detail=f"the chosen/rejected prefix probe could not run: {type(exc).__name__}: {exc}", + ) + + def measure( + self, column_roles: dict[str, str] + ) -> tuple[list[FeatureSchema], dict[str, ColumnStats], PartitionClassification]: + """Schema, stats and classification from what was folded. + + The columns have their own per-column guard inside the fold; this is the wide one, for + anything structural that no single column owns -- a schema that cannot be resolved, + a classifier that trips over a shape no detector anticipated. + """ + try: + if isinstance(self._columns, InferredColumnFold): + self.features, measured = self._columns.finalize() + else: + measured = self._columns.finalize() + classification = classify( + self.features, + measured.stats, + probes=measured.probes, + prefix_pair=self._prefix.result(), + column_roles=column_roles, + ) + quote_enumerations(self.features, measured.stats, measured.vocabularies) + classification.evidence.extend(_capped_columns_evidence(self.features)) + classification.evidence.extend(measured.errors) + if self._prefix_error is not None: + classification.evidence.append(self._prefix_error) + return self.features, measured.stats, classification + except Exception as exc: + detail = f"could not measure this partition: {type(exc).__name__}: {exc}" + return ( + [], + {}, + PartitionClassification(dataset_type="unknown", evidence=[Evidence(kind="error", detail=detail)]), + ) + + +def _profile_partition( + source: FileSource, + name: str, + entries: list[FileEntry], + row_budget: int | None, + column_roles: dict[str, str], + *, + all_paths: list[str], +) -> _PartitionOutcome: + """Profile one partition — the files of one source directory, whatever formats they are in. + + The reader is resolved per file rather than per partition. Format is a property of a file, and + a directory holding two of them is a stray file, not a second dataset; splitting the partition + to keep one scalar ``file_format`` true is what made partition names unstable. A partition whose + files do not all declare a schema simply infers one, from the rows, as it folds them. + + An unreadable file (or a format with no registered reader) is isolated: it is named on a + :class:`FileError` the envelope collects, contributes no rows, and flips ``scanned_all`` off — it + never aborts the profile. Files that read cleanly are counted, not listed. + """ + # Footers first, before a single row is read. A parquet file declares its schema and its exact + # row count there, so one seek per file establishes the partition's whole shape: what the columns + # are, and how many rows are coming. That is what a fold needs and cannot otherwise have -- the + # accumulators must exist before the first batch, and the quality stride must be placed before + # the column it strides has been seen. + previews = _peek_files(source, entries) + arrow_schemas = [preview.arrow_schema for preview in previews.values() if preview.arrow_schema is not None] + declared = _unify_schemas(arrow_schemas) if len(arrow_schemas) == len(entries) and arrow_schemas else None + # Declared or not, the partition folds. With a schema the accumulators are chosen up front and + # the exact row count places the quality stride; without one both wait for the data, which costs + # a deferred dtype per column and nothing else. + row_cap = _per_file_cap(row_budget, len(entries)) + folds = _PartitionFolds( + derive_features([], declared) if declared is not None else None, + expected_rows=_expected_rows(previews, row_cap), + ) + rows_scanned = 0 + files_read = 0 + rows_present: int | None = 0 + partition_scanned = True + file_errors: list[FileError] = [] + file_formats: set[str] = set() + split_profiles: list[SplitProfile] = [] + for split in resolve_splits(entries): + split_examples = 0 + split_counts_known = True # every file's exact total row count is known (footer or full scan) + split_scanned = True # every row of every file was actually parsed + for entry in split.entries: + file_formats.add(_format_of(entry.path)) + error: str | None = None + num_rows: int | None = None + scanned_all = False + scanned = 0 + try: + # Inside the guard: resolving the reader can fail too, and a format with no reader + # registered is a file the profiler could not use like any other. + reader = get_reader(_format_of(entry.path)) + preview = previews[entry.path] + read_errors: list[str] = [] + for batch in reader.batches(source, entry, row_cap=row_cap, errors=read_errors): + folds.update(batch) + scanned += len(batch) + # A file the reader only partly understood is named, the same as one it could not + # open at all. Folding it silently would make a corrupt shard look complete. + error = "; ".join(read_errors) or None + # A footer knows the count before the read; a line-delimited file only knows it by + # reaching the end, which a capped read does not do. + if preview.num_rows is not None: + num_rows = preview.num_rows + elif row_cap is None or scanned < row_cap: + num_rows = scanned + # Exhaustive requires parsing every row. A known count alone is not enough, and a + # partial read is not exhaustive however many rows it managed to get. + scanned_all = num_rows is not None and scanned >= num_rows and error is None + except Exception as exc: + # Failure isolation: an unreadable file (or missing reader) keeps its identity, + # skips the rest of its rows, and does not abort the profile. The reason is recorded + # rather than swallowed, so a consumer can tell corrupt input from a profiler bug. + error = f"{type(exc).__name__}: {exc}" + num_rows = None + scanned_all = False + # Counted for what was actually consumed, outside the guard, because a read is no longer + # all-or-nothing: a fold cannot give rows back, so a file that failed on its fifth batch + # still contributed four and the envelope has to say so. Accounting for it as unread + # would leave `rows_scanned` describing fewer rows than the stats were built from. + rows_scanned += scanned + if scanned or error is None: + files_read += 1 + if error is not None: + file_errors.append(FileError(path=entry.path, error=error)) + rows_present = _add_known(rows_present, num_rows) + if num_rows is None: + split_counts_known = False + else: + split_examples += num_rows + if not scanned_all: + split_scanned = False + partition_scanned = partition_scanned and split_scanned + split_profiles.append( + SplitProfile( + name=split.name, + canonical=split.canonical, + # Inferred from the same paths the split itself was read off, then verified against + # the whole listing; None when one pattern cannot express the split exactly. + data_files=infer_data_files(split.name, split.entries, all_paths), + num_files=len(split.entries), + # From the listing, not from reading, so a file that failed mid-read still weighs + # what it weighs — unlike `num_examples`, this never goes unknown. + size_bytes=sum(entry.size_bytes for entry in split.entries), + num_examples=split_examples if split_counts_known else None, + ) + ) + features, stats, classification = folds.measure(column_roles) + partition = PartitionProfile( + name=name, + # Observed, not chosen: the partition reports the formats its files turned out to be in + # rather than picking one and splitting to keep that true. + file_formats=sorted(file_formats), + splits=split_profiles, + features=features, + stats=stats, + # Scoped to this partition, which is where it was decided all along: `partition_scanned` is + # the value that already gated whether `categorical.values` could quote a proven enumeration. + rows_complete=partition_scanned, + classification=classification, + ) + return _PartitionOutcome( + partition=partition, + rows_scanned=rows_scanned, + files_read=files_read, + rows_present=rows_present, + file_errors=file_errors, + ) diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/base.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/base.py new file mode 100644 index 0000000000..38b4d3c5aa --- /dev/null +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/base.py @@ -0,0 +1,132 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Per-format reader contract and registry. + +Each reader is a stateless handler for one on-disk format, safe to reuse across files. Readers are +looked up by ``file_format`` (from :func:`detect_format`) so the pipeline never branches on format +itself. Built-in readers self-register (a ``register_reader`` call at the bottom of each module) and +are loaded lazily on the first :func:`get_reader` call, so importing this module does not pull in +pyarrow until a reader is actually needed. +""" + +from __future__ import annotations + +from collections.abc import Iterator +from dataclasses import dataclass +from pathlib import Path +from typing import Any, ClassVar, Protocol + +import pyarrow as pa +from nemo_datasets_plugin.profiler.file_source import FileEntry, FileSource + + +@dataclass(frozen=True) +class ReadResult: + """What a format reader returns for one file.""" + + rows: list[dict[str, Any]] # the rows read (a sample, or all of them) + rows_scanned: int # number of rows actually parsed + num_rows: int | None = None # exact total when cheaply known (e.g. a parquet footer), else None + arrow_schema: pa.Schema | None = None # the declared column schema, when the format carries one + # Why the read understood less than the whole file, when that happened. None means nothing was + # lost. This is the only channel a reader has to explain a partial result, so a consumer can tell + # "corrupt input" from "unsupported format" from "profiler bug" instead of seeing a silent gap. + error: str | None = None + + +@dataclass(frozen=True) +class FilePreview: + """What a reader can learn about a file without reading a single row. + + A parquet footer carries both; a line-delimited format carries neither. The pipeline asks every + file this before it reads any of them, because knowing the schema up front is what lets a + partition be measured without first being materialised, and knowing the row count up front is + what lets a quality sample be strided across a column the fold has not finished seeing. + """ + + arrow_schema: pa.Schema | None = None + num_rows: int | None = None + + +class FormatReader(Protocol): + """Reads schema and rows for one file format.""" + + file_format: ClassVar[str] + + def peek(self, source: FileSource, entry: FileEntry) -> FilePreview: + """What the file declares about itself, without reading its rows.""" + ... + + def read(self, source: FileSource, entry: FileEntry, *, row_cap: int | None = None) -> ReadResult: + """Read up to ``row_cap`` rows (all rows when None) plus whatever the format declares cheaply.""" + ... + + def batches( + self, + source: FileSource, + entry: FileEntry, + *, + row_cap: int | None = None, + errors: list[str] | None = None, + ) -> Iterator[list[dict[str, Any]]]: + """The same rows :meth:`read` would return, handed over in chunks and never all at once. + + ``errors`` collects any reason the read understood less than the whole file, the way + :attr:`ReadResult.error` does for the batched-up path. A generator cannot return one -- + by the time it knows, the caller has consumed everything it yielded -- and the caller has + to know, or a partially parsed file would fold silently and look complete. + """ + ... + + +_READERS: dict[str, FormatReader] = {} +_builtins_loaded = False + + +def _load_builtin_readers() -> None: + """Import the built-in reader modules so their self-registration runs (once). + + Deferred to call time — not import time — so there is no cycle with the reader modules that import + from this one, and pyarrow stays out of the import graph until a reader is actually resolved. + """ + global _builtins_loaded + if _builtins_loaded: + return + _builtins_loaded = True + from nemo_datasets_plugin.profiler.readers import jsonl, parquet # noqa: F401 self-registering + + +def register_reader(reader: FormatReader) -> None: + _READERS[reader.file_format] = reader + + +def get_reader(file_format: str) -> FormatReader: + _load_builtin_readers() + try: + return _READERS[file_format] + except KeyError: + raise KeyError(f"no reader registered for file format {file_format!r}") from None + + +_EXTENSION_FORMATS = { + ".parquet": "parquet", + ".jsonl": "jsonl", + ".ndjson": "jsonl", +} + + +def detect_format(path: str) -> str | None: + """Map a file path to a registered format by extension, or None when unrecognized.""" + return _EXTENSION_FORMATS.get(Path(path).suffix.lower()) + + +# Extensions that plainly hold dataset records but have no reader yet. Naming them explicitly is what +# lets the profiler say "there is data here I cannot read" instead of treating a dataset it does not +# understand as an empty one — a README or a LICENSE is genuinely not data and stays ignored. +_UNSUPPORTED_DATA_EXTENSIONS = {".csv", ".tsv", ".arrow", ".feather", ".json", ".avro", ".orc"} + + +def is_unsupported_data(path: str) -> bool: + """Whether a path looks like dataset records this profiler has no reader for.""" + return Path(path).suffix.lower() in _UNSUPPORTED_DATA_EXTENSIONS diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/jsonl.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/jsonl.py new file mode 100644 index 0000000000..1ea7da9809 --- /dev/null +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/jsonl.py @@ -0,0 +1,119 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Line-delimited JSON reader. No declared schema; the row count is exact only on a full read.""" + +from __future__ import annotations + +import json +from collections.abc import Iterator + +from nemo_datasets_plugin.profiler.file_source import FileEntry, FileSource +from nemo_datasets_plugin.profiler.readers.base import FilePreview, ReadResult, register_reader + +# Rows handed over at a time by :meth:`JsonlReader.batches`, matching the parquet reader so the +# caller's working set does not depend on which format it happens to be folding. +_BATCH_ROWS = 1024 + + +def _records(stream) -> Iterator[tuple[dict | None, str | None]]: + """Each line of the stream as either a record or a reason it was not one. + + Shared by :meth:`JsonlReader.read` and :meth:`JsonlReader.batches` so the two cannot drift on + what counts as a row -- a blank line, a stray scalar and a truncated line are three different + things and only one of them is an error. + """ + for line_number, raw_line in enumerate(stream, start=1): + stripped = raw_line.strip() + if not stripped: # tolerate blank lines between records + continue + try: + record = json.loads(stripped) + except ValueError as exc: + # A truncated or corrupt line costs that line, never the file. Dropping the whole file + # would erase its row count and every column it was the only witness for. + yield None, f"line {line_number}: {exc}" + continue + if isinstance(record, dict): + yield record, None + # a record is a column map; stray scalars/arrays are skipped rather than crash downstream + + +class JsonlReader: + file_format = "jsonl" + + def peek(self, source: FileSource, entry: FileEntry) -> FilePreview: + """Nothing. A line-delimited file declares no schema and carries no row count, which is why + a partition holding one cannot be folded without reading it first.""" + return FilePreview() + + def batches( + self, + source: FileSource, + entry: FileEntry, + *, + row_cap: int | None = None, + errors: list[str] | None = None, + ) -> Iterator[list[dict]]: + """Rows in chunks, reporting any line it could not read into ``errors``. + + A corrupt line costs that line, never the file -- dropping the whole file would erase its + row count and every column it was the only witness for -- but the caller has to be told, or + a partially parsed file would fold silently and look complete. + """ + rows: list[dict] = [] + scanned = 0 + unparseable = 0 + first_failure: str | None = None + with source.open(entry.path) as stream: + for record, failure in _records(stream): + if record is None: + unparseable += 1 + if first_failure is None: + first_failure = failure + continue + rows.append(record) + scanned += 1 + if len(rows) >= _BATCH_ROWS: + yield rows + rows = [] + if row_cap is not None and scanned >= row_cap: + break + if rows: + yield rows + if unparseable and errors is not None: + errors.append(f"skipped {unparseable} unparseable line(s); first at {first_failure}") + + def read(self, source: FileSource, entry: FileEntry, *, row_cap: int | None = None) -> ReadResult: + rows: list[dict] = [] + unparseable = 0 + first_failure: str | None = None + hit_cap = False + with source.open(entry.path) as stream: + for record, failure in _records(stream): + # Branching on the record rather than the failure: the two are exclusive, and this + # way the type narrows without an ignore standing in for the reasoning. + if record is None: + unparseable += 1 + if first_failure is None: + first_failure = failure + continue + rows.append(record) + if row_cap is not None and len(rows) >= row_cap: + hit_cap = True + break + + # `num_rows` counts records the reader could parse. A line of valid JSON that simply is not a + # row (a stray scalar or array) is not a row of this dataset, so it leaves the count exact and + # sets no error. An *unparseable* line is data we failed to read, so it is reported: the + # pipeline reads `error` to decide the file was not exhaustively scanned. + # + # A cap only costs the exact count when it actually stopped the read. A file smaller than the + # cap was still read to EOF, so it keeps an exact count — which is what lets a capped profile + # of a small dataset stay exhaustive instead of degrading every stat for no reason. + num_rows = None if hit_cap else len(rows) + error = f"skipped {unparseable} unparseable line(s); first at {first_failure}" if unparseable else None + return ReadResult(rows=rows, rows_scanned=len(rows), num_rows=num_rows, arrow_schema=None, error=error) + + +register_reader(JsonlReader()) diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/parquet.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/parquet.py new file mode 100644 index 0000000000..3eed3ccfae --- /dev/null +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/parquet.py @@ -0,0 +1,76 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Parquet reader — the footer gives an exact row count and the declared schema for free.""" + +from __future__ import annotations + +from collections.abc import Iterator + +import pyarrow.parquet as pq +from nemo_datasets_plugin.profiler.file_source import FileEntry, FileSource +from nemo_datasets_plugin.profiler.readers.base import FilePreview, ReadResult, register_reader + +# Rows handed over at a time. Small enough that the working set is a knob independent of the +# dataset, large enough that per-batch overhead stays invisible. +_BATCH_ROWS = 1024 + + +class ParquetReader: + file_format = "parquet" + + def peek(self, source: FileSource, entry: FileEntry) -> FilePreview: + """Schema and exact row count from the footer. Reads no rows, so a partition's whole shape is + knowable for the cost of one seek per file.""" + with source.open(entry.path) as stream: + parquet_file = pq.ParquetFile(stream) + return FilePreview(arrow_schema=parquet_file.schema_arrow, num_rows=parquet_file.metadata.num_rows) + + def batches( + self, + source: FileSource, + entry: FileEntry, + *, + row_cap: int | None = None, + errors: list[str] | None = None, + ) -> Iterator[list[dict]]: + """Rows in chunks, so the caller can fold them and let each chunk go. + + ``errors`` is never appended to: parquet either decodes a batch or raises, so there is no + partial understanding to report. + """ + scanned = 0 + with source.open(entry.path) as stream: + parquet_file = pq.ParquetFile(stream) + if row_cap == 0: + return + for batch in parquet_file.iter_batches(batch_size=min(row_cap or _BATCH_ROWS, _BATCH_ROWS)): + rows = batch.to_pylist() + if row_cap is not None and scanned + len(rows) > row_cap: + rows = rows[: row_cap - scanned] + scanned += len(rows) + if rows: + yield rows + if row_cap is not None and scanned >= row_cap: + return + + def read(self, source: FileSource, entry: FileEntry, *, row_cap: int | None = None) -> ReadResult: + with source.open(entry.path) as stream: + parquet_file = pq.ParquetFile(stream) + num_rows = parquet_file.metadata.num_rows + arrow_schema = parquet_file.schema_arrow + if row_cap == 0: + return ReadResult(rows=[], rows_scanned=0, num_rows=num_rows, arrow_schema=arrow_schema) + + rows: list[dict] = [] + for batch in parquet_file.iter_batches(batch_size=row_cap or 1024): + rows.extend(batch.to_pylist()) + if row_cap is not None and len(rows) >= row_cap: + break + + if row_cap is not None: + rows = rows[:row_cap] + return ReadResult(rows=rows, rows_scanned=len(rows), num_rows=num_rows, arrow_schema=arrow_schema) + + +register_reader(ParquetReader()) diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/schema.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/schema.py new file mode 100644 index 0000000000..3c95f744fb --- /dev/null +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/schema.py @@ -0,0 +1,235 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Row-schema derivation. + +Derive the ``features`` tree (a list of :class:`FeatureSchema`) de novo from the data. Parquet +carries a declared schema, so it is converted directly; formats without one (jsonl) are inferred +from the sampled rows by resolving each column's dtype. A list of ``{role, content}`` structs — or +ShareGPT's ``{from, value}`` spelling of the same thing — is recognized as the ``messages`` dtype, +and a list of ``{role, content}`` structs is recognized as the ``messages`` dtype. +""" + +from __future__ import annotations + +from typing import Any + +import pyarrow as pa +from nemo_platform_plugin.files.dataset_profile import FeatureSchema + +# A list element carrying at least one of these key sets is treated as a chat message (the messages +# dtype). ShareGPT-style data spells the same structure `{from, value}`; recognizing only +# `{role, content}` left a large slice of public chat data typed as a plain `list`, which then failed +# the `messages` dtype gate in classification and profiled as `unknown` with no stats at all. +_MESSAGE_KEY_SETS = ({"role", "content"}, {"from", "value"}) + + +# Columns a partition may have before the profiler stops describing it. Nothing legitimate reaches +# this: it is a guard against a malformed file whose rows carry per-row unique keys, which would +# otherwise mint a column -- and an accumulator -- for every row in the dataset. The row budget used +# to bound this by accident, since a schema inferred from at most N rows had at most N keys; an +# unbounded read has no such accident, so the bound is stated. +MAX_COLUMNS = 4096 + + +def derive_features(rows: list[dict[str, Any]], arrow_schema: pa.Schema | None = None) -> list[FeatureSchema]: + """The row schema. Uses the declared arrow schema when present, else infers from ``rows``. + + Truncated at :data:`MAX_COLUMNS`. Use :func:`columns_were_capped` to tell a dataset that really + is that wide from one whose keys are runaway; the caller reports the difference rather than + silently describing part of a file as though it were all of it. + """ + if arrow_schema is not None: + return [ + _feature_from_arrow(arrow_schema.field(i).name, arrow_schema.field(i).type) + for i in range(min(len(arrow_schema), MAX_COLUMNS)) + ] + return _features_from_rows(rows) + + +def columns_were_capped(features: list[FeatureSchema]) -> bool: + """Whether the schema stopped at the cap rather than at the end of the data.""" + return len(features) >= MAX_COLUMNS + + +def _is_message_struct(item: FeatureSchema) -> bool: + if item.dtype != "struct" or item.fields is None: + return False + names = {field.name for field in item.fields} + return any(keys <= names for keys in _MESSAGE_KEY_SETS) + + +# --- from a declared arrow schema (parquet) ------------------------------------------------------ + +_ARROW_SCALAR_DTYPES = [ + (pa.types.is_boolean, "bool"), + (pa.types.is_int8, "int8"), + (pa.types.is_int16, "int16"), + (pa.types.is_int32, "int32"), + (pa.types.is_int64, "int64"), + (pa.types.is_uint8, "uint8"), + (pa.types.is_uint16, "uint16"), + (pa.types.is_uint32, "uint32"), + (pa.types.is_uint64, "uint64"), + (pa.types.is_float16, "float16"), + (pa.types.is_float32, "float32"), + (pa.types.is_float64, "float64"), + (pa.types.is_string, "string"), + (pa.types.is_large_string, "string"), +] + + +def _arrow_scalar_dtype(arrow_type: pa.DataType) -> str: + for predicate, dtype in _ARROW_SCALAR_DTYPES: + if predicate(arrow_type): + return dtype + return "json" + + +def _feature_from_arrow(name: str, arrow_type: pa.DataType) -> FeatureSchema: + if pa.types.is_struct(arrow_type): + fields = [ + _feature_from_arrow(arrow_type.field(i).name, arrow_type.field(i).type) + for i in range(arrow_type.num_fields) + ] + return FeatureSchema(name=name, dtype="struct", fields=fields) + if pa.types.is_list(arrow_type) or pa.types.is_large_list(arrow_type) or pa.types.is_fixed_size_list(arrow_type): + item = _feature_from_arrow("", arrow_type.value_type) + dtype = "messages" if _is_message_struct(item) else "list" + return FeatureSchema(name=name, dtype=dtype, items=item) + return FeatureSchema(name=name, dtype=_arrow_scalar_dtype(arrow_type)) + + +# --- inferred from sampled rows (jsonl) ---------------------------------------------------------- + + +def _features_from_rows(rows: list[dict[str, Any]]) -> list[FeatureSchema]: + ordered_keys: list[str] = [] + seen: set[str] = set() + for row in rows: + for key in row: + if key not in seen: + seen.add(key) + ordered_keys.append(key) + if len(ordered_keys) >= MAX_COLUMNS: + break + if len(ordered_keys) >= MAX_COLUMNS: + break + return [_infer_feature(key, [row.get(key) for row in rows]) for key in ordered_keys] + + +def _infer_feature(name: str, values: list[Any]) -> FeatureSchema: + present = [value for value in values if value is not None] + if not present: + return FeatureSchema(name=name, dtype="json") + + if all(isinstance(value, dict) for value in present): + child_keys: list[str] = [] + seen: set[str] = set() + for record in present: + for key in record: + if key not in seen: + seen.add(key) + child_keys.append(key) + fields = [_infer_feature(key, [record.get(key) for record in present]) for key in child_keys] + return FeatureSchema(name=name, dtype="struct", fields=fields) + + if all(isinstance(value, list) for value in present): + item = _infer_feature("", [element for value in present for element in value]) + if _is_message_struct(item): + return FeatureSchema(name=name, dtype="messages", items=item) + return FeatureSchema(name=name, dtype="list", items=item) + + return FeatureSchema(name=name, dtype=_scalar_dtype(present)) + + +def _python_dtype(value: Any) -> str: + """The dtype one value implies, on its own.""" + if isinstance(value, bool): # bool before int: bool is a subclass of int + return "bool" + if isinstance(value, int): + return "int64" + if isinstance(value, float): + return "float64" + if isinstance(value, str): + return "string" + return "json" + + +def _resolve_scalar(dtypes: set[str]) -> str: + """The one dtype a column of these observed types has. Ints and floats widen; anything else in + disagreement is ``json``, which is the honest answer for a column that holds two shapes.""" + if dtypes <= {"int64", "float64"} and dtypes: + return "float64" if "float64" in dtypes else "int64" + if len(dtypes) == 1: + # `next(iter(...))`, never `pop()`: this set belongs to a SchemaFold that is still using it, + # and emptying it made a second call resolve the same column to `json`. + return next(iter(dtypes)) + return "json" + + +def _scalar_dtype(values: list[Any]) -> str: + return _resolve_scalar({_python_dtype(value) for value in values}) + + +class SchemaFold: + """One column's schema, folded from values as they arrive rather than decided over all of them. + + The dtype of an inferred column is a whole-column question -- the observed types are unioned and + a disagreement widens to ``json`` -- which is why a partition with no declared schema could not + be folded: an accumulator is chosen *by* dtype, and the dtype is not known until the last row. + + It is a fold, though, and always was. :func:`_infer_feature` is a set union over observed types, + a union over a struct's child keys, and a recursion over a list's flattened elements: state + proportional to the *schema*, not to the row count. This is that computation written incrementally + so a caller can hand over batches and keep none of them. + """ + + def __init__(self, name: str = "") -> None: + self._name = name + self._present = 0 + self._dicts = 0 + self._lists = 0 + self._dtypes: set[str] = set() + self._fields: dict[str, SchemaFold] = {} + self._field_order: list[str] = [] + self._item: SchemaFold | None = None + + def update(self, values: list[Any]) -> None: + for value in values: + if value is None: + continue # a null says nothing about the type; an all-null column resolves to json + self._present += 1 + self._dtypes.add(_python_dtype(value)) + if isinstance(value, dict): + self._dicts += 1 + for key, child in value.items(): + fold = self._fields.get(key) + if fold is None: + fold = SchemaFold(key) + self._fields[key] = fold + self._field_order.append(key) + fold.update([child]) + elif isinstance(value, list): + self._lists += 1 + if self._item is None: + self._item = SchemaFold() + self._item.update(value) + + def finalize(self) -> FeatureSchema: + if not self._present: + return FeatureSchema(name=self._name, dtype="json") + if self._dicts == self._present: + return FeatureSchema( + name=self._name, + dtype="struct", + fields=[self._fields[key].finalize() for key in self._field_order], + ) + if self._lists == self._present: + item = (self._item or SchemaFold()).finalize() + return FeatureSchema( + name=self._name, + dtype="messages" if _is_message_struct(item) else "list", + items=item, + ) + return FeatureSchema(name=self._name, dtype=_resolve_scalar(self._dtypes)) diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/splits.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/splits.py new file mode 100644 index 0000000000..dfb79c45ba --- /dev/null +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/splits.py @@ -0,0 +1,165 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Split resolution from file paths. + +Given the files in one partition, group them into splits by inferring each file's split from its +path — a split-named directory first, then the shard-stripped filename. A declared split map from a +dataset card would take precedence over this inference, but card parsing is not wired up yet, so +path inference is the only source today. +""" + +from __future__ import annotations + +import re +from dataclasses import dataclass +from pathlib import PurePosixPath + +from nemo_datasets_plugin.profiler.file_source import FileEntry + +# Strips a shard suffix like "-00000" or "-00000-of-00003" from a file stem. A bare trailing number +# must be zero-padded to count: `-\d{2,}` alone also matched years and versions, turning +# covid-19.jsonl into a "covid" split and data-2024.jsonl into "data". +_SHARD_SUFFIX = re.compile(r"-(?:\d{2,}-of-\d{2,}|0\d{3,})$") + +# Common on-disk split words -> the canonical concept they normalize to. +_CANONICAL_ALIASES = { + "train": "train", + "test": "test", + "validation": "validation", + "valid": "validation", + "val": "validation", + "dev": "validation", +} + + +@dataclass(frozen=True) +class ResolvedSplit: + """A split and the files that belong to it.""" + + name: str # the on-disk split name, e.g. "train" or "train_prefs" + canonical: str | None # normalized concept (train | validation | test), or None + entries: list[FileEntry] + + +def _canonical_for(split_name: str) -> str | None: + """Map a split name to its canonical concept, tolerating variant suffixes (train_prefs -> train).""" + lowered = split_name.lower() + for alias, canonical in _CANONICAL_ALIASES.items(): + if lowered == alias or lowered.startswith(f"{alias}_") or lowered.startswith(f"{alias}-"): + return canonical + return None + + +def is_split_directory(name: str) -> bool: + """Whether a directory name denotes a split rather than a partition. + + Partition grouping needs this to avoid turning ``train/`` and ``test/`` into two unrelated + partitions, each with its own schema and classification, when they are two splits of one dataset. + """ + return _canonical_for(name) is not None + + +def _split_name(path: str) -> str: + """The split a file belongs to. + + A split-named directory anywhere on the path wins over the filename, because the ``data/train/`` + layout names its shards ``0000.parquet`` — reading only the stem would file every split's shards + under the same meaningless name and collapse the whole dataset into one split. Nearest directory + to the file wins. Failing that, the shard-stripped stem: train-00000-of-00003.parquet -> "train". + """ + parts = PurePosixPath(path).parts + for directory in reversed(parts[:-1]): + if is_split_directory(directory): + return directory + return _SHARD_SUFFIX.sub("", parts[-1].split(".")[0]) + + +def _glob_matches(pattern: str, path: str) -> bool: + """Match ``path`` against ``pattern`` where ``*`` spans any run of characters except ``/``. + + Deliberately the narrowest dialect rather than the most expressive one. This single reading of + ``*`` is shared by shell globs, Python's :mod:`glob`, fsspec and HF ``data_files``, so a pattern + emitted here means the same thing wherever a consumer pastes it. ``**`` is not produced at all: + its semantics differ between those tools, and a pattern that silently selects a different set of + files in the reader than it did in the profiler is worse than no pattern. + """ + return re.fullmatch("[^/]*".join(re.escape(part) for part in pattern.split("*")), path) is not None + + +def infer_data_files(split_name: str, entries: list[FileEntry], all_paths: list[str]) -> str | None: + """A glob selecting exactly ``entries`` out of ``all_paths``, or None if no single one does. + + This is the inverse of the split resolution above: splits are read *off* the paths, so the + pattern is rebuilt from the same evidence — the directory the files share, and the split name + their filenames start with. It restores addressability of a split's files without restoring the + per-file manifest that was removed for scaling: one pattern per split, whatever the shard count. + + Candidates run most specific first, because the general ones over-match. In ``helpsteer2/`` a + ``train`` split wants ``helpsteer2/train*.parquet``; ``helpsteer2/*.parquet`` would swallow the + validation shards too, and only the ordering distinguishes them. The ``data/train/0000.parquet`` + layout is the other way round — no filename starts with "train", so the name-anchored candidates + are skipped and the directory-wide one is exactly right. + + Every candidate is matched back against *every* file the source listed, not just this split's or + even just this partition's, and the first that reproduces the split exactly wins. Nothing is + emitted on a near miss. A glob is an instruction to go read files, so an approximate one is not + a smaller version of the right answer -- it quietly pulls a README, or another split's shards, + into a training set. None says "these files are not expressible as one pattern", which is a + thing a consumer can handle; a wrong pattern is not. + """ + paths = [entry.path for entry in entries] + directories = {PurePosixPath(path).parent.as_posix() for path in paths} + if len(directories) != 1: + # Shards spread across subdirectories. Covering them needs `**`, whose meaning is not shared + # across glob implementations, so this reports no pattern rather than an ambiguous one. + return None + directory = directories.pop() + prefix = "" if directory == "." else f"{directory}/" + names = [PurePosixPath(path).name for path in paths] + suffixes = {PurePosixPath(path).suffix for path in paths} + # A partition may hold more than one format; only a single shared suffix can go in the pattern. + suffix = suffixes.pop() if len(suffixes) == 1 else None + + # Stems to anchor on: the bare split name first, then the same name keeping the separator that + # follows it. Plain `train*` is what a person would write and is tried first for that reason; + # the separator variants exist only for the sibling collision, where `train` sits beside + # `train_prefs` in one directory and `train*` swallows both. Verification is what demotes the + # simple form there, so the narrower `train-*` is reached only when it is actually needed. Each + # stem is offered only when every filename in the split carries it, so one can never exclude a + # file it ought to match. + stems = [split_name] + [f"{split_name}{sep}" for sep in ("-", ".", "_")] if split_name else [] + candidates: list[str] = [] + for stem in stems: + if not all(name.startswith(stem) for name in names): + continue + if suffix: + candidates.append(f"{prefix}{stem}*{suffix}") + candidates.append(f"{prefix}{stem}*") + if suffix: + candidates.append(f"{prefix}*{suffix}") + candidates.append(f"{prefix}*") + + target = set(paths) + for candidate in candidates: + if {path for path in all_paths if _glob_matches(candidate, path)} == target: + return candidate + return None + + +def resolve_splits(entries: list[FileEntry]) -> list[ResolvedSplit]: + """Group files into splits by path inference. + + Each file's split name comes from a split-named directory on its path, else its shard-stripped + stem; the canonical concept is matched against common aliases (val/valid/dev -> validation). When + no file carries a recognizable split, every file lands in one "default" split. + """ + grouped: dict[str, list[FileEntry]] = {} + for entry in entries: + grouped.setdefault(_split_name(entry.path), []).append(entry) + + canonicals = {name: _canonical_for(name) for name in grouped} + if not any(canonicals.values()): + return [ResolvedSplit(name="default", canonical=None, entries=list(entries))] + + return [ResolvedSplit(name=name, canonical=canonicals[name], entries=grouped[name]) for name in sorted(grouped)] diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py new file mode 100644 index 0000000000..071765282a --- /dev/null +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py @@ -0,0 +1,960 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Per-column statistics and content probes. + +Given a partition's features and its rows, :func:`measure_columns` measures each top-level column +according to its dtype: length quantiles and corruption signals for text, min/max/mean for numbers, +chat-shape signals for messages, and a bounded vocabulary where the column has one. The result is +sparse — a column with nothing worth measuring is omitted — and each column is isolated, so one the +detectors cannot handle costs only itself. Row values themselves are never stored here at all; a +small controlled vocabulary is added afterwards by :func:`quote_enumerations`, which gates on role. + +The same pass reads each column's *content* — answer markers, embedded transcripts — as plain +per-column counts (:class:`ColumnProbes`). Those are measurements, not interpretations: what they +mean is classification's job, and keeping the looking here is what stops a content signal from being +reachable only through a correctly named column. + +The measuring itself is done by a :class:`ColumnAccumulator` per column, chosen once on dtype. An +accumulator folds batches in and keeps no reference to them, so a column measured in pieces gives +the same answer as one measured whole — the property that lets a caller stop materialising a +partition before it can measure it. The base class is the entire measurement for a dtype with no +statistics of its own, because the probes run over every column whatever its type. + +Every measurement is O(1) in rows. Nothing is retained: the one thing that used to be — a string +column's values, held so a quality sample could be placed across them — went when the sample learned +to place itself as it goes, in contiguous blocks whose size does not depend on how long the column +turns out to be. +""" + +from __future__ import annotations + +import math +import re +from dataclasses import dataclass +from typing import Any + +from nemo_datasets_plugin.profiler.schema import MAX_COLUMNS, SchemaFold +from nemo_platform_plugin.files.dataset_profile import ( + CategoricalStats, + ColumnStats, + Evidence, + FeatureSchema, + MessageStats, + NumericStats, + Quantiles, + TextQuality, + TextStats, +) + +# A quotable enumeration holds at most this many distinct values. +_MAX_ENUM_VALUES = 32 + +# Where a column stops being a plausible controlled vocabulary, and so stops being worth counting. +# Three bounds because a count alone bounds cardinality but not bytes -- 1024 reasoning traces is +# 32 MB. `_MAX_VOCABULARY_VALUE_CHARS` is the one that matters most: it is a claim about what the +# column *is* rather than how big it is, so it settles a free-text column on the first value instead +# of after a thousand. Sized well above real vocabularies -- a `source` column spanning 500 datasets, +# a 200-class label set -- and far below anything that costs memory. +_MAX_VOCABULARY_VALUES = 1024 +_MAX_VOCABULARY_VALUE_CHARS = 256 +_MAX_VOCABULARY_BYTES = 64 * 1024 +# What a non-string distinct value is charged against the byte bound. Ints and bools are small and +# uniform, so an exact size buys nothing the count bound does not already give. +_NON_STRING_VALUE_BYTES = 8 + +# Roles that are controlled vocabularies by construction, and so are safe to quote at any dataset +# size. Everything else -- prompts, completions, chosen/rejected, context, chat -- is free text no +# matter how few distinct values a small sample happens to show, and unroled columns are unknown, +# which is the same thing for this purpose. An allowlist, so an unrecognized column fails to silence +# rather than to exposure. +_QUOTABLE_ROLES = frozenset({"label", "provenance", "meta", "rank"}) + + +@dataclass(frozen=True) +class ColumnMeasurements: + """Everything one pass over a partition's columns produced. + + A named result rather than a tuple because the vocabularies joined it: they are not part of the + stored profile, but :func:`quote_enumerations` needs them and can no longer go back to the rows + for them. + """ + + stats: dict[str, ColumnStats] + probes: dict[str, ColumnProbes] + vocabularies: dict[str, set[Any]] # name -> distinct values, only where the column stayed one + errors: list[Evidence] + + +class ColumnFold: + """The per-column accumulators for one partition, fed batch by batch. + + Each column is isolated. A value no detector anticipated -- a chat message whose ``role`` is a + number, a float where a string was declared -- costs that column its measurements and nothing + else, where previously it cost the partition every measurement it had. The failure is reported + as an ``error`` evidence rather than left as a silent gap, because a column absent from ``stats`` + is otherwise indistinguishable from one that simply had nothing worth measuring. It is caught per + column *per batch*, so a bad row in the middle of a file cannot take the rest of the file with it. + + This is the narrow half of the two guards the profiler runs. The wide one still wraps the whole + measure stage, and still catches anything structural -- schema derivation, classification -- that + is not attributable to a single column. + + Statistics and probes are folded together because they read the same values, and extracting a + column out of a batch costs more than either measurement. Neither fills in + ``categorical.values``: that needs the roles, which classification has not assigned yet, so + :func:`quote_enumerations` adds them afterwards, from the vocabulary this kept. + """ + + def __init__(self, features: list[FeatureSchema], expected_rows: int | None = None) -> None: + self._accumulators: dict[str, ColumnAccumulator] = {} + self._features: list[FeatureSchema] = [] + self._failed: dict[str, Evidence] = {} + for feature in features: + # Parquet permits duplicate field names, and every map here is keyed by name. Measuring + # the first and skipping the rest makes which one wins deterministic instead of + # "whichever came last", and keeps stats and probes agreeing on the same one. + if feature.name in self._accumulators: + continue + self._accumulators[feature.name] = _accumulator_for(feature, expected_rows) + self._features.append(feature) + + def update(self, rows: list[dict[str, Any]]) -> None: + """Fold one batch of rows into every column's accumulator.""" + for feature in self._features: + if feature.name in self._failed: + continue + try: + self._accumulators[feature.name].update([row.get(feature.name) for row in rows]) + except Exception as exc: + self._failed[feature.name] = Evidence( + kind="error", + detail=( + f"column {feature.name!r} ({feature.dtype}) could not be measured: {type(exc).__name__}: {exc}" + ), + ) + + def finalize(self) -> ColumnMeasurements: + stats: dict[str, ColumnStats] = {} + probes: dict[str, ColumnProbes] = {} + vocabularies: dict[str, set[Any]] = {} + errors: list[Evidence] = [] + for feature in self._features: + failure = self._failed.get(feature.name) + if failure is not None: + errors.append(failure) + continue + accumulator = self._accumulators[feature.name] + try: + column, probe = accumulator.finalize() + except Exception as exc: + errors.append( + Evidence( + kind="error", + detail=( + f"column {feature.name!r} ({feature.dtype}) could not be summarised: " + f"{type(exc).__name__}: {exc}" + ), + ) + ) + continue + probes[feature.name] = probe + vocabulary = accumulator.vocabulary() + if vocabulary is not None: + vocabularies[feature.name] = vocabulary + if column is not None: + stats[feature.name] = column + return ColumnMeasurements(stats=stats, probes=probes, vocabularies=vocabularies, errors=errors) + + +def measure_columns(features: list[FeatureSchema], rows: list[dict[str, Any]]) -> ColumnMeasurements: + """Measure every top-level column over rows already in hand. + + The whole partition as a single batch. :class:`ColumnFold` is the same measurement taken as the + rows arrive; this is the shape for a caller that has them all anyway. + """ + fold = ColumnFold(features) + fold.update(rows) + return fold.finalize() + + +def quote_enumerations( + features: list[FeatureSchema], stats: dict[str, ColumnStats], vocabularies: dict[str, set[Any]] +) -> None: + """Fill in ``categorical.values`` for columns whose role makes them a controlled vocabulary. + + Runs after classification, because it needs the roles it gates on, and mutates ``stats`` in place + the way classification mutates ``features``. Deliberately fills in rather than redacting: skip + this pass and no values are stored, where a redaction pass that got skipped would leak them. + + Reads what the accumulators already kept rather than going back to the rows. That second pass + was the last thing tying the measure stage to a materialised partition, and it was re-deriving a + set the vocabulary had built and thrown away. + + Cardinality is only the size bound. It cannot be the permission, because it inverts on small + data -- in a three-row dataset every column holds under 32 distinct values, free text included, + so an entire column of prompts was quotable. The role says what a column *is*, at any size. + """ + for feature in features: + if feature.semantic_role not in _QUOTABLE_ROLES: + continue + column = stats.get(feature.name) + if column is None or column.categorical is None or column.categorical.distinct_count > _MAX_ENUM_VALUES: + continue + values = vocabularies.get(feature.name) + if values is None: + continue + column.categorical.values = sorted(str(value) for value in values) + + +class ColumnAccumulator: + """Measures one top-level column, over however many batches it is handed. + + ``update`` folds a batch in and keeps no reference to it; ``finalize`` turns what was folded into + the stored blocks. Splitting a column across calls gives the same answer as one call with all of + it, which is the property that lets the caller stop materialising a partition before measuring it. + + The base class is the whole measurement for a dtype with no statistics of its own — a struct, a + list, anything the dispatch does not recognise — because the content probes run over every column + regardless of type. Subclasses add their dtype's state by overriding ``_observe`` and ``_blocks``. + """ + + def __init__(self) -> None: + self.rows = 0 + self._nulls = 0 + self._non_empty = 0 + self._texts = 0 + self._extractable_answer = 0 + self._transcript_marker = 0 + + def update(self, values: list[Any]) -> None: + """Fold one batch of this column's values in, one entry per row.""" + present: list[Any] = [] + for value in values: + self.rows += 1 + if value is None: + self._nulls += 1 + else: + present.append(value) + if value not in ("", [], {}): + self._non_empty += 1 + text = _probe_text(value) + if text is not None: + self._texts += 1 + if _GSM8K_ANSWER.search(text) or _BOXED_ANSWER.search(text): + self._extractable_answer += 1 + if _TRANSCRIPT_MARKER.search(text): + self._transcript_marker += 1 + self._observe(present) + + def finalize(self) -> tuple[ColumnStats | None, ColumnProbes]: + """The column's stored measurements, and its probe counts. + + Stats are None when there was nothing worth measuring, which keeps the map sparse. Probes are + always returned: a column of nothing is a finding classification is entitled to read. + """ + blocks = self._blocks() + null_rate = self._nulls / self.rows if self.rows else 0.0 + column = ColumnStats(null_rate=null_rate, **blocks) + if not any(blocks.values()) and null_rate == 0.0: + column = None + return column, ColumnProbes( + rows=self.rows, + non_empty=self._non_empty, + texts=self._texts, + extractable_answer=self._extractable_answer, + transcript_marker=self._transcript_marker, + ) + + def _observe(self, present: list[Any]) -> None: + """Fold this batch's non-null values into the dtype's own state. The base column has none.""" + + def _blocks(self) -> dict[str, Any]: + """The dtype-specific ``ColumnStats`` blocks. The base column contributes none.""" + return {} + + def backfill_nulls(self, count: int) -> None: + """Charge this column ``count`` rows in which it was absent. + + A column that first appears in the fiftieth batch was null for every row before it, which is + exactly what a materialising reader computes with ``row.get(name)``. Counted rather than fed + as values, so discovering a column late costs a pair of additions and not a pass. + """ + self.rows += count + self._nulls += count + + def vocabulary(self) -> set[Any] | None: + """The distinct values, for a column that is a bounded vocabulary. None for one that is not, + which is every dtype without a notion of cardinality.""" + return None + + +class _Vocabulary: + """Distinct values, for as long as the column still looks like a controlled vocabulary. + + Stops the moment it stops looking like one and drops what it had, which is the whole point: + counting distinct values exactly means *retaining* them, so on a free-text column this set grows + to hold the column. Today that costs little, because the rows are held anyway and the set stores + pointers into them -- 2.6 MB beside 61.4 MB of resident rows. It is the fold this is becoming + that makes it matter: once a batch is folded and discarded, this set is the *sole owner* of every + value it kept, and two text columns cost 46.8 MB against 0.163 MB for every other accumulator + combined. Unbounded, it is the one thing that would make the fold O(rows) again. + + Three bounds rather than one. A count alone bounds cardinality but not bytes, and 1024 reasoning + traces is 32 MB. The middle bound does the real work: it asks what the column *is* rather than + how many values it holds, in the same way the role gate on quoting does. A vocabulary member is + short by nature, so a single long value settles the question on sight, which is why free-text + columns stop here almost immediately instead of after 1024 values. + + The values themselves are never handed out here. They are row content, gated on role rather than + on size, and :func:`quote_enumerations` adds them once classification has assigned one. + """ + + def __init__(self) -> None: + self._values: set[Any] = set() + self._bytes = 0 + self._saturated = False + + def update(self, present: list[Any]) -> None: + if self._saturated: + return + for value in present: + if isinstance(value, str) and len(value) > _MAX_VOCABULARY_VALUE_CHARS: + return self._saturate() + try: + if value in self._values: + continue + self._values.add(value) + except TypeError: + return self._saturate() # unhashable values (dicts / lists) have no cardinality signal + self._bytes += len(value) if isinstance(value, str) else _NON_STRING_VALUE_BYTES + if len(self._values) > _MAX_VOCABULARY_VALUES or self._bytes > _MAX_VOCABULARY_BYTES: + return self._saturate() + + def _saturate(self) -> None: + self._values = set() # release what was held; holding it is the cost this bound exists to cap + self._saturated = True + + def finalize(self) -> CategoricalStats | None: + return None if self._saturated else CategoricalStats(distinct_count=len(self._values)) + + def values(self) -> set[Any] | None: + """What was kept, or None once the column stopped being a vocabulary. + + Handing this out is what lets :func:`quote_enumerations` fill in an enumeration without a + second pass over the rows -- which it could only do while the rows were still there. + """ + return None if self._saturated else self._values + + +class StringAccumulator(ColumnAccumulator): + """A ``string`` column: length quantiles, corruption ratios, and a vocabulary if it has one.""" + + def __init__(self, expected_rows: int | None = None) -> None: + super().__init__() + self._lengths = _LengthHistogram() + self._vocabulary = _Vocabulary() + self._quality = _TextQualityCounters() + self._seen = 0 + self._sampled = 0 + # With the row count known -- parquet footers give it before a row is read -- the cycle is + # fixed now and the blocks spread evenly over the whole column. Without it there is no length + # to spread over yet, so the cycle starts at one block and doubles each time the sample + # fills: every row eligible at first, thinning as the column turns out to be long. Each + # sampled row then stands for `cycle / block` rows, which is what keeps the estimate + # unbiased rather than weighted toward the head where sampling was densest. + self._block = _QUALITY_SAMPLE_BLOCK + self._cycle = _quality_cycle(expected_rows) if expected_rows is not None else self._block + self._adaptive = expected_rows is None + + def _observe(self, present: list[Any]) -> None: + for value in present: + if isinstance(value, str): + self._lengths.add(len(value)) + if self._seen % self._cycle < self._block: + self._quality.add(value, self._cycle / self._block) + self._sampled += 1 + if self._adaptive and self._sampled >= _QUALITY_SAMPLE_ROWS: + self._cycle *= 2 + self._sampled = 0 + self._seen += 1 + self._vocabulary.update(present) + + def _blocks(self) -> dict[str, Any]: + text = quality = None + if self._seen: + text = TextStats(chars=self._lengths.quantiles()) + quality = self._quality.finalize() + return {"text": text, "quality": quality, "categorical": self._vocabulary.finalize()} + + def vocabulary(self) -> set[Any] | None: + return self._vocabulary.values() + + +class NumericAccumulator(ColumnAccumulator): + """An ``int*`` / ``uint*`` / ``float*`` column: running extrema and mean, plus a vocabulary.""" + + def __init__(self) -> None: + super().__init__() + self._min = math.inf + self._max = -math.inf + self._sum = 0.0 + self._count = 0 + self._vocabulary = _Vocabulary() + + def _observe(self, present: list[Any]) -> None: + for value in present: + # Non-finite floats (NaN / +-inf) are dropped: they serialize to JSON null and then fail + # to re-validate against NumericStats' required floats, making the profile unreadable on + # its next load. bool is an int in Python and is not a number here. + if isinstance(value, (int, float)) and not isinstance(value, bool) and math.isfinite(value): + number = float(value) + self._min = min(self._min, number) + self._max = max(self._max, number) + self._sum += number + self._count += 1 + self._vocabulary.update(present) + + def _blocks(self) -> dict[str, Any]: + numeric = None + if self._count: + numeric = NumericStats(min=self._min, max=self._max, mean=self._sum / self._count) + return {"numeric": numeric, "categorical": self._vocabulary.finalize()} + + def vocabulary(self) -> set[Any] | None: + return self._vocabulary.values() + + +class BoolAccumulator(ColumnAccumulator): + """A ``bool`` column. The column that decides unpaired_preference deserves a measured class + balance rather than no stats at all, and two values is a vocabulary by any reading.""" + + def __init__(self) -> None: + super().__init__() + self._vocabulary = _Vocabulary() + + def _observe(self, present: list[Any]) -> None: + self._vocabulary.update(present) + + def _blocks(self) -> dict[str, Any]: + return {"categorical": self._vocabulary.finalize()} + + def vocabulary(self) -> set[Any] | None: + return self._vocabulary.values() + + +class MessageAccumulator(ColumnAccumulator): + """A ``messages`` column: turn and length distributions, the roles seen, and chat-shape rates.""" + + def __init__(self) -> None: + super().__init__() + self._conversations = 0 + self._turns = _LengthHistogram() + self._content_chars = _LengthHistogram() + self._roles_seen: list[str] = [] + self._ends_with_assistant = 0 + self._valid_alternation = 0 + self._has_tool_calls = False + + def _observe(self, present: list[Any]) -> None: + for messages in present: + if not isinstance(messages, list): + continue + self._conversations += 1 + self._turns.add(len(messages)) + total_content = 0 + for message in messages: + if not isinstance(message, dict): + continue + role = _role_of(message) + if role is not None: + # Coerced to str because roles_seen is typed list[str] and a non-string role + # would fail validation. Reported verbatim otherwise: the contract is explicit + # that an unexpected role is the finding worth surfacing, not something to + # normalize away. + role = role if isinstance(role, str) else str(role) + if role not in self._roles_seen and len(self._roles_seen) < _MAX_ROLES_SEEN: + self._roles_seen.append(role) + total_content += _content_len(_message_field(message, "content", "value")) + # `.get` truthiness, not `in`: parquet materializes every declared struct field, so a + # schema that merely declares tool_calls would otherwise report tool use on every row. + if message.get("tool_calls") or role == "tool": + self._has_tool_calls = True + self._content_chars.add(total_content) + if messages and isinstance(messages[-1], dict) and _is_assistant_role(_role_of(messages[-1])): + self._ends_with_assistant += 1 + if _valid_alternation(messages): + self._valid_alternation += 1 + + def _blocks(self) -> dict[str, Any]: + if not self._conversations: + return {"messages": None} + return { + "messages": MessageStats( + turns=self._turns.quantiles(), + content_chars=self._content_chars.quantiles(), + roles_seen=self._roles_seen, + ends_with_assistant_rate=self._ends_with_assistant / self._conversations, + valid_alternation_rate=self._valid_alternation / self._conversations, + has_tool_calls=self._has_tool_calls, + ) + } + + +def _accumulator_for(feature: FeatureSchema, expected_rows: int | None = None) -> ColumnAccumulator: + """The accumulator that knows how to measure this column, dispatched once on its dtype. + + ``expected_rows`` is the partition's row count when it is known before reading -- only a string + column uses it, to space its quality blocks across the whole of itself rather than thinning + as it goes. + """ + if feature.dtype == "string": + return StringAccumulator(expected_rows) + if feature.dtype == "bool": + return BoolAccumulator() + if feature.dtype == "messages": + return MessageAccumulator() + if _is_numeric(feature.dtype): + return NumericAccumulator() + return ColumnAccumulator() + + +class DeferredAccumulator(ColumnAccumulator): + """A column whose dtype is not known until the last row has gone by. + + An accumulator is normally chosen *by* dtype, which a declared schema gives up front. An inferred + one does not: the observed types are unioned over the whole column and a disagreement widens to + ``json``, so the choice cannot be made while the choosing still matters. Deferring it is the only + resolution that neither reads the data twice nor decides from a prefix and hopes. + + So every shape is measured at once and the answer picked at the end. It costs no more per value + than choosing would have -- a string only ever reaches the string state, an int only the numeric + -- and the state it costs is four bounded structures per column rather than one. A column that + resolves to a shape nothing measured, or to ``json``, simply has no blocks, which is what the + dispatch would have produced for it anyway. + """ + + def __init__(self, name: str, expected_rows: int | None = None) -> None: + super().__init__() + self._schema = SchemaFold(name) + self._string = StringAccumulator(expected_rows) + self._numeric = NumericAccumulator() + self._bool = BoolAccumulator() + self._messages = MessageAccumulator() + + def _observe(self, present: list[Any]) -> None: + self._schema.update(present) + # Routed by python type. Where a dtype resolves to something measurable, every present value + # is of that type by construction -- `_resolve_scalar` only returns `string` when the whole + # column was strings -- so this sees exactly what the chosen accumulator would have seen. + strings = [value for value in present if isinstance(value, str)] + if strings: + self._string._observe(strings) + numbers = [value for value in present if isinstance(value, (int, float)) and not isinstance(value, bool)] + if numbers: + self._numeric._observe(numbers) + bools = [value for value in present if isinstance(value, bool)] + if bools: + self._bool._observe(bools) + lists = [value for value in present if isinstance(value, list)] + if lists: + self._messages._observe(lists) + + def feature(self) -> FeatureSchema: + """The column's schema, as folded.""" + return self._schema.finalize() + + def _blocks(self) -> dict[str, Any]: + dtype = self.feature().dtype + if dtype == "string": + return self._string._blocks() + if dtype == "bool": + return self._bool._blocks() + if dtype == "messages": + return self._messages._blocks() + if _is_numeric(dtype): + return self._numeric._blocks() + return {} + + def vocabulary(self) -> set[Any] | None: + dtype = self.feature().dtype + if dtype == "string": + return self._string.vocabulary() + if dtype == "bool": + return self._bool.vocabulary() + if _is_numeric(dtype): + return self._numeric.vocabulary() + return None + + +class InferredColumnFold: + """A partition's columns, discovered as they appear and typed once they have all gone by. + + The counterpart to :class:`ColumnFold` for data that declares no schema. Columns are created on + first sight and back-filled with the rows they were absent for, which is what makes the result + identical to inferring the schema first and measuring second -- a row without the key genuinely + holds a null for it. + """ + + def __init__(self, expected_rows: int | None = None) -> None: + self._expected_rows = expected_rows + self._accumulators: dict[str, DeferredAccumulator] = {} + self._order: list[str] = [] + self._failed: dict[str, Evidence] = {} + self._rows_seen = 0 + + def update(self, rows: list[dict[str, Any]]) -> None: + for row in rows: + for name in row: + if name in self._accumulators or len(self._accumulators) >= MAX_COLUMNS: + continue + accumulator = DeferredAccumulator(name, self._expected_rows) + accumulator.backfill_nulls(self._rows_seen) + self._accumulators[name] = accumulator + self._order.append(name) + for name in self._order: + if name in self._failed: + continue + try: + self._accumulators[name].update([row.get(name) for row in rows]) + except Exception as exc: + self._failed[name] = Evidence( + kind="error", + detail=f"column {name!r} could not be measured: {type(exc).__name__}: {exc}", + ) + self._rows_seen += len(rows) + + def finalize(self) -> tuple[list[FeatureSchema], ColumnMeasurements]: + features: list[FeatureSchema] = [] + stats: dict[str, ColumnStats] = {} + probes: dict[str, ColumnProbes] = {} + vocabularies: dict[str, set[Any]] = {} + errors: list[Evidence] = [] + for name in self._order: + failure = self._failed.get(name) + if failure is not None: + errors.append(failure) + continue + accumulator = self._accumulators[name] + try: + features.append(accumulator.feature()) + column, probe = accumulator.finalize() + except Exception as exc: + errors.append( + Evidence( + kind="error", + detail=f"column {name!r} could not be summarised: {type(exc).__name__}: {exc}", + ) + ) + continue + probes[name] = probe + vocabulary = accumulator.vocabulary() + if vocabulary is not None: + vocabularies[name] = vocabulary + if column is not None: + stats[name] = column + return features, ColumnMeasurements(stats=stats, probes=probes, vocabularies=vocabularies, errors=errors) + + +def _is_numeric(dtype: str) -> bool: + return dtype.startswith(("int", "uint", "float")) + + +# How finely a length distribution is recorded. Lengths below the slice count get a counter each and +# are exact; above it, each octave is cut into this many slices, so a bucket spans a fixed *relative* +# width of 1/32. Reporting a bucket's midpoint then puts every estimate within ~1.6% of the truth, +# whatever the value's magnitude and however many rows there are. +_HISTOGRAM_SLICE_BITS = 5 +_HISTOGRAM_SLICES = 1 << _HISTOGRAM_SLICE_BITS + + +def _length_bucket(value: int) -> int: + """The counter a length belongs to.""" + if value < _HISTOGRAM_SLICES: + return value # small lengths get a counter each, so they are recorded exactly + shift = value.bit_length() - 1 - _HISTOGRAM_SLICE_BITS + return (shift + 1) * _HISTOGRAM_SLICES + ((value >> shift) - _HISTOGRAM_SLICES) + + +def _bucket_bounds(bucket: int) -> tuple[int, int]: + """The half-open range of lengths that land in ``bucket``. Inverse of :func:`_length_bucket`.""" + if bucket < _HISTOGRAM_SLICES: + return bucket, bucket + 1 + index, slice_index = divmod(bucket, _HISTOGRAM_SLICES) + shift = index - 1 + low = (_HISTOGRAM_SLICES + slice_index) << shift + return low, low + (1 << shift) + + +class _LengthHistogram: + """A per-row length distribution, held as counters rather than as the lengths themselves. + + This is what lets an accumulator stay O(1) in rows. Exact quantiles need every length kept and + sorted, which is a list that grows with the dataset; a reservoir of sampled lengths bounds that, + but buys the bound with an RNG -- and so with a seed back in the contract, and two runs over the + same bytes disagreeing. Counting into fixed buckets bounds it with neither. + + The two put their error in different places. A reservoir sees *some* rows exactly, so its error + is in which rows it happened to keep: probabilistic, and shrinking only with the sample size. + This sees *every* row imprecisely, so its error is in how finely each value was recorded: a hard + bound of half a bucket width, whatever the data does. Measured against exact quantiles on real + shards, ~2%. + + Rounding the value is the cheap error to accept here, because the number is read to pick a + sequence budget and gets rounded to a power of two by whoever reads it. ``max`` is kept exactly + and separately: it is the one value here a reader may treat as a hard bound. + """ + + def __init__(self) -> None: + self._counts: dict[int, int] = {} + self._rows = 0 + self._max = 0 + + def add(self, value: int) -> None: + bucket = _length_bucket(value) + self._counts[bucket] = self._counts.get(bucket, 0) + 1 + self._rows += 1 + if value > self._max: + self._max = value + + def quantiles(self) -> Quantiles: + return Quantiles(p50=self._at(50), p95=self._at(95), p99=self._at(99), max=self._max) + + def _at(self, percentile: int) -> int: + """Nearest-rank percentile: the bucket the p-th row falls in, reported at its midpoint. + + The rank is exact -- every row was counted, none sampled -- so only the value is approximate. + Midpoint rather than the bucket's low edge, which sits systematically under the truth and + roughly doubles the average error. + """ + if not self._rows: + return 0 + target = math.ceil(percentile / 100 * self._rows) + seen = 0 + for bucket in sorted(self._counts): + seen += self._counts[bucket] + if seen >= target: + low, high = _bucket_bounds(bucket) + # Never above `max`: a midpoint can overshoot the largest value actually present, + # and a p99 above the maximum would be nonsense. + return min((low + high) // 2, self._max) + return self._max + + +# --- text quality -------------------------------------------------------------------------------- + + +_WHITESPACE_RUN = re.compile(r"\s") +_NON_ASCII_RUN = re.compile(r"[^\x00-\x7f]") +# Any character repeated four or more times in a row. Scanning with the regex engine instead of a +# Python loop is what keeps this affordable: these three measurements used to run three interpreted +# passes over every character of every string and dominated total profiling time. +_REPEAT_RUN = re.compile(r"(.)\1{3,}", re.DOTALL) + +# What `\s` matches within ASCII, in the order `str.count` will be asked for them. +_ASCII_WHITESPACE = " \t\n\r\f\v" + +# Rows a column's quality ratios are measured over. These three are the only per-character work left +# in the profiler -- measured at 37x the cost of every content probe combined, and roughly fifteen +# times everything else in a column's measurement put together -- while every other statistic is +# O(1) per row. They are also ratios, which a sample of tens of thousands of rows pins down far past +# the precision anyone reads them to. Bounding them is what makes reading every row affordable. +_QUALITY_SAMPLE_ROWS = 50_000 + +# ...and the sample is taken in contiguous blocks of this many rows, not at an even step. +# +# A step aliases. Data is periodic more often than it looks -- a set that round-robins over ten +# sources, or carries k responses per prompt, is periodic by construction -- and a step that shares a +# factor with the period samples one phase and only that phase. Measured before this was blocks: +# 500,000 rows with every tenth corrupt gives a step of ten, which reported a repetition score of +# 1.000 against a truth of 0.100. Not noise; the wrong answer. +# +# A block longer than the period sees every phase of it, whatever the period is, and costs exactly +# the same. 512 covers anything plausible -- k-per-prompt is single digits, round-robin over sources +# is tens to low hundreds. +_QUALITY_SAMPLE_BLOCK = 512 + + +def _whitespace_count(text: str) -> int: + """Whitespace characters, matching ``\\s`` exactly. + + The ASCII branch is not merely faster, it is the only one that may take the shortcut: within + ASCII ``\\s`` is precisely :data:`_ASCII_WHITESPACE`, so counting those six literals in C is the + same measurement. Outside it, ``\\s`` also matches U+00A0 and the rest of Unicode's spaces, which + the literal count would silently miss -- so the regex is a correctness fallback, not a slow path + kept for tidiness. + """ + if text.isascii(): + return sum(text.count(char) for char in _ASCII_WHITESPACE) + return _count_matches(_WHITESPACE_RUN, text) + + +def _non_ascii_count(text: str) -> int: + """Characters outside ASCII. ``str.isascii`` settles the common case in C without a scan. + + Deliberately not ``len(text.encode()) - len(text)``, which is faster still and answers a + different question: that counts *bytes* of encoding overhead, so a three-byte codepoint would + contribute two where this contributes one. + """ + if text.isascii(): + return 0 + return _count_matches(_NON_ASCII_RUN, text) + + +class _TextQualityCounters: + """The three corruption ratios as running sums, so a sampled subset needs no storage. + + Every denominator is the sample's own, never the column's: each ratio is an estimate over the + rows actually scanned, which is what keeps it unbiased rather than diluted. + """ + + def __init__(self) -> None: + self._chars = 0 + self._whitespace = 0 + self._non_ascii = 0 + self._repetition = 0.0 + self._rows = 0 + + def add(self, text: str, weight: float = 1.0) -> None: + """Fold one sampled row in, standing for ``weight`` rows of the column. + + The weight is what keeps a *varying* sample rate honest. Sampling one row in four and + counting it once would let a densely sampled stretch outvote a sparsely sampled one; counting + it four times estimates the population sums instead, and the ratios come out unbiased. + """ + self._chars += weight * len(text) + self._whitespace += weight * _whitespace_count(text) + self._non_ascii += weight * _non_ascii_count(text) + self._repetition += weight * _repetition_score(text) + self._rows += weight + + def finalize(self) -> TextQuality: + return TextQuality( + whitespace_ratio=self._whitespace / self._chars if self._chars else 0.0, + non_ascii_ratio=self._non_ascii / self._chars if self._chars else 0.0, + repetition_score=self._repetition / self._rows if self._rows else 0.0, + ) + + +def _quality_cycle(rows: int) -> int: + """How many rows one sampled block stands in for, when the column's length is known. + + Every ``cycle`` rows, the first ``_QUALITY_SAMPLE_BLOCK`` of them are measured. Sized so the + blocks add up to the sample budget and spread across the whole column. + """ + return max(_QUALITY_SAMPLE_BLOCK, rows * _QUALITY_SAMPLE_BLOCK // _QUALITY_SAMPLE_ROWS) + + +def _count_matches(pattern: re.Pattern[str], text: str) -> int: + return sum(1 for _ in pattern.finditer(text)) + + +def _repetition_score(text: str) -> float: + """Fraction of characters inside a run of the same character of length >= 4. + + A cheap corruption proxy: near zero for natural text, high for scraping junk and degenerate + single-character loops (``"aaaaaa"``, long ``"------"`` separators). + """ + if not text: + return 0.0 + redundant = sum(len(match.group(0)) for match in _REPEAT_RUN.finditer(text)) + return redundant / len(text) + + +# --- messages ------------------------------------------------------------------------------------ + +# Role strings that mean "the turn the model is trained to produce". Matching only the literal +# "assistant" made every chat dataset using another convention (ShareGPT's gpt, or bot/model) look +# like it ended on a user turn, which classification reads as a prompt-only dataset with no training +# target — a false negative over a large slice of public chat data. +_ASSISTANT_ROLES = {"assistant", "gpt", "bot", "model", "chatbot", "ai"} + +# Distinct role strings a chat column may show before the list stops growing. It is fed straight from +# row content, so without a bound one malformed column could hold a string per message -- and since +# membership is checked against the list, that is quadratic as well as unbounded. The truncation +# costs nothing a reader would act on: the list exists to pick a chat template, and a column with +# more than this many roles is not a chat column, which the first few dozen already say. +_MAX_ROLES_SEEN = 64 + + +def _message_field(message: dict, *names: str) -> Any: + """The first present, non-null value among ``names``. + + Chat rows spell the same two fields either ``{role, content}`` or ``{from, value}``. Reading with + a plain ``.get`` default is not enough: parquet materializes *every* declared struct field, so an + absent field arrives as an explicit None rather than a missing key. + """ + for name in names: + value = message.get(name) + if value is not None: + return value + return None + + +def _role_of(message: dict) -> Any: + return _message_field(message, "role", "from") + + +def _is_assistant_role(role: Any) -> bool: + return isinstance(role, str) and role.lower() in _ASSISTANT_ROLES + + +def _content_len(content: Any) -> int: + if isinstance(content, str): + return len(content) + if isinstance(content, list): # VLM content as a list of typed parts + return sum( + len(part["text"]) for part in content if isinstance(part, dict) and isinstance(part.get("text"), str) + ) + return 0 + + +def _valid_alternation(messages: list) -> bool: + """True when user/assistant turns alternate (ignoring any leading system turns).""" + roles = [_role_of(m) for m in messages if isinstance(m, dict) and _role_of(m) != "system"] + return all(roles[i] != roles[i + 1] for i in range(len(roles) - 1)) + + +# --- content probes ------------------------------------------------------------------------------ + +# Probes run over *every* column, not only role-assigned ones. Gating them on roles made a content +# signal reachable only through a recognized column name: a dataset whose answer column is called +# `a` instead of `answer` lost verifiability entirely, even though the markers were sitting in the +# data and the regex would have matched them. Classification reads these counts and decides what +# they mean; it no longer does the looking. +_TRANSCRIPT_MARKER = re.compile(r"\n\n(?:Human|Assistant|User):") +_GSM8K_ANSWER = re.compile(r"####\s*-?[\d.,/]+\s*$") +_BOXED_ANSWER = re.compile(r"\\boxed\{") + + +@dataclass(frozen=True) +class ColumnProbes: + """What the content probes saw in one column across the sampled rows. + + Internal to the profiler rather than part of the stored contract: these are inputs to + classification, and promoting them to durable per-column facts is a separate contract change. + Counts, not rates — the caller divides, so a zero denominator stays visible instead of becoming + a silent 0.0. + """ + + rows: int # rows considered for this column + non_empty: int # value present and not "" / [] / {} — a usable target of any dtype + texts: int # rows that yielded text: a string, or a chat column's final turn + extractable_answer: int # of `texts`, how many carry `#### ` or `\boxed{` + transcript_marker: int # of `texts`, how many embed a Human:/Assistant: transcript + + +def _probe_text(value: Any) -> str | None: + """The text a probe reads from one cell: the string itself, or a chat column's final turn. + + The final turn is read through :func:`_message_field`, so ShareGPT's ``{from, value}`` spelling + works like ``{role, content}``. Both are handled everywhere else in this module and in schema + derivation; missing it here cost every ShareGPT-shaped dataset its verifiability. + """ + if isinstance(value, str): + return value + if isinstance(value, list) and value and isinstance(value[-1], dict): + content = _message_field(value[-1], "content", "value") + if isinstance(content, str): + return content + return None diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/tasks/profile/__main__.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/tasks/profile/__main__.py new file mode 100644 index 0000000000..b54c16c6ef --- /dev/null +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/tasks/profile/__main__.py @@ -0,0 +1,23 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Module entry point: ``python -m nemo_datasets_plugin.tasks.profile``.""" + +import logging +import signal +import sys +from types import FrameType + +from nemo_datasets_plugin.tasks.profile.run import run + +logger = logging.getLogger(__name__) + + +def _shutdown_handler(signum: int, frame: FrameType | None) -> None: + logger.warning("Received shutdown signal (%s). Shutting down gracefully.", signum) + sys.exit(128 + signum) + + +if __name__ == "__main__": + signal.signal(signal.SIGTERM, _shutdown_handler) + sys.exit(run()) diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/tasks/profile/run.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/tasks/profile/run.py new file mode 100644 index 0000000000..6c977d4751 --- /dev/null +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/tasks/profile/run.py @@ -0,0 +1,174 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Dataset-profiler task. + +Runs as a platform job: it reads a step config naming what to profile, runs the profiler, and +publishes the resulting ``DatasetProfile`` as a job result artifact named ``profile`` +(``profile.json``). + +This is deliberately *not* a ``nemo`` CLI subcommand. The profiler is new enough that its inputs and +its output contract are both still moving, and a published subcommand is a promise to keep them +still. A task module is invoked by the platform and by tests, which is the whole audience today. + +Only a local directory is profiled here. Reading a platform fileset through ranged requests, and +storing the profile back onto that fileset, both need Files-service surface this plugin does not +depend on; they arrive with the Files integration and change only :func:`_build_source` and the +publish step. The profiler core stays blind to where its bytes come from — that is what the +``FileSource`` seam is for. +""" + +from __future__ import annotations + +import json +import logging +import os +import tempfile +from pathlib import Path + +from nemo_datasets_plugin.profiler.file_source import FileSource, LocalFileSource +from nemo_datasets_plugin.profiler.pipeline import DEFAULT_ROW_BUDGET, profile +from nemo_platform import NeMoPlatform +from nemo_platform_plugin.job_results import PlatformJobResults +from nemo_platform_plugin.jobs.constants import ( + EPHEMERAL_TASK_STORAGE_PATH_ENVVAR, + NEMO_JOB_ID_ENVVAR, + NEMO_JOB_STEP_CONFIG_FILE_PATH_ENVVAR, + NEMO_JOB_WORKSPACE_ENVVAR, +) +from nemo_platform_plugin.sdk_provider import get_platform_sdk + +logger = logging.getLogger(__name__) + +# The service identity the task authenticates as. Any ``service:*`` principal is granted the internal +# ``ServiceSystem`` role, so no registration is required; a dedicated name just keeps audit logs and +# traces attributable. +_SERVICE_IDENTITY = "datasets" + +# Result artifact published back to the job's fileset. +_PROFILE_RESULT_NAME = "profile" +_PROFILE_FILE_NAME = "profile.json" + + +def run(sdk: NeMoPlatform | None = None) -> int: + """Entry point for the profiler task. Returns a process exit code.""" + _configure_logging() + try: + service_sdk = sdk or get_platform_sdk(as_service=_SERVICE_IDENTITY) + config = _load_step_config() + return _profile_and_publish( + service_sdk, + source=_build_source(config), + workspace=config.get("workspace") or _required_env(NEMO_JOB_WORKSPACE_ENVVAR), + job_name=_required_env(NEMO_JOB_ID_ENVVAR), + row_budget=_resolve_row_budget(config), + column_roles=_resolve_column_roles(config), + ) + except Exception: + logger.exception("Dataset profiler task failed") + return 1 + + +def _profile_and_publish( + sdk: NeMoPlatform, + *, + source: FileSource, + workspace: str, + job_name: str, + row_budget: int | None, + column_roles: dict[str, str], +) -> int: + logger.info("Profiling with a row budget of %s per partition", row_budget if row_budget else "unbounded") + dataset_profile = profile(source, row_budget=row_budget, column_roles=column_roles) + + # Scoped to the job's ephemeral storage when the platform provided one, and cleaned up either + # way — under the local subprocess backend this runs on a developer's machine, where an + # abandoned mkdtemp accumulates one directory per profiling run. + with tempfile.TemporaryDirectory( + prefix="dataset-profile-", + dir=os.environ.get(EPHEMERAL_TASK_STORAGE_PATH_ENVVAR) or None, + ) as scratch: + result_dir = Path(scratch) / _PROFILE_RESULT_NAME + result_dir.mkdir(parents=True) + (result_dir / _PROFILE_FILE_NAME).write_text(dataset_profile.model_dump_json(indent=2)) + + results = PlatformJobResults(job_name=job_name, workspace=workspace, sdk=sdk) + ref = results.save(_PROFILE_RESULT_NAME, result_dir) + logger.info("Published dataset profile: %s", ref.artifact_url) + return 0 + + +def _build_source(config: dict) -> FileSource: + """The files to profile, as named by the step config.""" + path = _required_config(config, "path") + try: + return LocalFileSource(path) + except NotADirectoryError as exc: + raise RuntimeError(f"step config 'path' must name a directory: {exc}") from exc + + +def _resolve_row_budget(config: dict) -> int | None: + """Rows the profiler may read per partition, from the step config. + + Defaults to reading everything, which it did not used to. A partition was materialised before it + was measured, so an uncapped run held every row of every file at roughly 20x the on-disk parquet + size and a large fileset killed the job outright. The profiler folds now: memory is flat in rows, + so the only thing a budget buys is a shorter run, and the default should not be to answer a + question worse than it can be answered. + + ``0`` and ``null`` both ask for every row -- the same thing the default does -- and are kept so a + caller that was setting them explicitly still means what it meant. + """ + if "row_budget" not in config: + return DEFAULT_ROW_BUDGET + requested = config["row_budget"] + if requested is None: + return None + # Validated here as well as at any API boundary that produced it: this reads a file off disk, so + # nothing upstream is guaranteed to have checked it. + budget = int(requested) + if budget < 0: + raise ValueError(f"row_budget must be >= 0, got {budget}") + return budget or None + + +def _resolve_column_roles(config: dict) -> dict[str, str]: + """Caller-declared column roles, for datasets whose column names the role table does not know. + + Not validated against the role vocabulary here. The profiler applies its own dtype gates and + reports a hint the data cannot support as evidence on the profile, which is a better place for + the finding than a task that fails before producing anything. + """ + roles = config.get("column_roles") or {} + if not isinstance(roles, dict): + raise ValueError(f"column_roles must map column name to role, got {type(roles).__name__}") + return {str(name): str(role) for name, role in roles.items()} + + +def _load_step_config() -> dict: + path = os.environ.get(NEMO_JOB_STEP_CONFIG_FILE_PATH_ENVVAR) + if not path: + raise RuntimeError(f"{NEMO_JOB_STEP_CONFIG_FILE_PATH_ENVVAR} not set; running outside the platform?") + with open(path, "r", encoding="utf-8") as f: + return json.load(f) + + +def _required_config(config: dict, key: str) -> str: + value = config.get(key) + if not value: + raise RuntimeError(f"Step config is missing '{key}'; nothing says what to profile.") + return str(value) + + +def _required_env(name: str) -> str: + value = os.environ.get(name) + if not value: + raise RuntimeError(f"Missing required job environment variable: {name}") + return value + + +def _configure_logging() -> None: + logging.basicConfig( + level=logging.INFO, + format="%(asctime)s %(levelname)s %(name)s: %(message)s", + ) diff --git a/plugins/nemo-datasets/tests/test_classify.py b/plugins/nemo-datasets/tests/test_classify.py new file mode 100644 index 0000000000..304852b5d8 --- /dev/null +++ b/plugins/nemo-datasets/tests/test_classify.py @@ -0,0 +1,395 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for classification: role assignment, format/prompt-form axes, and dataset type.""" + +from nemo_datasets_plugin.profiler.classify import PrefixPairFold, classify +from nemo_datasets_plugin.profiler.stats import measure_columns +from nemo_platform_plugin.files.dataset_profile import ( + CategoricalStats, + ColumnStats, + FeatureSchema, + MessageStats, + Quantiles, +) + + +def _probes(features, rows): + return measure_columns(features, rows).probes + + +def classify_rows(features, stats, rows, **kwargs): + """Classify from rows the way the pipeline does: fold first, then interpret the folds.""" + prefix = PrefixPairFold() + prefix.update(rows) + return classify(features, stats, probes=_probes(features, rows), prefix_pair=prefix.result(), **kwargs) + + +def _f(name, dtype): + return FeatureSchema(name=name, dtype=dtype) + + +def _binary_column(): + """A column observed to hold two distinct values -- what makes an int/bool a real label.""" + return ColumnStats(categorical=CategoricalStats(distinct_count=2)) + + +def _messages_column(ends_with_assistant_rate): + q = Quantiles(p50=1, p95=1, p99=1, max=1) + return ColumnStats( + messages=MessageStats( + turns=q, + content_chars=q, + roles_seen=["user", "assistant"], + ends_with_assistant_rate=ends_with_assistant_rate, + valid_alternation_rate=1.0, + ) + ) + + +# --- roles --------------------------------------------------------------------------------------- + + +def test_roles_assigned_by_name_and_dtype(): + features = [_f("prompt", "string"), _f("response", "string"), _f("helpfulness", "int64")] + classify(features, {}) + assert [f.semantic_role for f in features] == ["prompt", "completion", "score"] + + +def test_dtype_gate_rejects_mismatched_aliases(): + # "label" only counts as a label when boolean; a string column named "messages" is not messages. + features = [_f("label", "string"), _f("messages", "string")] + classify(features, {}) + assert all(f.semantic_role is None for f in features) + + +def test_physical_name_differs_from_role(): + features = [_f("response", "string")] + classify(features, {}) + assert features[0].semantic_role == "completion" + + +# --- format axis --------------------------------------------------------------------------------- + + +def test_format_standard_conversational_and_mixed(): + assert classify([_f("prompt", "string"), _f("completion", "string")], {}).format == "standard" + assert classify([_f("prompt", "messages"), _f("completion", "messages")], {}).format == "conversational" + mixed = [_f("prompt", "string"), _f("chosen", "messages"), _f("rejected", "messages")] + assert classify(mixed, {}).format == "mixed" + + +# --- dataset type + prompt form ------------------------------------------------------------------ + + +def test_prompt_completion_with_explicit_prompt(): + result = classify([_f("prompt", "string"), _f("completion", "string")], {}) + assert result.dataset_type == "prompt_completion" + assert result.prompt_form == "explicit" + + +def test_preference_pair_is_implicit_without_a_prompt(): + result = classify([_f("chosen", "string"), _f("rejected", "string")], {}) + assert result.dataset_type == "preference_pair" + assert result.prompt_form == "implicit" + + +def test_scored_response_beats_prompt_completion(): + features = [ + _f("prompt", "string"), + _f("response", "string"), + _f("helpfulness", "int64"), + _f("correctness", "int64"), + ] + assert classify(features, {}).dataset_type == "scored_response" + + +def test_unpaired_preference_accepts_a_boolean_label(): + features = [_f("prompt", "string"), _f("completion", "string"), _f("label", "bool")] + assert classify(features, {}).dataset_type == "unpaired_preference" + + +def test_unpaired_preference_accepts_a_binary_integer_label(): + # 0/1 is the usual on-disk encoding; requiring a bool made unpaired_preference unreachable for + # most real datasets. + features = [_f("prompt", "string"), _f("completion", "string"), _f("label", "int64")] + stats = {"label": ColumnStats(categorical=CategoricalStats(distinct_count=2))} + assert classify(features, stats).dataset_type == "unpaired_preference" + assert features[2].semantic_role == "label" + + +def test_wide_integer_label_is_not_a_preference_label(): + # A multi-class index or a rating is a different claim from a binary preference. + features = [_f("prompt", "string"), _f("completion", "string"), _f("label", "int64")] + stats = {"label": ColumnStats(categorical=CategoricalStats(distinct_count=7))} + assert classify(features, stats).dataset_type == "prompt_completion" + assert features[2].semantic_role is None + + +# --- rank ------------------------------------------------------------------------------------ + + +def test_rank_needs_something_to_rank(): + # A lone numeric column named "rank" used to short-circuit every more specific structure. + features = [_f("rank", "int64")] + assert classify(features, {}).dataset_type == "unknown" + + +def test_rank_does_not_override_a_preference_pair(): + features = [_f("chosen", "string"), _f("rejected", "string"), _f("rank", "int64")] + assert classify(features, {}).dataset_type == "preference_pair" + + +def test_rank_does_not_override_scored_responses(): + features = [_f("prompt", "string"), _f("response", "string"), _f("helpfulness", "int64"), _f("rank", "int64")] + assert classify(features, {}).dataset_type == "scored_response" + + +def test_rank_alongside_a_completion_is_ranked_responses(): + features = [_f("prompt", "string"), _f("completion", "string"), _f("rank", "int64")] + assert classify(features, {}).dataset_type == "ranked_responses" + + +def test_messages_ending_on_assistant_is_messages_type(): + result = classify([_f("messages", "messages")], {"messages": _messages_column(1.0)}) + assert result.dataset_type == "messages" + assert result.prompt_form == "n/a" + + +def test_messages_ending_on_user_is_prompt_only(): + result = classify([_f("messages", "messages")], {"messages": _messages_column(0.0)}) + assert result.dataset_type == "prompt_only" + + +def test_single_text_column_is_text(): + assert classify([_f("text", "string")], {}).dataset_type == "text" + + +def test_unrecognized_columns_are_unknown(): + result = classify([_f("foo", "int64"), _f("bar", "int64")], {}) + assert result.dataset_type == "unknown" + assert result.prompt_form is None # no axes asserted for unknown data + + +# --- evidence ------------------------------------------------------------------------------------ + + +def test_classification_records_evidence(): + result = classify([_f("prompt", "string"), _f("completion", "string")], {}) + assert {e.kind for e in result.evidence} >= {"column_name", "column_dtype"} + + +# --- verifiability + content probes -------------------------------------------------------------- + + +def test_verifiability_extractable_gsm8k_answer(): + features = [_f("problem", "string"), _f("solution", "string")] + rows = [{"problem": "q", "solution": "steps #### 18"}, {"problem": "q", "solution": "no final answer"}] + result = classify_rows(features, {}, rows) + assert result.verifiability.method == "extractable_final_answer" + assert result.verifiability.coverage == 0.5 + + +def test_verifiability_boxed_answer(): + features = [_f("prompt", "string"), _f("completion", "string")] + result = classify_rows(features, {}, [{"prompt": "q", "completion": r"reasoning \boxed{42}"}]) + assert result.verifiability.method == "extractable_final_answer" + assert result.verifiability.coverage == 1.0 + + +def test_verifiability_ground_truth_column_coverage(): + features = [_f("prompt", "string"), _f("ground_truth", "string")] + rows = [{"prompt": "q", "ground_truth": "42"}, {"prompt": "q", "ground_truth": None}] + result = classify_rows(features, {}, rows) + assert result.verifiability.method == "ground_truth_column" + assert result.verifiability.coverage == 0.5 + + +def test_no_verifiability_without_a_target(): + features = [_f("prompt", "string"), _f("completion", "string")] + result = classify_rows(features, {}, [{"prompt": "q", "completion": "just prose, no answer"}]) + assert result.verifiability is None + + +def test_verifiability_ignores_below_threshold_extractable_noise(): + # One coincidental "#### " in a large sample is noise, not a verifiable dataset (kto-mix-14k). + features = [_f("prompt", "string"), _f("completion", "string")] + rows = [{"prompt": "q", "completion": "just prose"} for _ in range(100)] + rows[0]["completion"] = "the answer is #### 7" # 1/100 = 1% < 5% floor + assert classify_rows(features, {}, rows).verifiability is None + + +def test_verifiability_asserted_above_coverage_floor(): + features = [_f("prompt", "string"), _f("completion", "string")] + rows = [{"prompt": "q", "completion": "just prose"} for _ in range(10)] + for row in rows[:2]: + row["completion"] = "answer #### 7" # 2/10 = 20% >= 5% floor + result = classify_rows(features, {}, rows) + assert result.verifiability.method == "extractable_final_answer" + assert result.verifiability.coverage == 0.2 + + +def test_sparse_ground_truth_falls_through_to_extractable_answer(): + # A ground_truth column present in too few rows must not mask a strong extractable-answer signal. + features = [_f("completion", "string"), _f("ground_truth", "string")] + rows = [{"completion": "reasoning #### 5", "ground_truth": None} for _ in range(100)] + rows[0]["ground_truth"] = "5" # 1/100 ground_truth coverage -> below floor, must fall through + result = classify_rows(features, {}, rows) + assert result.verifiability.method == "extractable_final_answer" + assert result.verifiability.coverage == 1.0 + + +def test_implicit_prompt_evidence_from_embedded_transcript(): + features = [_f("chosen", "string"), _f("rejected", "string")] + rows = [{"chosen": "\n\nHuman: hi\n\nAssistant: hello", "rejected": "\n\nHuman: hi\n\nAssistant: hey"}] + result = classify_rows(features, {}, rows) + assert result.prompt_form == "implicit" + assert any(e.kind == "content_probe" for e in result.evidence) + + +def test_ground_truth_may_be_a_container_dtype(): + # test_cases (list) and verification_info (struct) are verification targets, not free text, + # so the text-only dtype gate must not drop them. + features = [_f("prompt", "string"), _f("test_cases", "list"), _f("verification_info", "struct")] + classify(features, {}) + assert features[1].semantic_role == "ground_truth" + assert features[2].semantic_role == "ground_truth" + + +def test_container_ground_truth_drives_verifiability(): + features = [_f("prompt", "string"), _f("test_cases", "list")] + rows = [{"prompt": "q", "test_cases": [{"in": "1", "out": "2"}]}, {"prompt": "q2", "test_cases": []}] + result = classify_rows(features, {}, rows) + assert result.verifiability.method == "ground_truth_column" + assert result.verifiability.coverage == 0.5 # the empty test_cases list is not a usable target + + +def test_bare_scalar_ground_truth_alias_is_still_rejected(): + # A numeric column named "ground_truth" is far more likely a label/score than a target. + features = [_f("ground_truth", "int64")] + classify(features, {}) + assert features[0].semantic_role is None + + +def test_verifiability_survives_an_unrecognized_column_name(): + # Gating the probes on roles made a content signal reachable only through a recognized column + # name: the `#### ` markers were in the data and the regex would have matched them, but + # nothing knew where to look. The finding must name the column it came from. + features = [_f("q", "string"), _f("a", "string")] + rows = [{"q": "what is 2+2?", "a": f"add them #### {i}"} for i in range(10)] + result = classify_rows(features, {}, rows) + + assert {f.semantic_role for f in features} == {None} # still unroled, and honest about it + assert result.dataset_type == "unknown" + assert result.verifiability.method == "extractable_final_answer" + assert result.verifiability.coverage == 1.0 + assert "'a'" in result.verifiability.evidence[0].detail + + +def test_a_named_completion_still_decides_where_to_look(): + # Roles order the interpretation even though they no longer gate it: a column *known* to be the + # completion is a better answer than one that merely looks like it. + features = [_f("completion", "string"), _f("notes", "string")] + rows = [{"completion": "just prose", "notes": "scratch #### 9"} for _ in range(10)] + assert classify_rows(features, {}, rows).verifiability is None + + +def test_verifiability_reads_a_sharegpt_conversational_completion(): + features = [_f("prompt", "string"), _f("completion", "messages")] + rows = [{"prompt": "q", "completion": [{"from": "human", "value": "q"}, {"from": "gpt", "value": "#### 4"}]}] + result = classify_rows(features, {}, rows) + assert result.verifiability.method == "extractable_final_answer" + assert result.verifiability.coverage == 1.0 + + +def test_classification_without_probes_claims_nothing_rather_than_guessing(): + # classify() used to derive probes from rows when it was handed none. It no longer sees rows at + # all, so absent probes must read as "nothing was measured" -- never as "nothing is there". + rows = [{"prompt": "q", "completion": f"steps #### {i}"} for i in range(10)] + features = [_f("prompt", "string"), _f("completion", "string")] + + blind = classify(features, {}) + assert blind.verifiability is None + + measured = classify([_f("prompt", "string"), _f("completion", "string")], {}, probes=_probes(features, rows)) + assert measured.verifiability.method == "extractable_final_answer" + assert measured.verifiability.coverage == 1.0 + + +# --- candidates ------------------------------------------------------------------------------------ + + +def test_candidates_list_every_structure_the_roles_satisfy(): + # prompt + completion + score + label is genuinely both scored_response and unpaired_preference. + # Reporting only the first made rule order an invisible tie-break. + features = [_f("prompt", "string"), _f("completion", "string"), _f("score", "float64"), _f("label", "bool")] + result = classify(features, {"label": _binary_column()}) + + assert result.candidates == ["scored_response", "unpaired_preference", "prompt_completion"] + assert result.dataset_type == result.candidates[0] # the summary is the head, never more + + +def test_candidates_collapse_to_one_when_the_structure_is_unambiguous(): + result = classify([_f("prompt", "string"), _f("completion", "string")], {}) + assert result.candidates == ["prompt_completion"] + + +def test_unknown_is_still_reported_as_a_candidate(): + result = classify([_f("foo", "int64"), _f("bar", "int64")], {}) + assert result.dataset_type == "unknown" + assert result.candidates == ["unknown"] + + +def test_prompt_only_is_not_claimed_alongside_a_training_target(): + # Collecting candidates rather than returning early risks a prompt+completion set also claiming + # prompt_only, which asserts the opposite of what the data holds. + result = classify([_f("prompt", "string"), _f("completion", "string")], {}) + assert "prompt_only" not in result.candidates + + +# --- declared roles (hints) ------------------------------------------------------------------------ + + +def test_a_hint_names_a_column_the_alias_table_does_not_know(): + features = [_f("q", "string"), _f("a", "string")] + result = classify(features, {}, column_roles={"q": "prompt", "a": "completion"}) + + assert [(f.semantic_role, f.semantic_role_source) for f in features] == [ + ("prompt", "declared"), + ("completion", "declared"), + ] + assert result.dataset_type == "prompt_completion" + + +def test_a_hint_takes_precedence_over_the_name_alias(): + # The caller knows their schema; the table is ~35 English names. + features = [_f("prompt", "string")] + classify(features, {}, column_roles={"prompt": "context"}) + assert features[0].semantic_role == "context" + assert features[0].semantic_role_source == "declared" + + +def test_a_hint_the_dtype_cannot_support_is_rejected_loudly(): + # A hint says which column, not what the data is. Accepting it unconditionally would let one + # typo produce a nonsense classification, and silence is what made the table's misses costly. + features = [_f("n", "int64")] + result = classify(features, {}, column_roles={"n": "prompt"}) + + assert features[0].semantic_role is None + rejections = [e for e in result.evidence if e.kind == "user_hint"] + assert len(rejections) == 1 + assert "n -> prompt" in rejections[0].detail and "int64" in rejections[0].detail + + +def test_a_rejected_hint_falls_back_to_detection(): + # `answer` is a known alias; a bad hint on it must not cost the role the table would have found. + features = [_f("answer", "string")] + classify(features, {}, column_roles={"answer": "messages"}) # messages needs the messages dtype + assert features[0].semantic_role == "completion" + assert features[0].semantic_role_source == "detected" + + +def test_detected_roles_are_marked_as_detected(): + features = [_f("prompt", "string")] + classify(features, {}) + assert features[0].semantic_role_source == "detected" diff --git a/plugins/nemo-datasets/tests/test_pipeline.py b/plugins/nemo-datasets/tests/test_pipeline.py new file mode 100644 index 0000000000..5718dc6c8c --- /dev/null +++ b/plugins/nemo-datasets/tests/test_pipeline.py @@ -0,0 +1,919 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the profiling pipeline: split/partition resolution and envelope assembly.""" + +import json +from datetime import datetime, timezone +from pathlib import Path + +import pyarrow as pa +import pyarrow.parquet as pq +import pytest +from nemo_datasets_plugin.profiler.file_source import FileEntry, LocalFileSource +from nemo_datasets_plugin.profiler.partition import group_partitions +from nemo_datasets_plugin.profiler.pipeline import _expected_rows, _peek_files, profile +from nemo_datasets_plugin.profiler.readers.base import FilePreview +from nemo_datasets_plugin.profiler.splits import infer_data_files, resolve_splits +from nemo_platform_plugin.files.dataset_profile import DatasetProfile + +FIXED_TIME = datetime(2026, 7, 13, 12, 0, 0, tzinfo=timezone.utc) + + +def _write_parquet(path, rows): + path.parent.mkdir(parents=True, exist_ok=True) + pq.write_table(pa.Table.from_pylist(rows), path) + + +def _entries(*paths): + return [FileEntry(path=p, size_bytes=100) for p in paths] + + +# --- split resolution ---------------------------------------------------------------------------- + + +def test_resolve_splits_infers_canonical_from_sharded_names(): + entries = _entries( + "train-00000-of-00002.parquet", + "train-00001-of-00002.parquet", + "validation-00000-of-00001.parquet", + ) + splits = {s.name: s for s in resolve_splits(entries)} + assert set(splits) == {"train", "validation"} + assert splits["train"].canonical == "train" + assert splits["validation"].canonical == "validation" + assert len(splits["train"].entries) == 2 + + +def test_resolve_splits_normalizes_aliases(): + splits = {s.name: s.canonical for s in resolve_splits(_entries("val.jsonl", "dev.jsonl"))} + assert splits == {"val": "validation", "dev": "validation"} + + +def test_resolve_splits_does_not_mistake_years_for_shard_numbers(): + # A bare trailing number only reads as a shard when it is zero-padded; otherwise dates and + # versions were being stripped, e.g. covid-19.jsonl -> a "covid" split. + assert [s.name for s in resolve_splits(_entries("covid-19.jsonl"))] == ["default"] + names = {s.name for s in resolve_splits(_entries("train-00000-of-00002.parquet", "data-2024.jsonl"))} + assert names == {"train", "data-2024"} + + +def test_resolve_splits_falls_back_to_single_default(): + splits = resolve_splits(_entries("shard-00000.parquet", "shard-00001.parquet")) + assert len(splits) == 1 + assert splits[0].name == "default" + assert splits[0].canonical is None + assert len(splits[0].entries) == 2 + + +# --- data_files glob inference -------------------------------------------------------------------- + + +def _globs(*paths): + """Infer a glob per split over ``paths``, verifying against the full listing (README included).""" + data = [e for e in _entries(*paths) if e.path.endswith((".parquet", ".jsonl"))] + return {s.name: infer_data_files(s.name, s.entries, list(paths)) for s in resolve_splits(data)} + + +@pytest.mark.parametrize( + "label,paths,expected", + [ + ( + "shards in one directory", + ( + "data/train-00000-of-00002.parquet", + "data/train-00001-of-00002.parquet", + "data/test-00000-of-00001.parquet", + ), + {"train": "data/train*.parquet", "test": "data/test*.parquet"}, + ), + ( + "a directory per split", + ("default/train/0000.parquet", "default/train/0001.parquet", "default/test/0000.parquet"), + {"train": "default/train/*.parquet", "test": "default/test/*.parquet"}, + ), + ( + "files at the fileset root", + ("train.jsonl", "validation.jsonl"), + {"train": "train*.jsonl", "validation": "validation*.jsonl"}, + ), + ( + "no split detected: the glob covers the partition", + ("shard-00000.parquet", "shard-00001.parquet"), + {"default": "*.parquet"}, + ), + ( + "mixed formats drop the suffix rather than losing a file", + ("train-00000-of-00002.parquet", "train-00001-of-00002.jsonl"), + {"train": "train*"}, + ), + ], +) +def test_data_files_glob_per_layout(label, paths, expected): + assert _globs(*paths) == expected, label + + +def test_data_files_glob_excludes_a_readme_beside_the_shards(): + # `data/*` would sweep the card into the split. The suffix-qualified candidate is what survives + # verification, and verification runs against every listed file, not just the data ones. + assert _globs("data/train-00000-of-00001.parquet", "data/README.md") == {"train": "data/train*.parquet"} + + +def test_data_files_glob_keeps_a_separator_to_beat_a_sibling_split(): + # `train*` would also match train_prefs, so the simple form loses verification and the narrower + # `train-*` is reached. Both splits still get a pattern; neither over-matches the other. + assert _globs( + "train-00000-of-00002.parquet", "train-00001-of-00002.parquet", "train_prefs-00000-of-00001.parquet" + ) == {"train": "train-*.parquet", "train_prefs": "train_prefs*.parquet"} + + +def test_data_files_glob_refuses_rather_than_sweep_in_a_non_data_file(): + # Mixed suffixes leave no suffix to qualify with, and an unsplit-named set leaves no stem to + # anchor on, so the only candidate left is `data/*` -- which would hand a reader the README as + # if it were a shard. Verification is the whole of what stops that, and None is the answer. + assert _globs("data/shard-00000.parquet", "data/shard-00001.jsonl", "data/README.md") == {"default": None} + + +def test_data_files_glob_is_none_when_shards_span_subdirectories(): + # Expressing this needs `**`, whose meaning differs between glob implementations. None is the + # honest answer; a pattern that selects a different set in the reader than here would not be. + assert _globs("train/part-a/0000.parquet", "train/part-b/0000.parquet") == {"train": None} + + +def test_data_files_glob_means_the_same_thing_to_pythons_own_glob(tmp_path): + """The dialect claim, checked against an independent implementation rather than our matcher. + + A pattern is only worth emitting if a consumer resolves it to the files we said it selects. + """ + for rel in ( + "data/train-00000-of-00002.parquet", + "data/train-00001-of-00002.parquet", + "data/test-00000-of-00001.parquet", + ): + _write_parquet(tmp_path / rel, [{"a": 1}]) + (tmp_path / "data" / "README.md").write_text("card") + + for split in profile(LocalFileSource(tmp_path), created_at=FIXED_TIME).partitions[0].splits: + resolved = sorted(p.relative_to(tmp_path).as_posix() for p in Path(tmp_path).glob(split.data_files)) + assert len(resolved) == split.num_files, f"{split.name}: {split.data_files} -> {resolved}" + assert all(name.startswith(f"data/{split.name}") for name in resolved) + + +# --- the fold ------------------------------------------------------------------------------------ + + +def test_a_parquet_footer_declares_enough_to_fold_without_reading_rows(tmp_path): + # The footer is what makes a fold possible at all: the schema, so accumulators can exist before + # the first batch, and the exact row count, so the quality stride can be placed before the column + # it strides has been seen. A line-delimited file declares neither, which is why it materialises. + _write_parquet(tmp_path / "train.parquet", [{"a": i} for i in range(7)]) + (tmp_path / "extra.jsonl").write_text('{"a": 1}\n') + source = LocalFileSource(tmp_path) + + previews = _peek_files(source, source.list_files()) + + assert previews["train.parquet"].num_rows == 7 + assert previews["train.parquet"].arrow_schema is not None + assert previews["extra.jsonl"] == FilePreview() # declares nothing, so the partition cannot fold + + +def test_expected_rows_counts_what_the_read_will_actually_scan(): + # The stride has to be placed over the rows that will be *scanned*, not the rows the dataset + # holds, or a budgeted run would stride far too coarsely and sample almost nothing. + previews = {"a": FilePreview(num_rows=100), "b": FilePreview(num_rows=100)} + assert _expected_rows(previews, None) == 200 + assert _expected_rows(previews, 30) == 60 # capped per file, exactly as the read will be + assert _expected_rows({"a": FilePreview(num_rows=100), "b": FilePreview()}, None) is None + + +def test_the_folded_and_materialised_paths_measure_the_same_thing(tmp_path): + # Parquet declares a schema and is folded batch by batch; jsonl declares none and is + # materialised. The same rows have to measure the same either way, or the batch size -- an + # implementation detail no reader of a profile can see -- would be visible in the numbers. + rows = [{"prompt": f"question {i}", "completion": "answer " * (i % 7 + 1), "score": i % 5} for i in range(200)] + _write_parquet(tmp_path / "pq" / "train.parquet", rows) + (tmp_path / "jl").mkdir() + (tmp_path / "jl" / "train.jsonl").write_text("\n".join(json.dumps(row) for row in rows)) + + folded = profile(LocalFileSource(tmp_path / "pq"), created_at=FIXED_TIME).partitions[0] + materialised = profile(LocalFileSource(tmp_path / "jl"), created_at=FIXED_TIME).partitions[0] + + assert folded.stats == materialised.stats + assert folded.classification == materialised.classification + assert [f.model_dump() for f in folded.features] == [f.model_dump() for f in materialised.features] + + +def test_an_exhaustive_fold_does_not_cost_more_than_a_budgeted_one(tmp_path): + # The point of the whole exercise: reading every row costs what reading some of them costs, so + # the budget stops being a memory guard. Same measurements, and `rows_complete` finally true. + _write_parquet(tmp_path / "train.parquet", [{"t": f"row {i}" * (i % 5 + 1)} for i in range(5000)]) + + budgeted = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME, row_budget=500) + exhaustive = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME, row_budget=None) + + assert budgeted.sampling.rows_scanned == 500 + assert exhaustive.sampling.rows_scanned == 5000 + assert budgeted.partitions[0].rows_complete is False + assert exhaustive.partitions[0].rows_complete is True + # Exact where it claims to be exact: the longest row is found by reading all of them. + assert exhaustive.partitions[0].stats["t"].text.chars.max >= budgeted.partitions[0].stats["t"].text.chars.max + + +# --- partition grouping -------------------------------------------------------------------------- + + +def test_group_partitions_names_the_root_partition_with_the_empty_prefix(): + # "" is the path prefix root-level files share, and no directory can be named it -- which is what + # keeps root files distinct from a directory literally called "default". + assert group_partitions(_entries("train.parquet", "test.parquet")) == [ + ("", _entries("train.parquet", "test.parquet")) + ] + + +def test_group_partitions_collapses_single_container_dir(): + # One container directory is still one partition, and it keeps that directory as its name. + # Reporting "default" here discarded the only thing identifying the partition. + parts = group_partitions(_entries("data/train.parquet", "data/test.parquet")) + assert [name for name, _ in parts] == ["data"] + + +def test_group_partitions_splits_multiple_top_dirs(): + parts = group_partitions(_entries("main/train.parquet", "socratic/train.parquet")) + assert [name for name, _ in parts] == ["main", "socratic"] + + +def test_group_partitions_does_not_treat_split_dirs_as_partitions(): + # train/ and test/ are one dataset's splits, not two datasets. + parts = group_partitions(_entries("train/data.parquet", "test/data.parquet")) + assert [name for name, _ in parts] == [""] + + +def test_resolve_splits_reads_the_split_directory(): + # The data// layout names every shard the same thing; only the directory carries + # the split, so reading the stem alone would collapse the dataset into one split. + splits = {s.name: s for s in resolve_splits(_entries("data/train/0000.parquet", "data/test/0000.parquet"))} + assert set(splits) == {"train", "test"} + assert splits["train"].canonical == "train" + assert splits["test"].canonical == "test" + + +def test_resolve_splits_prefers_directory_over_stem(): + splits = resolve_splits(_entries("main/train-00000-of-00001.parquet")) + assert [s.name for s in splits] == ["train"] # no split dir on the path, so the stem is used + + +# --- end-to-end profile() ------------------------------------------------------------------------ + + +def test_profile_parquet_dataset_builds_envelope(tmp_path): + _write_parquet(tmp_path / "train-00000-of-00001.parquet", [{"prompt": "a"}, {"prompt": "b"}]) + _write_parquet(tmp_path / "validation-00000-of-00001.parquet", [{"prompt": "c"}]) + + result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) + + assert result.profiler_info["name"] == "nemo-dataset-profiler" + assert len(result.partitions) == 1 + partition = result.partitions[0] + assert partition.name == "" # root-level files: the empty path prefix + assert partition.file_formats == ["parquet"] + + splits = {s.name: s for s in partition.splits} + assert set(splits) == {"train", "validation"} + assert splits["train"].canonical == "train" + assert splits["train"].num_examples == 2 + assert splits["validation"].num_examples == 1 + assert splits["train"].num_files == 1 + + # Row schema, stats, and classification are all derived now. + assert [f.name for f in partition.features] == ["prompt"] + assert partition.features[0].dtype == "string" + assert partition.features[0].semantic_role == "prompt" + assert partition.stats["prompt"].text is not None + assert partition.classification.dataset_type == "prompt_only" # a lone prompt column, no target + + # A budgeted run over files that all fit under their share is still a complete scan, which is + # why the budget and the outcome are separate fields. + assert partition.rows_complete is True + assert result.sampling.rows_scanned == result.sampling.rows_present # exhaustive by default + assert result.sampling.rows_scanned == 3 + assert result.sampling.rows_present == 3 + assert result.sampling.files_read == result.sampling.files_present == 2 + + +def test_profile_jsonl_dataset_counts_rows_exactly(tmp_path): + (tmp_path / "train.jsonl").write_text('{"a": 1}\n{"a": 2}\n{"a": 3}\n') + result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) + + partition = result.partitions[0] + assert partition.file_formats == ["jsonl"] + assert partition.splits[0].name == "train" + assert partition.splits[0].num_examples == 3 + assert result.sampling.rows_scanned == 3 + + +def test_profile_multiple_directories_become_partitions(tmp_path): + _write_parquet(tmp_path / "main" / "train-00000-of-00001.parquet", [{"q": "1"}]) + _write_parquet(tmp_path / "socratic" / "train-00000-of-00001.parquet", [{"q": "2"}]) + + result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) + + assert [p.name for p in result.partitions] == ["main", "socratic"] + assert all(p.file_formats == ["parquet"] for p in result.partitions) + + +def test_profile_top_level_split_dirs_become_one_partition(tmp_path): + # train/ + test/ is one dataset with two splits, not two datasets. As separate partitions each + # would derive its own schema and classification, and the split structure would disappear. + _write_parquet(tmp_path / "train" / "data.parquet", [{"prompt": "a", "completion": "b"}]) + _write_parquet(tmp_path / "test" / "data.parquet", [{"prompt": "c", "completion": "d"}]) + + result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) + + assert [p.name for p in result.partitions] == [""] + splits = {s.name: s for s in result.partitions[0].splits} + assert set(splits) == {"train", "test"} + assert splits["train"].canonical == "train" + assert splits["test"].num_examples == 1 + + +def test_profile_nested_split_dirs_keep_splits_apart(tmp_path): + # data// shards are all named alike, so only the directory distinguishes them. + # Reading the stem alone pooled train and test into a single "default" split. + _write_parquet(tmp_path / "data" / "train" / "0000.parquet", [{"prompt": "a"}, {"prompt": "b"}]) + _write_parquet(tmp_path / "data" / "test" / "0000.parquet", [{"prompt": "c"}]) + + result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) + + assert [p.name for p in result.partitions] == ["data"] # the container directory, not "default" + splits = {s.name: s for s in result.partitions[0].splits} + assert set(splits) == {"train", "test"} + assert splits["train"].num_examples == 2 + assert splits["test"].num_examples == 1 + + +def test_profile_keeps_a_mixed_format_directory_as_one_partition(tmp_path): + # A stray .jsonl beside .parquet shards is noise, not a second dataset. Splitting the partition + # to keep a scalar `file_format` true invented structure that is not in the data *and* renamed + # the real partition (default -> default:parquet). Format is a per-file fact instead. + _write_parquet(tmp_path / "data" / "train-00000-of-00001.parquet", [{"prompt": "a"}]) + (tmp_path / "data" / "extra.jsonl").write_text('{"question": "b"}\n{"question": "c"}\n') + + result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) + + assert len(result.partitions) == 1 + partition = result.partitions[0] + assert partition.name == "data" + assert partition.file_formats == ["jsonl", "parquet"] + # Both formats' columns reach features. Trusting the declared parquet schema would have erased + # `question`, which only the schemaless file witnesses -- the defect the split worked around. + assert sorted(f.name for f in partition.features) == ["prompt", "question"] + assert result.sampling.rows_scanned == 3 # 1 parquet + 2 jsonl, each counted once + assert partition.rows_complete is True + + +def test_root_files_and_a_directory_named_default_stay_distinct(): + # The collision the empty-string sentinel exists to prevent: a derived label collapsed both to + # "default", leaving two partitions with one name and no way to reference either. + parts = group_partitions(_entries("root.parquet", "default/inner.parquet")) + assert [name for name, _ in parts] == ["", "default"] + + +def test_an_unrelated_file_does_not_rename_a_partition(tmp_path): + # Dropping a stray .jsonl into main/ used to turn partition "main" into "main:parquet" -- not + # renamed, *gone*, so a stored reference resolved to nothing. + _write_parquet(tmp_path / "main" / "train.parquet", [{"q": "a"}]) + _write_parquet(tmp_path / "socratic" / "train.parquet", [{"q": "b"}]) + before = [p.name for p in profile(LocalFileSource(tmp_path), created_at=FIXED_TIME).partitions] + + (tmp_path / "main" / "notes.jsonl").write_text('{"note": "someone dropped this here"}\n') + after = [p.name for p in profile(LocalFileSource(tmp_path), created_at=FIXED_TIME).partitions] + + assert before == after == ["main", "socratic"] + + +def test_profile_unions_columns_across_shards(tmp_path): + # A column that appears only in a later shard must still reach features/stats. Taking the first + # shard's schema would drop it entirely. + _write_parquet(tmp_path / "train-00000-of-00002.parquet", [{"prompt": "a", "completion": "b"}]) + _write_parquet(tmp_path / "train-00001-of-00002.parquet", [{"prompt": "c", "completion": "d", "score": 3}]) + + partition = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME).partitions[0] + + assert [f.name for f in partition.features] == ["prompt", "completion", "score"] + assert partition.stats["score"].null_rate == 0.5 # absent in the first shard, and said so + + +def test_profile_is_invariant_to_shard_order(tmp_path, tmp_path_factory): + # The same rows must profile the same way regardless of which shard sorts first. First-wins + # schema selection made this data classify as prompt_completion or scored_response depending + # purely on filename order. + narrow = [{"prompt": "a", "completion": "b"}] + wide = [{"prompt": "c", "completion": "d", "score": 3}] + + forward = tmp_path_factory.mktemp("forward") + _write_parquet(forward / "train-00000-of-00002.parquet", narrow) + _write_parquet(forward / "train-00001-of-00002.parquet", wide) + + reverse = tmp_path_factory.mktemp("reverse") + _write_parquet(reverse / "train-00000-of-00002.parquet", wide) + _write_parquet(reverse / "train-00001-of-00002.parquet", narrow) + + first = profile(LocalFileSource(forward), created_at=FIXED_TIME).partitions[0] + second = profile(LocalFileSource(reverse), created_at=FIXED_TIME).partitions[0] + + assert [(f.name, f.dtype) for f in first.features] == [(f.name, f.dtype) for f in second.features] + assert first.classification.dataset_type == second.classification.dataset_type == "scored_response" + + +def test_profile_survives_conflicting_shard_schemas(tmp_path): + # Two shards disagreeing on a column's type has no right answer at the schema level; fall back to + # inferring from the rows (which widens to json) rather than asserting one shard over the other. + _write_parquet(tmp_path / "train-00000-of-00002.parquet", [{"score": 1}]) + _write_parquet(tmp_path / "train-00001-of-00002.parquet", [{"score": "high"}]) + + partition = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME).partitions[0] + + assert [f.name for f in partition.features] == ["score"] + assert partition.features[0].dtype == "json" # mixed, and honest about it + + +def test_profile_isolates_unreadable_files(tmp_path): + _write_parquet(tmp_path / "train-00000-of-00001.parquet", [{"a": 1}]) + (tmp_path / "test-00000-of-00001.parquet").write_bytes(b"not a real parquet file") + + result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) + + splits = {s.name: s for s in result.partitions[0].splits} + assert splits["train"].num_examples == 1 + assert splits["test"].num_examples is None # unreadable -> count unknown, not a crash + assert [e.path for e in result.file_errors] == ["test-00000-of-00001.parquet"] # named, with a reason + assert result.file_errors[0].error + assert result.partitions[0].rows_complete is False # a file could not be fully parsed + assert result.sampling.rows_present is None + assert result.sampling.files_read == 1 # one file was actually read; the other never opened + assert result.sampling.files_present == 2 # ...out of two that were there to read + + +def test_profile_row_budget_bounds_reads_and_says_so(tmp_path): + _write_parquet(tmp_path / "train-00000-of-00001.parquet", [{"a": i} for i in range(10)]) + + result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME, row_budget=4) + + assert result.sampling.rows_scanned == 4 + assert result.partitions[0].rows_complete is False # 4 of 10 rows is not a full scan + # The footer knows the total even though the cap stopped the read. Gating this on completeness + # nulled it exactly when it carried information: "4 of 10" is a ratio, "4 of unknown" is not. + assert result.sampling.rows_present == 10 + assert result.partitions[0].splits[0].num_examples == 10 # the footer count survives sampling + + +def test_profile_uncapped_read_is_a_full_scan(tmp_path): + _write_parquet(tmp_path / "train-00000-of-00001.parquet", [{"a": i} for i in range(10)]) + + result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME, row_budget=None) + + assert result.partitions[0].rows_complete is True + assert result.sampling.rows_scanned == result.sampling.rows_present == 10 + + +def test_profile_cap_larger_than_a_jsonl_file_keeps_it_exhaustive(tmp_path): + # jsonl has no footer, so a cap could easily cost the exact count on files that never hit it. + # Reading to EOF under the cap must stay exact, or capping would degrade every small dataset. + (tmp_path / "train.jsonl").write_text('{"a": 1}\n{"a": 2}\n') + + result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME, row_budget=1000) + + assert result.partitions[0].splits[0].num_examples == 2 + assert result.partitions[0].rows_complete is True + assert result.sampling.rows_present == 2 + + +def test_profile_reports_unsupported_data_files(tmp_path): + # A directory of formats we cannot read must not profile as an exhaustively scanned empty + # dataset — that is indistinguishable from a dataset that really is empty. + (tmp_path / "train.csv").write_text("prompt,completion\na,b\n") + (tmp_path / "test.arrow").write_bytes(b"\x00") + + result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) + + assert result.partitions == [] + assert result.sampling.rows_present is None # not 0: "empty" would be a lie + assert result.sampling.files_read == 0 + assert result.sampling.files_present == 2 # both are data; neither could be read + # ...and they still weigh what they weigh. This is the case `bytes_present` exists for: no + # partition grouped these files, so summing the splits reports zero -- the same lie as "empty". + on_disk = (tmp_path / "train.csv").stat().st_size + (tmp_path / "test.arrow").stat().st_size + assert result.sampling.bytes_present == on_disk + assert sum(s.size_bytes for p in result.partitions for s in p.splits) == 0 + # Typed records now, each saying why -- not bare paths tucked into a free-form dict. + assert [e.path for e in result.file_errors] == ["test.arrow", "train.csv"] + assert all("no reader" in e.error for e in result.file_errors) + + +def test_profile_ignores_non_data_files_without_penalty(tmp_path): + # A README is genuinely not data, so it must not cost exhaustiveness the way a .csv does. + _write_parquet(tmp_path / "train-00000-of-00001.parquet", [{"a": 1}]) + (tmp_path / "README.md").write_text("a dataset card") + (tmp_path / "LICENSE").write_text("apache") + + result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) + + assert result.partitions[0].rows_complete is True + assert result.file_errors == [] + assert result.sampling.files_present == 1 # the README and LICENSE are not data, counted nowhere + # Nor does their weight land on the dataset: a card is not part of what has to be moved. + assert result.sampling.bytes_present == (tmp_path / "train-00000-of-00001.parquet").stat().st_size + + +def test_split_size_survives_a_shard_it_could_not_read(tmp_path): + # Size comes from the listing and a row count from reading, so they go unknown independently. + # A shard that fails to parse still weighs what it weighs, where the split's row count cannot. + _write_parquet(tmp_path / "train-00000-of-00002.parquet", [{"a": 1}, {"a": 2}]) + (tmp_path / "train-00001-of-00002.parquet").write_bytes(b"not parquet") + + result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) + + on_disk = sum(p.stat().st_size for p in tmp_path.glob("*.parquet")) + split = result.partitions[0].splits[0] + assert split.num_files == 2 + assert split.size_bytes == on_disk # both shards, including the one that would not open + assert split.num_examples is None # the broken shard's rows are unknowable... + assert split.size_bytes > 0 # ...but its bytes are not + assert result.sampling.bytes_present == on_disk + + +def test_profile_records_a_partial_jsonl_read(tmp_path): + # One corrupt line costs that line, not the file — but the profile must still say the file was + # only partly understood, rather than presenting a clean-looking count. + (tmp_path / "train.jsonl").write_text('{"a": 1}\n{"a": 2\n{"a": 3}\n') + + result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) + + assert result.partitions[0].splits[0].num_examples == 2 # the readable rows survived + assert [e.path for e in result.file_errors] == ["train.jsonl"] + assert "line 2" in result.file_errors[0].error + assert result.partitions[0].rows_complete is False # a line was lost, so not a full scan + + +def test_profile_classifies_roles_type_and_verifiability(tmp_path): + _write_parquet( + tmp_path / "train-00000-of-00001.parquet", + [{"problem": "q1", "solution": "steps #### 5"}, {"problem": "q2", "solution": "steps #### 6"}], + ) + partition = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME).partitions[0] + + assert {f.semantic_role for f in partition.features} == {"prompt", "completion"} + assert partition.classification.dataset_type == "prompt_completion" + assert partition.classification.verifiability.method == "extractable_final_answer" + assert partition.classification.verifiability.coverage == 1.0 + + +def test_profile_sharegpt_dataset_is_a_chat_dataset(tmp_path): + # End to end: {from, value} must reach the messages dtype, carry stats, and classify as chat + # rather than falling through to `unknown` with nothing measured. + conversation = [{"from": "human", "value": "hi"}, {"from": "gpt", "value": "hello"}] + (tmp_path / "train.jsonl").write_text(json.dumps({"conversations": conversation}) + "\n") + + partition = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME).partitions[0] + + assert partition.features[0].dtype == "messages" + assert partition.features[0].semantic_role == "messages" + assert partition.classification.dataset_type == "messages" + assert partition.stats["conversations"].messages.roles_seen == ["human", "gpt"] + + +def test_profile_degrades_one_partition_when_measurement_fails(tmp_path, monkeypatch): + # Reads are isolated per file, but schema/stats/classification ran unguarded, so one odd value + # could abort an otherwise complete profile. Structure must survive a measurement failure. + from nemo_datasets_plugin.profiler import pipeline as pipeline_module + + _write_parquet(tmp_path / "train-00000-of-00001.parquet", [{"a": 1}, {"a": 2}]) + monkeypatch.setattr(pipeline_module, "classify", lambda *a, **k: (_ for _ in ()).throw(RuntimeError("boom"))) + + result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) # must not raise + + partition = result.partitions[0] + assert partition.splits[0].num_examples == 2 # structure survives + assert partition.stats == {} + assert partition.classification.dataset_type == "unknown" + assert [e.kind for e in partition.classification.evidence] == ["error"] + assert "RuntimeError" in partition.classification.evidence[0].detail # says what failed + + +def test_a_read_failure_does_not_look_like_a_measurement_failure(tmp_path): + # The two failure domains have to stay distinguishable: a bad *file* is a FileError, and the + # rows that were readable still measure and classify normally. Folding the read and measure + # loops together is what would blur this, so it is pinned before that happens. + _write_parquet(tmp_path / "train-00000-of-00002.parquet", [{"prompt": "q", "completion": "a"}]) + (tmp_path / "train-00001-of-00002.parquet").write_bytes(b"not parquet") + + result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) + + part = result.partitions[0] + assert [e.path for e in result.file_errors] == ["train-00001-of-00002.parquet"] + assert part.classification.dataset_type == "prompt_completion" # the readable rows still classify + assert "error" not in {e.kind for e in part.classification.evidence} + assert part.stats # ...and are still measured + + +def test_a_measurement_failure_does_not_look_like_a_read_failure(tmp_path, monkeypatch): + from nemo_datasets_plugin.profiler import pipeline as pipeline_module + + _write_parquet(tmp_path / "train-00000-of-00001.parquet", [{"a": 1}]) + monkeypatch.setattr(pipeline_module, "classify", lambda *a, **k: (_ for _ in ()).throw(RuntimeError("boom"))) + + result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) + + assert result.file_errors == [] # the file was fine; the data was odd + assert [e.kind for e in result.partitions[0].classification.evidence] == ["error"] + # `rows_complete` speaks to rows read, and every row *was* read -- so it stays True even though + # there are no stats. Pinned as it stands; the field means what it says once Phase 5 renames it. + assert result.partitions[0].rows_complete is True + + +def test_one_unmeasurable_column_does_not_cost_the_partition_its_classification(tmp_path, monkeypatch): + # The narrow guard, end to end. The column's failure reaches the profile as evidence, and + # everything the partition could still establish -- the other column's stats, the roles, the + # dataset type -- survives it. + from nemo_datasets_plugin.profiler import stats as stats_module + + real_accumulator_for = stats_module._accumulator_for + + class Boom(stats_module.ColumnAccumulator): + def _observe(self, present): + raise RuntimeError("boom") + + monkeypatch.setattr( + stats_module, + "_accumulator_for", + lambda feature, expected_rows=None: ( + Boom() if feature.name == "completion" else real_accumulator_for(feature, expected_rows) + ), + ) + _write_parquet(tmp_path / "train.parquet", [{"prompt": "q", "completion": "a"}]) + + part = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME).partitions[0] + + assert "prompt" in part.stats and "completion" not in part.stats + assert part.classification.dataset_type == "prompt_completion" # roles come from names, not stats + assert any(e.kind == "error" and "'completion'" in e.detail for e in part.classification.evidence) + # The reasoning for the classification still reads first; the failure is a caveat on it. + assert part.classification.evidence[0].kind != "error" + + +def test_a_measurement_failure_is_scoped_to_its_own_partition(tmp_path, monkeypatch): + from nemo_datasets_plugin.profiler import pipeline as pipeline_module + + real_classify = pipeline_module.classify + + def poison_one_partition(features, stats, **kwargs): + if any(feature.name == "poison" for feature in features): + raise RuntimeError("boom") + return real_classify(features, stats, **kwargs) + + _write_parquet(tmp_path / "good" / "train.parquet", [{"prompt": "q", "completion": "a"}]) + _write_parquet(tmp_path / "bad" / "train.parquet", [{"poison": 1}]) + monkeypatch.setattr(pipeline_module, "classify", poison_one_partition) + + partitions = {p.name: p for p in profile(LocalFileSource(tmp_path), created_at=FIXED_TIME).partitions} + + assert partitions["bad"].classification.dataset_type == "unknown" + assert partitions["good"].classification.dataset_type == "prompt_completion" + assert partitions["good"].stats # a neighbour's bad data costs this partition nothing + + +def test_a_file_that_fails_partway_still_counts_what_it_contributed(tmp_path, monkeypatch): + # A read used to be all-or-nothing, so a failure meant no rows at all and the envelope could be + # written after it. A fold cannot give rows back: batches already folded are in the statistics + # whatever happens next, and counting the file as unread left `rows_scanned` describing fewer + # rows than the stats were built from. + from nemo_datasets_plugin.profiler import pipeline as pipeline_module + + _write_parquet(tmp_path / "train.parquet", [{"a": i} for i in range(4000)]) + real_update = pipeline_module._PartitionFolds.update + calls = {"n": 0} + + def fail_on_the_third_batch(self, rows): + calls["n"] += 1 + if calls["n"] == 3: + raise RuntimeError("boom mid-file") + return real_update(self, rows) + + monkeypatch.setattr(pipeline_module._PartitionFolds, "update", fail_on_the_third_batch) + result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) + + assert result.sampling.rows_scanned == 2048 # two batches of 1024 were folded before it failed + assert result.sampling.files_read == 1 # the file *was* read from, just not to its end + assert result.partitions[0].stats["a"].numeric is not None # and those rows shaped the stats + assert [e.path for e in result.file_errors] == ["train.parquet"] + assert result.partitions[0].rows_complete is False + + +def test_reading_everything_is_the_default(tmp_path): + # The point of the whole exercise. The budget existed to keep a materialised partition off the + # heap; nothing is materialised, so the default should not answer the question worse than it can + # be answered. + _write_parquet(tmp_path / "train.parquet", [{"t": f"row {i}"} for i in range(5_000)]) + + result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) + + assert result.sampling.rows_scanned == 5_000 == result.sampling.rows_present + assert result.partitions[0].rows_complete is True + + +def test_rows_complete_speaks_to_rows_read_not_to_exactness(tmp_path): + # It was `stats_complete`, which promised more than it delivered: quantiles and quality ratios + # are estimates by construction, whatever it says. Renamed to what it actually measures. + _write_parquet(tmp_path / "train.parquet", [{"t": f"row {i}"} for i in range(100)]) + + short = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME, row_budget=10) + whole = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) + + assert short.partitions[0].rows_complete is False # ten of a hundred rows + assert whole.partitions[0].rows_complete is True + # True either way, and it is the measurements themselves that say whether they are exact. + assert whole.partitions[0].stats["t"].text.chars.max == max(len(f"row {i}") for i in range(100)) + + +def test_profile_is_deterministic(tmp_path): + _write_parquet(tmp_path / "train-00000-of-00001.parquet", [{"a": 1}, {"a": 2}]) + source = LocalFileSource(tmp_path) + first = profile(source, created_at=FIXED_TIME) + second = profile(source, created_at=FIXED_TIME) + assert first.model_dump_json() == second.model_dump_json() + + +def test_profile_tolerates_non_object_jsonl_lines(tmp_path): + # A valid-JSON-but-non-object line parses cleanly, so the reader (not the read) must handle it; + # otherwise it would poison the unprotected schema/stats stage and abort the whole profile. + (tmp_path / "train.jsonl").write_text('{"a": 1}\n[1, 2, 3]\n{"a": 2}\n') + result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) + + assert result.partitions[0].splits[0].num_examples == 2 # objects counted, stray array dropped + assert result.partitions[0].rows_complete is True + + +def test_profile_survives_a_hostile_directory(tmp_path): + """Everything that can go wrong at once must still yield a profile that says what went wrong. + + Each of these individually used to either abort the run or vanish silently; this is the shape of + bug that got through, so it is worth asserting as one scenario rather than only in isolation. + """ + _write_parquet(tmp_path / "train-00000-of-00002.parquet", [{"prompt": "a", "completion": "b"}]) + (tmp_path / "train-00001-of-00002.parquet").write_bytes(b"not a parquet file") # corrupt + (tmp_path / "extra.jsonl").write_text( + '{"messages": [{"role": 1, "content": "hi"}]}\n' # non-string role + '{"messages": [{"role": "user"\n' # truncated line + "[1, 2, 3]\n" # valid JSON, not a row + '{"messages": [{"role": "user", "content": "ok"}]}\n' + ) + (tmp_path / "leftovers.csv").write_text("a,b\n1,2\n") # recognizable data, no reader + (tmp_path / "README.md").write_text("a dataset card") # not data at all + + result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) # must not raise + + # Nothing here is exhaustive, and the profile says so rather than looking clean. + assert result.partitions[0].rows_complete is False + assert result.sampling.rows_present is None + # One channel for every file the profiler could not use, whether or not a partition grouped it: + # the .csv it never read, the corrupt shard, and the jsonl it only partly parsed. + assert [e.path for e in result.file_errors] == [ + "extra.jsonl", + "leftovers.csv", + "train-00001-of-00002.parquet", + ] + + # One partition, not one per format: the stray .jsonl is noise, not a second dataset. + assert len(result.partitions) == 1 + partition = result.partitions[0] + assert partition.file_formats == ["jsonl", "parquet"] + # The readable parquet rows still produced a real classification... + assert partition.classification.dataset_type == "prompt_completion" + # ...and the odd jsonl rows were measured rather than aborting the run. + assert partition.stats["messages"].messages.roles_seen == ["1", "user"] + + # The whole thing still round-trips as a stored profile. + assert DatasetProfile.model_validate_json(result.model_dump_json()) == result + + +def test_split_file_counts_account_for_every_data_file(tmp_path): + # The contract promises split membership is exhaustive and disjoint. With per-file records gone + # the counts are all that carries it, so the invariant is worth asserting on them directly -- + # a count that silently dropped a file would look exactly like a smaller dataset. + _write_parquet(tmp_path / "train-00000-of-00002.parquet", [{"a": 1}]) + _write_parquet(tmp_path / "train-00001-of-00002.parquet", [{"a": 2}]) + _write_parquet(tmp_path / "test-00000-of-00001.parquet", [{"a": 3}]) + (tmp_path / "README.md").write_text("a dataset card") # not data; never becomes a FileRecord + + source = LocalFileSource(tmp_path) + result = profile(source, created_at=FIXED_TIME) + + counted = sum(split.num_files for partition in result.partitions for split in partition.splits) + listed = [e.path for e in source.list_files() if e.path.endswith(".parquet")] + assert counted == len(listed) # exhaustive and disjoint: each file lands in exactly one split + + +def test_profile_isolates_detected_format_with_no_reader(tmp_path, monkeypatch): + # If detect_format recognizes an extension the registry has no reader for, that file must be + # isolated like a corrupt one, not crash the whole profile. + from nemo_datasets_plugin.profiler.readers import base + + monkeypatch.setitem(base._EXTENSION_FORMATS, ".xyz", "xyz-no-reader") + _write_parquet(tmp_path / "train-00000-of-00001.parquet", [{"a": 1}]) + (tmp_path / "extra.xyz").write_text("whatever") + + result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) # must not raise + + assert "extra.xyz" in {e.path for e in result.file_errors} # named, not silently dropped + assert result.partitions[0].rows_complete is False + + +def test_a_column_only_a_schemaless_file_witnessed_survives(tmp_path): + # A group where only some files declare a schema used to trust that schema and erase every column + # the schemaless files were the sole witness for. Now the partition infers its schema from the + # rows as it folds them, so the sole witness is heard. + _write_parquet(tmp_path / "train-00000-of-00002.parquet", [{"prompt": "a"}]) + (tmp_path / "train-00001-of-00002.jsonl").write_text(json.dumps({"prompt": "b", "extra": "only here"})) + + part = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME).partitions[0] + + assert [f.name for f in part.features] == ["prompt", "extra"] + assert set(part.stats) <= {f.name for f in part.features} + + +def test_a_declared_schema_is_trusted_when_it_covers_every_file(tmp_path): + # The other half: when every file declares one, the schema is authoritative and the rows are not + # consulted for it. That is what lets the partition fold with its accumulators chosen up front. + _write_parquet(tmp_path / "train-00000-of-00002.parquet", [{"prompt": "a"}]) + _write_parquet(tmp_path / "train-00001-of-00002.parquet", [{"prompt": "b"}]) + + part = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME).partitions[0] + + assert [f.name for f in part.features] == ["prompt"] + + +def test_rows_completeness_is_per_partition(tmp_path): + # A corrupt shard in one partition says nothing about the measurements in another, but a + # fileset-wide flag downgraded every partition to the worst one. It was never even the value + # that gated quoting a proven enumeration -- that was decided per partition and never stored. + rows = [{"label": t} for t in (True, False, True)] + _write_parquet(tmp_path / "main" / "train.parquet", rows) + _write_parquet(tmp_path / "socratic" / "train.parquet", rows) + (tmp_path / "socratic" / "broken.parquet").write_bytes(b"not a parquet file") + + partitions = {p.name: p for p in profile(LocalFileSource(tmp_path), created_at=FIXED_TIME).partitions} + + assert partitions["main"].rows_complete is True + assert partitions["socratic"].rows_complete is False + # Quoting is decided by role, not by completeness, so both keep their label vocabulary -- + # rows_complete is what tells a consumer whether socratic's list is the whole of it. + assert partitions["main"].stats["label"].categorical.values == ["False", "True"] + assert partitions["socratic"].stats["label"].categorical.values == ["False", "True"] + + +def test_dataset_wide_completeness_is_one_expression(tmp_path): + # SamplingInfo no longer carries `exhaustive`; the contract documents this derivation in its + # place. It has to keep working, or dropping the flag cost consumers something -- and it now + # says *which* half failed, which the single bit could not. + _write_parquet(tmp_path / "train.parquet", [{"a": 1}]) + clean = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) + assert all(p.rows_complete for p in clean.partitions) and not clean.file_errors + + (tmp_path / "extra.csv").write_text("a,b\n1,2\n") + with_csv = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) + assert all(p.rows_complete for p in with_csv.partitions) # the parquet rows are still complete + assert with_csv.file_errors # but there is data here that went unprofiled + assert with_csv.sampling.files_read == 1 and with_csv.sampling.files_present == 2 + + +def test_row_budget_is_divided_across_a_partitions_files(tmp_path): + for shard in range(4): + _write_parquet(tmp_path / f"train-{shard:05d}-of-00004.parquet", [{"a": i} for i in range(200)]) + + result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME, row_budget=400) + + assert result.sampling.rows_scanned == 400 # 400 / 4 files = 100 rows each + + +def test_rows_read_do_not_grow_when_a_dataset_is_resharded(tmp_path_factory): + # The property the budget exists for. Under a per-file cap the same data split ten ways further + # cost ten times the peak memory while describing exactly the same rows. + def rows_read(shards, per_shard): + root = tmp_path_factory.mktemp(f"shards{shards}") + for shard in range(shards): + _write_parquet(root / f"train-{shard:05d}-of-{shards:05d}.parquet", [{"a": i} for i in range(per_shard)]) + return profile(LocalFileSource(root), created_at=FIXED_TIME, row_budget=400).sampling.rows_scanned + + assert rows_read(4, 200) == rows_read(40, 20) == 400 + + +def test_row_budget_keeps_a_floor_under_very_thin_shards(tmp_path): + # Below the floor a file cannot witness the columns only it holds, which is the reason every file + # is opened rather than a subset sampled. Overshooting the budget there is the right trade, and + # the arithmetic share would be 1, so the floor holds and the budget is deliberately exceeded. + for shard in range(10): + _write_parquet(tmp_path / f"train-{shard:05d}-of-00010.parquet", [{"a": i} for i in range(50)]) + + result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME, row_budget=10) + + assert result.sampling.rows_scanned == 100 # 10 files x the 10-row floor, over the budget of 10 diff --git a/plugins/nemo-datasets/tests/test_profile_task.py b/plugins/nemo-datasets/tests/test_profile_task.py new file mode 100644 index 0000000000..43fd69078e --- /dev/null +++ b/plugins/nemo-datasets/tests/test_profile_task.py @@ -0,0 +1,145 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for the dataset-profiler job task.""" + +import json +from pathlib import Path +from typing import cast + +import nemo_datasets_plugin.tasks.profile.run as run_mod +import pyarrow as pa +import pyarrow.parquet as pq +from nemo_platform import NeMoPlatform +from nemo_platform_plugin.job_results import ResultRef +from nemo_platform_plugin.jobs.constants import ( + NEMO_JOB_ID_ENVVAR, + NEMO_JOB_STEP_CONFIG_FILE_PATH_ENVVAR, + NEMO_JOB_WORKSPACE_ENVVAR, +) + +# The task touches the sdk only through PlatformJobResults, which is patched below, so a bare object +# stands in for it. +_SDK = cast(NeMoPlatform, object()) + + +def _dataset(root: Path, rows=None) -> Path: + root.mkdir(parents=True, exist_ok=True) + rows = rows or [{"q": "why?", "a": "because #### 4"}] + pq.write_table(pa.Table.from_pylist(rows), root / "train-00000-of-00001.parquet") + return root + + +def _install(monkeypatch, tmp_path: Path, config: dict) -> dict: + """Point the task at a step config and capture what it publishes.""" + published: dict = {} + + class _Results: + def __init__(self, *, job_name, workspace, sdk): + published.update(job_name=job_name, workspace=workspace) + + def save(self, name, local_path): + published["name"] = name + published["profile"] = json.loads((Path(local_path) / "profile.json").read_text()) + return ResultRef(name=name, artifact_url=f"file://{local_path}") + + config_path = tmp_path / "step-config.json" + config_path.write_text(json.dumps(config)) + monkeypatch.setenv(NEMO_JOB_STEP_CONFIG_FILE_PATH_ENVVAR, str(config_path)) + monkeypatch.setenv(NEMO_JOB_WORKSPACE_ENVVAR, "ws1") + monkeypatch.setenv(NEMO_JOB_ID_ENVVAR, "job-1") + monkeypatch.setattr(run_mod, "PlatformJobResults", _Results) + return published + + +def test_task_profiles_a_directory_and_publishes_the_profile(tmp_path, monkeypatch): + data = _dataset(tmp_path / "data") + published = _install(monkeypatch, tmp_path, {"path": str(data)}) + + assert run_mod.run(_SDK) == 0 + + assert published["job_name"] == "job-1" + assert published["workspace"] == "ws1" + assert published["name"] == "profile" + profile = published["profile"] + assert profile["partitions"][0]["file_formats"] == ["parquet"] + assert profile["sampling"]["files_read"] == 1 + + +def test_task_prefers_the_step_configs_workspace_over_the_environment(tmp_path, monkeypatch): + data = _dataset(tmp_path / "data") + published = _install(monkeypatch, tmp_path, {"path": str(data), "workspace": "explicit"}) + + assert run_mod.run(_SDK) == 0 + assert published["workspace"] == "explicit" + + +def test_task_passes_column_role_hints_to_the_profiler(tmp_path, monkeypatch): + # The step config is the profiler's hint channel now that there is no CLI to carry --column-role. + data = _dataset(tmp_path / "data") + published = _install(monkeypatch, tmp_path, {"path": str(data), "column_roles": {"q": "prompt", "a": "completion"}}) + + assert run_mod.run(_SDK) == 0 + classification = published["profile"]["partitions"][0]["classification"] + assert classification["dataset_type"] == "prompt_completion" + + +def test_task_reads_everything_by_default(tmp_path, monkeypatch): + data = _dataset(tmp_path / "data") + published = _install(monkeypatch, tmp_path, {"path": str(data)}) + + assert run_mod.run(_SDK) == 0 + sampling = published["profile"]["sampling"] + assert sampling["rows_scanned"] == sampling["rows_present"] # nothing left unread + + +def test_task_honours_an_explicit_row_budget(tmp_path, monkeypatch): + data = _dataset(tmp_path / "data", rows=[{"a": i} for i in range(20)]) + published = _install(monkeypatch, tmp_path, {"path": str(data), "row_budget": 5}) + + assert run_mod.run(_SDK) == 0 + assert published["profile"]["sampling"]["rows_scanned"] == 5 + assert published["profile"]["sampling"]["rows_scanned"] == 5 + + +def test_row_budget_zero_asks_for_every_row(tmp_path, monkeypatch): + data = _dataset(tmp_path / "data", rows=[{"a": i} for i in range(20)]) + published = _install(monkeypatch, tmp_path, {"path": str(data), "row_budget": 0}) + + assert run_mod.run(_SDK) == 0 + sampling = published["profile"]["sampling"] + assert sampling["rows_scanned"] == sampling["rows_present"] # 0 means "all of them" + assert published["profile"]["partitions"][0]["rows_complete"] is True + + +def test_task_fails_when_the_step_config_says_nothing_to_profile(tmp_path, monkeypatch): + published = _install(monkeypatch, tmp_path, {}) + assert run_mod.run(_SDK) == 1 # a nonzero exit, not a traceback out of the container + assert published == {} + + +def test_task_fails_when_the_path_is_not_a_directory(tmp_path, monkeypatch): + target = tmp_path / "a-file" + target.write_text("x") + published = _install(monkeypatch, tmp_path, {"path": str(target)}) + + assert run_mod.run(_SDK) == 1 + assert published == {} + + +def test_task_fails_without_a_step_config(tmp_path, monkeypatch): + _install(monkeypatch, tmp_path, {"path": str(_dataset(tmp_path / "data"))}) + monkeypatch.delenv(NEMO_JOB_STEP_CONFIG_FILE_PATH_ENVVAR) + assert run_mod.run(_SDK) == 1 + + +def test_task_rejects_a_negative_row_budget(tmp_path, monkeypatch): + data = _dataset(tmp_path / "data") + _install(monkeypatch, tmp_path, {"path": str(data), "row_budget": -1}) + assert run_mod.run(_SDK) == 1 + + +def test_task_rejects_column_roles_that_are_not_a_mapping(tmp_path, monkeypatch): + data = _dataset(tmp_path / "data") + _install(monkeypatch, tmp_path, {"path": str(data), "column_roles": ["q=prompt"]}) + assert run_mod.run(_SDK) == 1 diff --git a/plugins/nemo-datasets/tests/test_readers.py b/plugins/nemo-datasets/tests/test_readers.py new file mode 100644 index 0000000000..391b11aa5a --- /dev/null +++ b/plugins/nemo-datasets/tests/test_readers.py @@ -0,0 +1,147 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the file-source seam and the per-format readers.""" + +import pyarrow as pa +import pyarrow.parquet as pq +import pytest +from nemo_datasets_plugin.profiler.file_source import FileEntry, LocalFileSource +from nemo_datasets_plugin.profiler.readers.base import detect_format, get_reader + +PARQUET_ROWS = [ + {"prompt": "a", "score": 1}, + {"prompt": "b", "score": 2}, + {"prompt": "c", "score": 3}, +] + + +def _write_parquet(path, rows): + pq.write_table(pa.Table.from_pylist(rows), path) + + +# --- file source --------------------------------------------------------------------------------- + + +def test_local_file_source_lists_sorted_with_sizes(tmp_path): + (tmp_path / "b.jsonl").write_text('{"x": 1}\n') + (tmp_path / "a.parquet").write_bytes(b"not-real-parquet") # only listed here, not parsed + sub = tmp_path / "sub" + sub.mkdir() + (sub / "c.jsonl").write_text("{}\n") + + entries = LocalFileSource(tmp_path).list_files() + + assert [e.path for e in entries] == ["a.parquet", "b.jsonl", "sub/c.jsonl"] + assert all(e.size_bytes > 0 for e in entries) + assert all(e.checksum is None for e in entries) # local sources report no checksum + + +def test_local_file_source_open_reads_bytes(tmp_path): + (tmp_path / "f.jsonl").write_text('{"x": 1}\n') + with LocalFileSource(tmp_path).open("f.jsonl") as stream: + assert stream.read() == b'{"x": 1}\n' + + +def test_local_file_source_rejects_non_directory(tmp_path): + target = tmp_path / "f" + target.write_text("x") + with pytest.raises(NotADirectoryError): + LocalFileSource(target) + + +# --- registry ------------------------------------------------------------------------------------ + + +def test_detect_format_by_extension(): + assert detect_format("data/train-00000-of-00003.parquet") == "parquet" + assert detect_format("x.jsonl") == "jsonl" + assert detect_format("x.ndjson") == "jsonl" + assert detect_format("README.md") is None + + +def test_get_reader_unknown_format_raises(): + with pytest.raises(KeyError): + get_reader("arrow") + + +# --- parquet reader ------------------------------------------------------------------------------ + + +def test_parquet_reader_reads_schema_rows_and_exact_count(tmp_path): + _write_parquet(tmp_path / "d.parquet", PARQUET_ROWS) + result = get_reader("parquet").read(LocalFileSource(tmp_path), FileEntry("d.parquet", 0)) + + assert result.num_rows == 3 # exact, from the footer + assert result.rows_scanned == 3 + assert result.rows == PARQUET_ROWS + assert result.arrow_schema is not None + assert set(result.arrow_schema.names) == {"prompt", "score"} + + +def test_parquet_reader_row_cap_bounds_rows_but_keeps_exact_count(tmp_path): + _write_parquet(tmp_path / "d.parquet", PARQUET_ROWS) + result = get_reader("parquet").read(LocalFileSource(tmp_path), FileEntry("d.parquet", 0), row_cap=2) + + assert result.num_rows == 3 # footer count is unaffected by sampling + assert result.rows_scanned == 2 + assert result.rows == PARQUET_ROWS[:2] + + +def test_parquet_reader_zero_cap_reads_no_rows(tmp_path): + _write_parquet(tmp_path / "d.parquet", PARQUET_ROWS) + result = get_reader("parquet").read(LocalFileSource(tmp_path), FileEntry("d.parquet", 0), row_cap=0) + + assert result.rows == [] + assert result.num_rows == 3 + assert result.arrow_schema is not None # schema is still known without reading rows + + +# --- jsonl reader -------------------------------------------------------------------------------- + + +def test_jsonl_reader_full_read_is_exact_and_skips_blanks(tmp_path): + (tmp_path / "d.jsonl").write_text('{"a": 1}\n\n{"a": 2}\n') + result = get_reader("jsonl").read(LocalFileSource(tmp_path), FileEntry("d.jsonl", 0)) + + assert result.rows == [{"a": 1}, {"a": 2}] + assert result.num_rows == 2 # exact on a full read + assert result.arrow_schema is None # jsonl declares no schema + + +def test_jsonl_reader_row_cap_leaves_count_unknown(tmp_path): + (tmp_path / "d.jsonl").write_text('{"a": 1}\n{"a": 2}\n{"a": 3}\n') + result = get_reader("jsonl").read(LocalFileSource(tmp_path), FileEntry("d.jsonl", 0), row_cap=2) + + assert result.rows == [{"a": 1}, {"a": 2}] + assert result.rows_scanned == 2 + assert result.num_rows is None # a partial read can't assert the total + + +def test_jsonl_reader_skips_non_object_lines(tmp_path): + # A record is a column map; valid JSON that is a scalar or array is not a row. + (tmp_path / "d.jsonl").write_text('{"a": 1}\n[1, 2, 3]\n42\n"loose"\n{"a": 2}\n') + result = get_reader("jsonl").read(LocalFileSource(tmp_path), FileEntry("d.jsonl", 0)) + + assert result.rows == [{"a": 1}, {"a": 2}] # stray non-object lines dropped, objects kept + assert result.num_rows == 2 + # Not a read failure: those lines are not rows of this dataset, so the count stays exact and the + # file is still exhaustively scanned. Only an *unparseable* line is an error. + assert result.error is None + + +def test_jsonl_reader_survives_an_unparseable_line(tmp_path): + # One truncated line must cost that line, not the file. Dropping the whole file would erase its + # row count and any column it was the only witness for. + (tmp_path / "d.jsonl").write_text('{"a": 1}\n{"a": 2\n{"a": 3}\n') + result = get_reader("jsonl").read(LocalFileSource(tmp_path), FileEntry("d.jsonl", 0)) + + assert result.rows == [{"a": 1}, {"a": 3}] # the readable rows survive + assert result.rows_scanned == 2 + assert result.error is not None + assert "line 2" in result.error # self-describing: which line, and why + + +def test_jsonl_reader_clean_read_reports_no_error(tmp_path): + (tmp_path / "d.jsonl").write_text('{"a": 1}\n{"a": 2}\n') + assert get_reader("jsonl").read(LocalFileSource(tmp_path), FileEntry("d.jsonl", 0)).error is None diff --git a/plugins/nemo-datasets/tests/test_schema.py b/plugins/nemo-datasets/tests/test_schema.py new file mode 100644 index 0000000000..d27b84b91d --- /dev/null +++ b/plugins/nemo-datasets/tests/test_schema.py @@ -0,0 +1,100 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for row-schema derivation (from a declared arrow schema and from sampled rows).""" + +import pyarrow as pa +from nemo_datasets_plugin.profiler.schema import MAX_COLUMNS, columns_were_capped, derive_features + +# --- from a declared arrow schema (parquet) ------------------------------------------------------ + + +def test_from_arrow_scalars_keep_declared_widths(): + schema = pa.schema( + [("s", pa.string()), ("i", pa.int64()), ("i32", pa.int32()), ("b", pa.bool_()), ("f", pa.float64())] + ) + features = {f.name: f.dtype for f in derive_features([], schema)} + assert features == {"s": "string", "i": "int64", "i32": "int32", "b": "bool", "f": "float64"} + + +def test_from_arrow_list_of_role_content_structs_is_messages(): + schema = pa.schema([("prompt", pa.list_(pa.struct([("role", pa.string()), ("content", pa.string())])))]) + feature = derive_features([], schema)[0] + assert feature.dtype == "messages" + assert feature.items.dtype == "struct" + assert [f.name for f in feature.items.fields] == ["role", "content"] + + +def test_from_arrow_fixed_and_variable_lists_agree_on_shape(): + # A fixed-size list is still a list of its element type. The constant length itself is no longer + # recorded, so the two cases must be indistinguishable rather than one silently losing `items`. + fixed = derive_features([], pa.schema([("embedding", pa.list_(pa.float32(), 768))]))[0] + assert (fixed.dtype, fixed.items.dtype) == ("list", "float32") + + variable = derive_features([], pa.schema([("tags", pa.list_(pa.string()))]))[0] + assert (variable.dtype, variable.items.dtype) == ("list", "string") + + +# --- inferred from sampled rows (jsonl) ---------------------------------------------------------- + + +def test_from_rows_scalars_widen_int_and_float(): + rows = [{"a": 1, "b": 1.5, "c": "x", "d": True}, {"a": 2, "b": 2, "c": "y", "d": False}] + features = {f.name: f.dtype for f in derive_features(rows)} + assert features == {"a": "int64", "b": "float64", "c": "string", "d": "bool"} + + +def test_from_rows_list_of_role_content_structs_is_messages(): + rows = [{"conv": [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "yo"}]}] + feature = derive_features(rows)[0] + assert feature.dtype == "messages" + assert {f.name for f in feature.items.fields} == {"role", "content"} + + +def test_from_rows_sharegpt_from_value_is_messages(): + # ShareGPT spells the same structure {from, value}. Recognizing only {role, content} left it a + # plain list, which then failed the messages dtype gate and profiled as `unknown` with no stats. + rows = [{"conversations": [{"from": "human", "value": "hi"}, {"from": "gpt", "value": "yo"}]}] + feature = derive_features(rows)[0] + assert feature.dtype == "messages" + assert {f.name for f in feature.items.fields} == {"from", "value"} + + +def test_from_arrow_sharegpt_from_value_is_messages(): + schema = pa.schema([("conversations", pa.list_(pa.struct([("from", pa.string()), ("value", pa.string())])))]) + assert derive_features([], schema)[0].dtype == "messages" + + +def test_from_rows_lists_infer_their_element_type(): + constant = derive_features([{"e": [0.1, 0.2, 0.3]}, {"e": [0.4, 0.5, 0.6]}])[0] + assert (constant.dtype, constant.items.dtype) == ("list", "float64") + + variable = derive_features([{"e": [1, 2]}, {"e": [1, 2, 3]}])[0] + assert (variable.dtype, variable.items.dtype) == ("list", "int64") + + +def test_from_rows_nested_struct(): + feature = derive_features([{"meta": {"id": 1, "src": "a"}}, {"meta": {"id": 2, "src": "b"}}])[0] + assert feature.dtype == "struct" + assert {f.name for f in feature.fields} == {"id", "src"} + + +def test_from_rows_all_null_column_is_json(): + assert derive_features([{"x": None}, {"x": None}])[0].dtype == "json" + + +def test_derive_features_prefers_declared_arrow_schema(): + feature = derive_features([{"x": 1}], pa.schema([("x", pa.int32())]))[0] + assert feature.dtype == "int32" # declared width beats the int64 inference from rows + + +def test_column_count_is_bounded_and_says_when_it_stopped(): + # A malformed file whose rows carry unique keys would otherwise mint a column -- and later an + # accumulator -- for every row. The row budget used to bound this by accident; an unbounded read + # does not, so the bound is stated and the truncation is reported rather than silent. + rows = [{f"col{i}": i} for i in range(MAX_COLUMNS + 500)] + features = derive_features(rows) + assert len(features) == MAX_COLUMNS + assert columns_were_capped(features) + + assert not columns_were_capped(derive_features([{"a": 1, "b": 2}])) diff --git a/plugins/nemo-datasets/tests/test_stats.py b/plugins/nemo-datasets/tests/test_stats.py new file mode 100644 index 0000000000..fe83c8a0bf --- /dev/null +++ b/plugins/nemo-datasets/tests/test_stats.py @@ -0,0 +1,574 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for per-column statistics.""" + +import math + +import pytest +from nemo_datasets_plugin.profiler import stats as stats_module +from nemo_datasets_plugin.profiler.stats import ( + _MAX_VOCABULARY_BYTES, + _MAX_VOCABULARY_VALUE_CHARS, + _MAX_VOCABULARY_VALUES, + _NON_ASCII_RUN, + _WHITESPACE_RUN, + _non_ascii_count, + _whitespace_count, + measure_columns, + quote_enumerations, +) +from nemo_platform_plugin.files.dataset_profile import ColumnStats, FeatureSchema + + +def _stats(features, rows): + """Statistics only. Asserts nothing failed: these tests measure values, not the guard, and a + swallowed exception would surface here as a confusing KeyError instead of its own message.""" + measured = measure_columns(features, rows) + assert not measured.errors, measured.errors + return measured.stats + + +def _probes(features, rows): + """The content probes alone. `measure_columns` measures both in one pass; these tests want one.""" + return measure_columns(features, rows).probes + + +def _feature(name, dtype): + return FeatureSchema(name=name, dtype=dtype) + + +def _rows(name, values): + return [{name: value} for value in values] + + +# --- length histogram ---------------------------------------------------------------------------- + + +@pytest.mark.parametrize("value", [0, 1, 31, 32, 33, 63, 64, 255, 256, 1_000, 1_300, 65_535, 1_000_000, 33_554_432]) +def test_every_length_lands_in_a_bucket_that_contains_it(value): + # The bounds are what a quantile is read off, so they have to invert the bucketing exactly. A + # bucket whose range does not contain its own values would report a plausible wrong number. + low, high = stats_module._bucket_bounds(stats_module._length_bucket(value)) + assert low <= value < high + + +def test_short_lengths_are_recorded_exactly(): + # Below the slice count every length gets its own counter. That is what keeps the small fixtures + # in this file exact, and it is why a column of short strings loses nothing to bucketing. + hist = stats_module._LengthHistogram() + for n in range(stats_module._HISTOGRAM_SLICES): + hist.add(n) + quantiles = hist.quantiles() + assert (quantiles.p50, quantiles.p95, quantiles.p99, quantiles.max) == (15, 30, 31, 31) + + +def test_quantiles_stay_within_the_bound_on_a_heavy_tail(): + # The shape that matters: most rows short, a thin long tail. It is also the shape a mean cannot + # describe, which is why the distribution is carried at all. + values = [10] * 5000 + [200] * 3000 + [4000] * 1500 + [90_000] * 500 + hist = stats_module._LengthHistogram() + for value in values: + hist.add(value) + quantiles = hist.quantiles() + + ordered = sorted(values) + for percentile, got in ((50, quantiles.p50), (95, quantiles.p95), (99, quantiles.p99)): + want = ordered[min(math.ceil(percentile / 100 * len(ordered)), len(ordered)) - 1] + assert abs(got - want) / want <= 0.02, (percentile, want, got) + assert quantiles.max == 90_000 # exact, never rounded to a bucket + assert quantiles.p50 <= quantiles.p95 <= quantiles.p99 <= quantiles.max + + +def test_an_empty_histogram_reports_zeros(): + quantiles = stats_module._LengthHistogram().quantiles() + assert (quantiles.p50, quantiles.p95, quantiles.p99, quantiles.max) == (0, 0, 0, 0) + + +def test_roles_seen_stops_growing(): + # Fed straight from row content, so without a bound one malformed column could hold a string per + # message -- and membership is checked against the list, so it is quadratic as well as unbounded. + rows = _rows("m", [[{"role": f"role-{i}", "content": "x"}] for i in range(stats_module._MAX_ROLES_SEEN * 3)]) + measured = _stats([_feature("m", "messages")], rows)["m"] + assert len(measured.messages.roles_seen) == stats_module._MAX_ROLES_SEEN + # The rates still count every row: only the vocabulary of roles is bounded, not the measurement. + assert measured.messages.turns.max == 1 + + +# --- text ---------------------------------------------------------------------------------------- + + +def test_text_stats_length_quantiles_and_quality(): + values = ["a", "bb", "ccc", "dddd"] + stats = _stats([_feature("t", "string")], _rows("t", values))["t"] + assert stats.text.chars.max == 4 + assert stats.text.chars.p50 in {2, 3} # nearest-rank over 4 values + assert stats.quality is not None + assert stats.quality.whitespace_ratio == 0.0 + + +def test_text_quality_flags_repetition_and_non_ascii(): + stats = _stats([_feature("t", "string")], _rows("t", ["aaaaaaaa", "héllo wörld"]))["t"] + assert stats.quality.repetition_score > 0.0 # the "aaaaaaaa" run + assert stats.quality.non_ascii_ratio > 0.0 # accented characters + + +def test_one_bad_column_costs_only_itself(monkeypatch): + # The narrow guard. A value no detector anticipated used to cost the partition every measurement + # it had; it now costs its own column, and says so rather than leaving a silent gap. + real_accumulator_for = stats_module._accumulator_for + + class Boom(stats_module.ColumnAccumulator): + def _observe(self, present): + raise RuntimeError("boom") + + monkeypatch.setattr( + stats_module, + "_accumulator_for", + lambda feature, expected_rows=None: ( + Boom() if feature.name == "bad" else real_accumulator_for(feature, expected_rows) + ), + ) + + features = [_feature("good", "string"), _feature("bad", "string")] + result = measure_columns(features, [{"good": "x", "bad": "y"}]) + measured, probes, errors = result.stats, result.probes, result.errors + + assert "good" in measured and "bad" not in measured + assert "good" in probes and "bad" not in probes # probes go with the column that failed + assert [e.kind for e in errors] == ["error"] + assert "'bad'" in errors[0].detail and "RuntimeError" in errors[0].detail + + +# One column per dtype the dispatch knows, each carrying the awkward cases: nulls, empties, a +# non-finite float, a value long enough to saturate a vocabulary, both chat spellings. +_DTYPE_VALUES = { + "string": ["a prompt #### 42", "héllo wörld", "", "aaaaaaaa", None, "x" * 300, "yes", "yes"], + "int64": [1, 2, 2, None, 3, -5], + "float64": [1.5, float("nan"), 2.5, None, float("inf"), 0.0], + "bool": [True, False, True, None], + "messages": [ + [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "there"}], + [{"from": "human", "value": "q"}, {"from": "gpt", "value": "\\boxed{4}"}], + [{"role": "user", "content": "x", "tool_calls": [{"n": 1}]}], + None, + [], + ], + "struct": [{"a": 1}, None, {"b": 2}], +} + + +@pytest.mark.parametrize("dtype", sorted(_DTYPE_VALUES)) +@pytest.mark.parametrize("chunks", [1, 2, 3, 7]) +def test_an_accumulator_folds_rather_than_buffers(dtype, chunks): + # The property the fold rests on: a column split across calls has to measure the same as one + # handed over whole. Without it, batching would quietly change the numbers -- and the batch size + # is an implementation detail no reader of the profile could see. + values = _DTYPE_VALUES[dtype] * 5 + feature = _feature("c", dtype) + + whole = stats_module._accumulator_for(feature) + whole.update(values) + + in_pieces = stats_module._accumulator_for(feature) + step = max(1, -(-len(values) // chunks)) + for start in range(0, len(values), step): + in_pieces.update(values[start : start + step]) + + assert in_pieces.finalize() == whole.finalize() + + +def test_the_typed_accumulators_probe_exactly_as_the_bare_one_does(): + # Probe counting lives on the base class, so every dtype gets it for free. If a subclass ever + # shadows that, a chat column would quietly stop contributing verifiability evidence. + features = [_feature(dtype, dtype) for dtype in sorted(_DTYPE_VALUES)] + rows = [dict(zip(sorted(_DTYPE_VALUES), values)) for values in zip(*_DTYPE_VALUES.values())] + + probes, errors = measure_columns(features, rows).probes, measure_columns(features, rows).errors + + assert errors == [] + assert probes == _probes(features, rows) # probes come off the base class + + +def test_measure_columns_measures_every_dtype_the_dispatch_knows(): + features = [ + _feature("text", "string"), + _feature("count", "int64"), + _feature("score", "float64"), + _feature("flag", "bool"), + _feature("chat", "messages"), + _feature("meta", "struct"), + _feature("missing", "string"), + ] + rows = [ + { + "text": "a prompt ending in #### 42", + "count": i, + "score": i / 3, + "flag": i % 2 == 0, + "chat": [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "there"}], + "meta": {"src": "x"}, + "missing": None, + } + for i in range(5) + ] + result = measure_columns(features, rows) + measured, probes, errors = result.stats, result.probes, result.errors + + assert errors == [] + assert measured["text"].text is not None and measured["text"].quality is not None + assert measured["count"].numeric.min == 0.0 + assert measured["count"].categorical.distinct_count == 5 + assert measured["score"].numeric.mean == pytest.approx(sum(i / 3 for i in range(5)) / 5) + assert measured["flag"].categorical.distinct_count == 2 + assert measured["chat"].messages.roles_seen == ["user", "assistant"] + assert "meta" not in measured # a struct with no nulls has nothing worth measuring + assert measured["missing"].null_rate == 1.0 # all-null, kept for the null rate alone + assert set(probes) == {feature.name for feature in features} # every column, typed or not + + +def test_a_column_with_nothing_to_measure_is_not_reported_as_an_error(): + # Absence from `stats` is the normal sparse case. Only a *failure* earns an error, or the two + # would be indistinguishable and the guard would cry wolf on every well-formed struct column. + result = measure_columns([_feature("s", "struct")], [{"s": {"a": 1}}]) + measured, probes, errors = result.stats, result.probes, result.errors + assert measured == {} and errors == [] + assert "s" in probes + + +@pytest.mark.parametrize( + "text", + [ + "", + "plain ascii", + "tabs\tand\nnewlines\r\f\v", + "café naïve", # non-ascii letters, ascii spaces + "a b", # NO-BREAK SPACE: \s matches it, counting six ascii literals would not + "　

", # ideographic space, line separator, paragraph separator + "\U0001f600 emoji beyond the BMP", # 4-byte codepoint: byte overhead != character count + "mixed   café\tend", + ], +) +def test_quality_fast_paths_are_the_same_measurement_as_the_regexes(text): + # The whole risk in this change. `str.count` over six ascii literals is not `\s`, and + # `len(encode) - len` is not a character count -- either would be faster while quietly + # measuring something else. The fast path is only allowed where it is provably identical. + assert _whitespace_count(text) == sum(1 for _ in _WHITESPACE_RUN.finditer(text)) + assert _non_ascii_count(text) == sum(1 for _ in _NON_ASCII_RUN.finditer(text)) + + +def test_the_quality_sample_does_not_alias_against_periodic_data(monkeypatch): + # A set that round-robins over sources, or carries k responses per prompt, is periodic by + # construction. An evenly-spaced step whose spacing shares a factor with that period samples one + # phase and only that phase: 500,000 rows with every tenth corrupt gave a step of ten and a + # repetition score of 1.000 against a truth of 0.100. A contiguous block longer than the period + # sees every phase of it, whatever the period turns out to be. + monkeypatch.setattr(stats_module, "_QUALITY_SAMPLE_ROWS", 1_000) + monkeypatch.setattr(stats_module, "_QUALITY_SAMPLE_BLOCK", 64) + period, n = 10, 100_000 + values = ["aaaaaaaaaaaa" if i % period == 0 else "the quick brown fox" for i in range(n)] + + known = stats_module.StringAccumulator(n) + known.update(values) + assert known.finalize()[0].quality.repetition_score == pytest.approx(1 / period, abs=0.02) + + unknown = stats_module.StringAccumulator(None) + unknown.update(values) + assert unknown.finalize()[0].quality.repetition_score == pytest.approx(1 / period, abs=0.02) + + +def test_a_known_row_count_strides_evenly_and_deterministically(monkeypatch): + # With the row count known up front the stride is fixed, so the sample is spread evenly over the + # whole column -- and two runs over the same bytes agree, which is why no RNG is involved. + monkeypatch.setattr(stats_module, "_QUALITY_SAMPLE_ROWS", 10) + values = ["clean text"] * 100 + ["aaaaaaaaaaaa"] * 100 + + def quality(expected_rows): + acc = stats_module.StringAccumulator(expected_rows) + acc.update(values) + return acc.finalize()[0].quality.repetition_score + + assert quality(len(values)) == quality(len(values)) # deterministic + # Half the column is corrupt and the stride spans it, so the estimate lands near a half. + assert 0.4 <= quality(len(values)) <= 0.6 + + +def test_an_unknown_row_count_thins_as_it_goes_and_stays_unbiased(monkeypatch): + # No footer, so no length to spread blocks over: the cycle starts at one block and doubles as the + # sample fills. Sampling is then densest at the head, which would skew the answer -- weighting + # each sampled row by the rows its block stood for is what corrects it. + monkeypatch.setattr(stats_module, "_QUALITY_SAMPLE_ROWS", 40) + monkeypatch.setattr(stats_module, "_QUALITY_SAMPLE_BLOCK", 8) + values = ["clean text"] * 500 + ["aaaaaaaaaaaa"] * 500 + + acc = stats_module.StringAccumulator(None) + acc.update(values) + score = acc.finalize()[0].quality.repetition_score + + assert acc._cycle > acc._block # it did thin + assert 0.35 <= score <= 0.65 # ...and still found roughly half the column corrupt + + +def test_quality_is_measured_across_the_column_not_its_head(monkeypatch): + # Corruption confined to the second half. A head sample would report a clean column; a stride + # sees it. This is why the sample is strided and not simply the first N rows. + monkeypatch.setattr(stats_module, "_QUALITY_SAMPLE_ROWS", 4) + values = ["ordinary sentence"] * 8 + ["aaaaaaaaaaaa"] * 8 + quality = _stats([_feature("t", "string")], _rows("t", values))["t"].quality + assert quality.repetition_score > 0.4 + + +def test_a_column_under_the_bound_is_measured_exactly(monkeypatch): + monkeypatch.setattr(stats_module, "_QUALITY_SAMPLE_ROWS", 100) + values = ["ordinary sentence"] * 9 + ["aaaaaaaaaaaa"] + quality = _stats([_feature("t", "string")], _rows("t", values))["t"].quality + assert quality.repetition_score == pytest.approx(0.1) # exactly one corrupt row in ten + + +def test_cardinality_is_counted_while_the_column_is_a_vocabulary(): + labels = _stats([_feature("c", "string")], _rows("c", ["yes", "no", "yes", "no"])) + assert labels["c"].categorical.distinct_count == 2 + + # Still counted well past the point where every value is distinct: it is size, not repetition, + # that decides whether a column is a vocabulary. + many = _stats([_feature("t", "string")], _rows("t", [f"unique-{i}" for i in range(50)])) + assert many["t"].categorical.distinct_count == 50 + + +def test_cardinality_stops_at_too_many_values(): + over = _stats([_feature("t", "string")], _rows("t", [f"v{i}" for i in range(_MAX_VOCABULARY_VALUES + 1)])) + assert over["t"].categorical is None # absence is the claim: not a vocabulary + assert over["t"].text is not None # ...but the column is still measured + + at_cap = _stats([_feature("t", "string")], _rows("t", [f"v{i}" for i in range(_MAX_VOCABULARY_VALUES)])) + assert at_cap["t"].categorical.distinct_count == _MAX_VOCABULARY_VALUES + + +def test_one_long_value_settles_it_without_counting(): + # The rule that does the real work: a vocabulary member is short by nature, so a single long + # value proves the column is not one -- on sight, rather than after a thousand of them. + values = ["yes", "no", "x" * (_MAX_VOCABULARY_VALUE_CHARS + 1)] + assert _stats([_feature("t", "string")], _rows("t", values))["t"].categorical is None + + still_short = ["yes", "no", "x" * _MAX_VOCABULARY_VALUE_CHARS] + assert _stats([_feature("t", "string")], _rows("t", still_short))["t"].categorical.distinct_count == 3 + + +def test_cardinality_stops_on_total_bytes_before_the_count(): + # Values individually short enough and few enough, but heavy in aggregate. Without this bound + # the other two would admit 1024 x 256 chars -- four times the byte budget. + values = [f"{i:04d}" + "x" * 200 for i in range(_MAX_VOCABULARY_BYTES // 200)] + assert len(values) < _MAX_VOCABULARY_VALUES # the count bound is not what stops this + assert _stats([_feature("t", "string")], _rows("t", values))["t"].categorical is None + + +def test_derive_stats_never_quotes_values(): + # Quoting needs a role, and roles are not assigned when stats are measured. Filling them in + # afterwards rather than redacting means a skipped pass stores nothing instead of leaking. + stats = _stats([_feature("c", "string")], _rows("c", ["yes", "no"])) + assert stats["c"].categorical.values is None + + +def test_bool_column_gets_a_measured_class_balance(): + stats = _stats([_feature("label", "bool")], _rows("label", [True, False, True])) + assert stats["label"].categorical.distinct_count == 2 + + +# --- numeric ------------------------------------------------------------------------------------- + + +def test_numeric_stats_and_cardinality(): + stats = _stats([_feature("n", "int64")], _rows("n", [0, 4, 2, 2, 3]))["n"] + assert (stats.numeric.min, stats.numeric.max) == (0.0, 4.0) + assert stats.numeric.mean == 2.2 + assert stats.categorical.distinct_count == 4 # {0, 2, 3, 4} + + +def test_numeric_cardinality_counts_without_quoting(): + stats = _stats([_feature("n", "int64")], _rows("n", [1, 2, 3]))["n"] + assert stats.categorical.distinct_count == 3 + assert stats.categorical.values is None + + +def test_numeric_stats_ignore_non_finite_values(): + # NaN / +-inf poison min/max/mean and serialize to JSON null, which then fails to re-validate + # against NumericStats' required floats -- making the whole profile unreadable. Drop them. + values = [1.0, float("nan"), 3.0, float("inf"), float("-inf"), 5.0] + stats = _stats([_feature("n", "float64")], _rows("n", values))["n"] + assert (stats.numeric.min, stats.numeric.max, stats.numeric.mean) == (1.0, 5.0, 3.0) + ColumnStats.model_validate_json(stats.model_dump_json()) # round-trips: no NaN/inf leaked into JSON + + +def test_numeric_all_non_finite_yields_no_numeric_summary(): + stats = _stats([_feature("n", "float64")], _rows("n", [float("nan"), float("inf")])) + assert stats.get("n") is None or stats["n"].numeric is None + + +# --- messages ------------------------------------------------------------------------------------ + + +def test_message_stats_shape_signals(): + rows = [ + {"m": [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "hello there"}]}, + {"m": [{"role": "user", "content": "again"}, {"role": "assistant", "content": "yes"}]}, + ] + stats = _stats([_feature("m", "messages")], rows)["m"] + assert stats.messages.turns.max == 2 + assert stats.messages.roles_seen == ["user", "assistant"] # first-seen order + assert stats.messages.ends_with_assistant_rate == 1.0 + assert stats.messages.valid_alternation_rate == 1.0 + assert stats.messages.has_tool_calls is False + + +def test_message_stats_detects_tool_calls_and_user_ending(): + rows = [{"m": [{"role": "user", "content": "run"}, {"role": "assistant", "tool_calls": [{"id": "1"}]}]}] + stats = _stats([_feature("m", "messages")], rows)["m"] + assert stats.messages.has_tool_calls is True + assert stats.messages.ends_with_assistant_rate == 1.0 # last turn is the assistant tool call + + +def test_message_ends_with_user_turn_is_prompt_only_signal(): + rows = [{"m": [{"role": "user", "content": "solve"}]}] + stats = _stats([_feature("m", "messages")], rows)["m"] + assert stats.messages.ends_with_assistant_rate == 0.0 + + +def test_message_stats_read_sharegpt_from_value(): + rows = [{"m": [{"from": "human", "value": "hi"}, {"from": "gpt", "value": "hello there"}]}] + stats = _stats([_feature("m", "messages")], rows)["m"] + assert stats.messages.roles_seen == ["human", "gpt"] # verbatim, not normalized + assert stats.messages.content_chars.max == len("hi") + len("hello there") + assert stats.messages.ends_with_assistant_rate == 1.0 # "gpt" is the responder turn + + +def test_assistant_equivalent_roles_count_as_the_training_target(): + # Matching only the literal "assistant" made every other convention look prompt-only. + for responder in ("assistant", "gpt", "bot", "model", "AI"): + rows = [{"m": [{"role": "user", "content": "q"}, {"role": responder, "content": "a"}]}] + stats = _stats([_feature("m", "messages")], rows)["m"] + assert stats.messages.ends_with_assistant_rate == 1.0, responder + + +def test_non_string_role_does_not_break_measurement(): + # roles_seen is typed list[str]; a numeric role used to raise a ValidationError from inside the + # one stage the pipeline did not guard, aborting the whole profile. + rows = [{"m": [{"role": 1, "content": "hi"}]}] + stats = _stats([_feature("m", "messages")], rows)["m"] + assert stats.messages.roles_seen == ["1"] + + +def test_declared_but_unset_tool_calls_is_not_tool_use(): + # parquet materializes every declared struct field, so `"tool_calls" in message` reported tool + # use for any schema that merely declares the field. + rows = [{"m": [{"role": "user", "content": "hi", "tool_calls": None}]}] + stats = _stats([_feature("m", "messages")], rows)["m"] + assert stats.messages.has_tool_calls is False + + +def test_message_content_parts_tolerate_non_string_text(): + # A VLM-style content part whose "text" key is present but not a string must not crash measurement. + rows = [{"m": [{"role": "user", "content": [{"type": "image"}, {"type": "text", "text": None}]}]}] + stats = _stats([_feature("m", "messages")], rows)["m"] + assert stats.messages.content_chars.max == 0 # no measurable text, and no crash + + +# --- sparsity and null rate ---------------------------------------------------------------------- + + +def test_unmeasured_dtypes_are_omitted(): + features = [_feature("s", "struct"), _feature("j", "json")] + rows = [{"s": {"a": 1}, "j": object()}] + assert _stats(features, rows) == {} + + +def test_null_rate_is_reported(): + stats = _stats([_feature("t", "string")], _rows("t", ["a", None, "c", None]))["t"] + assert stats.null_rate == 0.5 + + +# --- content probes ------------------------------------------------------------------------------ + + +def test_probes_are_measured_for_every_column_not_just_named_ones(): + # The whole point of measuring probes here rather than in classify: a column whose name the + # alias table does not know still gets its content read. + features = [_feature("q", "string"), _feature("a", "string")] + rows = [{"q": "what is 2+2?", "a": "add them #### 4"}, {"q": "and 3+3?", "a": "no final answer"}] + probes = _probes(features, rows) + + assert set(probes) == {"q", "a"} + assert probes["a"].texts == 2 + assert probes["a"].extractable_answer == 1 + assert probes["q"].extractable_answer == 0 + + +def test_probes_read_the_final_turn_of_a_chat_column(): + rows = [{"m": [{"role": "user", "content": "q"}, {"role": "assistant", "content": "steps #### 7"}]}] + probes = _probes([_feature("m", "messages")], rows) + assert probes["m"].texts == 1 + assert probes["m"].extractable_answer == 1 + + +def test_probes_read_the_sharegpt_message_spelling(): + # {from, value} is handled in schema derivation and message stats; reading only {role, content} + # here cost every ShareGPT-shaped dataset its verifiability. + rows = [{"m": [{"from": "human", "value": "q"}, {"from": "gpt", "value": "steps #### 7"}]}] + probes = _probes([_feature("m", "messages")], rows) + assert probes["m"].texts == 1 + assert probes["m"].extractable_answer == 1 + + +def test_probes_count_non_empty_across_container_dtypes(): + # `non_empty` is what a ground_truth column's coverage is measured from, and a verification + # target is just as often a list or struct as a string. + features = [_feature("gt", "list")] + rows = [{"gt": [{"in": "1"}]}, {"gt": []}, {"gt": None}] + probes = _probes(features, rows) + assert probes["gt"].rows == 3 + assert probes["gt"].non_empty == 1 + + +def test_probes_detect_embedded_transcripts(): + rows = [{"c": "\n\nHuman: hi\n\nAssistant: hello"}, {"c": "plain prose"}] + probes = _probes([_feature("c", "string")], rows) + assert probes["c"].transcript_marker == 1 + + +# --- quoting a controlled vocabulary --------------------------------------------------------------- + + +def _quoted(name, dtype, values, role): + """Run the real two-step: measure, then quote by role, and report what was stored.""" + feature = _feature(name, dtype) + feature.semantic_role = role + rows = _rows(name, values) + measured = measure_columns([feature], rows) + quote_enumerations([feature], measured.stats, measured.vocabularies) + return measured.stats[name].categorical.values + + +def test_quotes_a_controlled_vocabulary_role(): + assert _quoted("label", "bool", [True, False, True], "label") == ["False", "True"] + assert _quoted("source", "string", ["gsm8k", "math", "gsm8k"], "provenance") == ["gsm8k", "math"] + assert _quoted("category", "string", ["code", "math"], "meta") == ["code", "math"] + + +def test_refuses_to_quote_free_text_however_few_distinct_values(): + # The failure the cardinality gate could not see: in a tiny dataset every column holds under the + # cap, so a whole column of prompts was quotable and the profile stored it verbatim. + rows = ["Patient Alice, SSN 123-45-6000", "Patient Bob, SSN 123-45-6001"] + assert _quoted("prompt", "string", rows, "prompt") is None + assert _quoted("completion", "string", rows, "completion") is None + assert _quoted("chosen", "string", rows, "chosen") is None + + +def test_refuses_to_quote_an_unroled_column(): + # An unrecognized column is unknown, which here is the same as free text: an allowlist means it + # fails to silence rather than to exposure. + assert _quoted("mystery", "string", ["a", "b"], None) is None + + +def test_refuses_to_quote_a_vocabulary_larger_than_the_cap(): + # Role grants permission; cardinality still bounds the size, so a provenance column holding a + # URL list is not mistaken for an enumeration. + assert _quoted("source", "string", [f"src-{i}" for i in range(40)], "provenance") is None diff --git a/pyproject.toml b/pyproject.toml index 8d030337b7..5217d7b606 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -184,6 +184,7 @@ enabled-plugins = [ "nemo-switchyard", "nemo-agents-plugin", "nemo-deployments-plugin[docker,k8s]", + "nemo-datasets-plugin", "nemo-customizer-plugin", "nemo-automodel-plugin", "nemo-optimization-plugin", @@ -405,6 +406,7 @@ nemo-safe-synthesizer-plugin = { workspace = true } nemo-switchyard = { workspace = true } nemo-agents-plugin = { workspace = true } nemo-deployments-plugin = { workspace = true } +nemo-datasets-plugin = { workspace = true } nemo-agents-example-calculator = { workspace = true } nemo-agents-example-email-phishing = { workspace = true } nemo-agents-example-email-security = { workspace = true } @@ -463,6 +465,7 @@ members = [ "plugins/nemo-safe-synthesizer", "plugins/nemo-switchyard", "plugins/nemo-agents", + "plugins/nemo-datasets", "plugins/nemo-deployments", "plugins/nemo-insights", "plugins/nemo-eval-author", @@ -585,6 +588,9 @@ extra-paths = [ # example-plugin is not a workspace member; add its src so ty can # resolve nmp.example_plugin imports when checking plugin test files. "plugins/example-plugin/src", + # nemo-datasets plugin is a workspace member; add its src so ty can + # resolve nemo_datasets_plugin imports from plugin test files. + "plugins/nemo-datasets/src", # nemo-guardrails plugin is not a workspace member; add its src so ty can # resolve nemo_guardrails_plugin imports when checking plugin test files. "plugins/nemo-guardrails/src", diff --git a/pytest.ini b/pytest.ini index 0568982260..6fbd695fc8 100644 --- a/pytest.ini +++ b/pytest.ini @@ -9,6 +9,7 @@ python_functions = test_* pythonpath = . plugins/example-plugin/src + plugins/nemo-datasets/src plugins/nemo-deployments/src plugins/nemo-deployments/tests/unit plugins/nemo-deployments/tests/integration diff --git a/uv.lock b/uv.lock index cf41b73710..270c23e799 100644 --- a/uv.lock +++ b/uv.lock @@ -31,6 +31,7 @@ members = [ "nemo-automodel-plugin", "nemo-customizer-plugin", "nemo-data-designer-plugin", + "nemo-datasets-plugin", "nemo-deployments-plugin", "nemo-eval-author-plugin", "nemo-evaluator-plugin", @@ -4397,6 +4398,25 @@ requires-dist = [ ] provides-extras = ["test", "data-designer-nemo", "nemo-platform-plugin"] +[[package]] +name = "nemo-datasets-plugin" +version = "0.1.0" +source = { editable = "plugins/nemo-datasets" } +dependencies = [ + { name = "nemo-platform-plugin", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nemo-platform-sdk", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pyarrow", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pydantic", extra = ["email"], marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] + +[package.metadata] +requires-dist = [ + { name = "nemo-platform-plugin", editable = "packages/nemo_platform_plugin" }, + { name = "nemo-platform-sdk", editable = "sdk/python/nemo-platform" }, + { name = "pyarrow", specifier = ">=19.0.1" }, + { name = "pydantic", specifier = ">=2.10.3" }, +] + [[package]] name = "nemo-deployments-plugin" version = "0.0.0" @@ -6641,6 +6661,7 @@ core-services = [ { name = "nemo-automodel-plugin", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "nemo-customizer-plugin", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "nemo-data-designer-plugin", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nemo-datasets-plugin", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "nemo-deployments-plugin", extra = ["docker", "k8s"], marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "nemo-eval-author-plugin", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "nemo-evaluator-plugin", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, @@ -6740,6 +6761,7 @@ enabled-plugins = [ { name = "nemo-automodel-plugin", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "nemo-customizer-plugin", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "nemo-data-designer-plugin", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nemo-datasets-plugin", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "nemo-deployments-plugin", extra = ["docker", "k8s"], marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "nemo-eval-author-plugin", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "nemo-evaluator-plugin", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, @@ -6763,6 +6785,7 @@ functional-services = [ { name = "nemo-automodel-plugin", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "nemo-customizer-plugin", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "nemo-data-designer-plugin", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nemo-datasets-plugin", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "nemo-deployments-plugin", extra = ["docker", "k8s"], marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "nemo-eval-author-plugin", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "nemo-evaluator-plugin", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, @@ -6861,6 +6884,7 @@ core-services = [ { name = "nemo-automodel-plugin", editable = "plugins/nemo-automodel" }, { name = "nemo-customizer-plugin", editable = "plugins/nemo-customizer" }, { name = "nemo-data-designer-plugin", editable = "plugins/nemo-data-designer" }, + { name = "nemo-datasets-plugin", editable = "plugins/nemo-datasets" }, { name = "nemo-deployments-plugin", extras = ["docker", "k8s"], editable = "plugins/nemo-deployments" }, { name = "nemo-eval-author-plugin", marker = "python_full_version < '3.14'", editable = "plugins/nemo-eval-author" }, { name = "nemo-evaluator-plugin", editable = "plugins/nemo-evaluator" }, @@ -6963,6 +6987,7 @@ enabled-plugins = [ { name = "nemo-automodel-plugin", editable = "plugins/nemo-automodel" }, { name = "nemo-customizer-plugin", editable = "plugins/nemo-customizer" }, { name = "nemo-data-designer-plugin", editable = "plugins/nemo-data-designer" }, + { name = "nemo-datasets-plugin", editable = "plugins/nemo-datasets" }, { name = "nemo-deployments-plugin", extras = ["docker", "k8s"], editable = "plugins/nemo-deployments" }, { name = "nemo-eval-author-plugin", marker = "python_full_version < '3.14'", editable = "plugins/nemo-eval-author" }, { name = "nemo-evaluator-plugin", editable = "plugins/nemo-evaluator" }, @@ -6986,6 +7011,7 @@ functional-services = [ { name = "nemo-automodel-plugin", editable = "plugins/nemo-automodel" }, { name = "nemo-customizer-plugin", editable = "plugins/nemo-customizer" }, { name = "nemo-data-designer-plugin", editable = "plugins/nemo-data-designer" }, + { name = "nemo-datasets-plugin", editable = "plugins/nemo-datasets" }, { name = "nemo-deployments-plugin", extras = ["docker", "k8s"], editable = "plugins/nemo-deployments" }, { name = "nemo-eval-author-plugin", marker = "python_full_version < '3.14'", editable = "plugins/nemo-eval-author" }, { name = "nemo-evaluator-plugin", editable = "plugins/nemo-evaluator" },