datasets-lint 0.1.0__py3-none-any.whl

datasets_lint/__init__.py

@@ -0,0 +1,21 @@
+"""datasets-lint — ruff for training data.
+
+Audit a dataset for duplicates, PII, distribution problems and eval-set
+contamination. Runs offline, on CPU, and exits non-zero when thresholds
+are breached so it can gate CI.
+"""
+
+__version__ = "0.1.0"
+
+from datasets_lint.dataset import Dataset, load_dataset  # noqa: E402
+from datasets_lint.result import CheckResult, Finding, Report, Severity  # noqa: E402
+
+__all__ = [
+    "__version__",
+    "CheckResult",
+    "Dataset",
+    "Finding",
+    "Report",
+    "Severity",
+    "load_dataset",
+]

datasets_lint/baseline.py

@@ -0,0 +1,71 @@
+"""Compare a run against a previous JSON report (`--baseline`)."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any
+
+from datasets_lint.result import Report, Severity
+
+_RANK = {"ok": 0, "warn": 1, "error": 2}
+
+
+def load_baseline(path: str | Path) -> dict[str, Any]:
+    data = json.loads(Path(path).read_text(encoding="utf-8"))
+    if "checks" not in data:
+        raise ValueError(f"{path} is not a datasets-lint JSON report")
+    return data
+
+
+def _num(value: Any) -> float | None:
+    if isinstance(value, bool) or not isinstance(value, (int, float)):
+        return None
+    return float(value)
+
+
+def diff_lines(report: Report, baseline: dict[str, Any]) -> list[str]:
+    """Human-readable deltas: `dup.near_dup_pct 18.2 → 12.0 (-6.2) better`."""
+    old = {c.get("id"): c for c in baseline.get("checks", [])}
+    lines: list[str] = []
+
+    old_rows = baseline.get("dataset", {}).get("audited_rows")
+    if isinstance(old_rows, int) and old_rows != report.audited_rows:
+        lines.append(f"audited rows: {old_rows:,} → {report.audited_rows:,}")
+
+    for result in report.results:
+        prev = old.get(result.check_id)
+        if not prev:
+            lines.append(f"{result.check_id}: new check (no baseline)")
+            continue
+
+        prev_sev = str(prev.get("severity", "ok"))
+        now_sev = result.severity.value
+        if _RANK[now_sev] != _RANK.get(prev_sev, 0):
+            arrow = "better" if _RANK[now_sev] < _RANK.get(prev_sev, 0) else "WORSE"
+            lines.append(f"{result.check_id}: severity {prev_sev} → {now_sev} ({arrow})")
+
+        prev_metrics = prev.get("metrics", {})
+        for key, value in result.metrics.items():
+            before = _num(prev_metrics.get(key))
+            after = _num(value)
+            if before is None or after is None or before == after:
+                continue
+            delta = after - before
+            lines.append(
+                f"{result.check_id}.{key}: {before:,.3f} → {after:,.3f} "
+                f"({delta:+,.3f}) {'better' if delta < 0 else 'worse'}"
+            )
+
+    if not lines:
+        lines.append("no metric changes against baseline")
+    return lines
+
+
+def verdict_delta(report: Report, baseline: dict[str, Any]) -> Severity:
+    """Return the baseline verdict if the run did not improve on it, else the new one."""
+    try:
+        prev = Severity(str(baseline.get("verdict", "ok")))
+    except ValueError:
+        return report.worst
+    return prev if prev.rank > report.worst.rank else report.worst

datasets_lint/checks/__init__.py

@@ -0,0 +1,18 @@
+"""Check registry. Add a new check by importing it and appending an instance."""
+
+from __future__ import annotations
+
+from datasets_lint.checks.base import Check, Context
+from datasets_lint.checks.dup import DupCheck
+from datasets_lint.checks.pii import PiiCheck
+from datasets_lint.checks.stats import StatsCheck
+
+ALL_CHECKS: list[Check] = [
+    StatsCheck(),
+    DupCheck(),
+    PiiCheck(),
+]
+
+CHECKS_BY_ID: dict[str, Check] = {c.id: c for c in ALL_CHECKS}
+
+__all__ = ["ALL_CHECKS", "CHECKS_BY_ID", "Check", "Context"]

datasets_lint/checks/base.py

@@ -0,0 +1,38 @@
+"""Check protocol and the context every check receives."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Any
+
+from datasets_lint.config import Config
+from datasets_lint.dataset import Dataset
+from datasets_lint.result import CheckResult
+
+
+@dataclass
+class Context:
+    config: Config
+    seed: int = 0
+
+    def threshold(self, check_id: str, key: str, default: Any) -> Any:
+        return self.config.threshold(check_id, key, default)
+
+
+class Check(ABC):
+    """A single data-quality check.
+
+    Subclass, set ``id``/``title``/``description``, implement ``run`` and add an
+    instance to ``datasets_lint.checks.ALL_CHECKS``. That is the whole plugin API.
+    """
+
+    id: str = ""
+    title: str = ""
+    description: str = ""
+
+    @abstractmethod
+    def run(self, ds: Dataset, ctx: Context) -> CheckResult: ...
+
+    def __repr__(self) -> str:  # pragma: no cover - debugging aid
+        return f"<Check {self.id}>"

datasets_lint/checks/dup.py

@@ -0,0 +1,187 @@
+"""Duplicate detection: exact hashing plus MinHash/LSH near-duplicates.
+
+Design note: MinHash *estimates* Jaccard, and on short rows (instruction-style
+data) a 64-permutation sketch has enough variance to miss an obvious duplicate.
+So LSH is used only to generate candidate pairs (with a deliberately loose
+threshold) and every candidate is then verified with an exact Jaccard on the
+shingle sets. Result: no false negatives from sketch noise, no O(n^2) scan.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import sys
+from collections import defaultdict
+from typing import Any
+
+from datasketch import MinHash, MinHashLSH
+
+from datasets_lint.checks.base import Check, Context
+from datasets_lint.dataset import Dataset, normalize
+from datasets_lint.result import CheckResult, Finding, Severity
+
+
+class _UnionFind:
+    def __init__(self) -> None:
+        self.parent: dict[int, int] = {}
+
+    def find(self, x: int) -> int:
+        self.parent.setdefault(x, x)
+        root = x
+        while self.parent[root] != root:
+            root = self.parent[root]
+        while self.parent[x] != root:
+            self.parent[x], x = root, self.parent[x]
+        return root
+
+    def union(self, a: int, b: int) -> None:
+        ra, rb = self.find(a), self.find(b)
+        if ra != rb:
+            self.parent[rb] = ra
+
+    def clusters(self) -> dict[int, list[int]]:
+        groups: dict[int, list[int]] = defaultdict(list)
+        for x in self.parent:
+            groups[self.find(x)].append(x)
+        return groups
+
+
+def _shingles(text: str, size: int) -> frozenset[str]:
+    words = text.split()
+    if not words:
+        return frozenset()
+    if len(words) <= size:
+        return frozenset({" ".join(words)})
+    # interning keeps memory flat: corpora repeat the same shingles constantly
+    return frozenset(
+        sys.intern(" ".join(words[i : i + size])) for i in range(len(words) - size + 1)
+    )
+
+
+class DupCheck(Check):
+    id = "dup"
+    title = "Duplicates"
+    description = "Exact and near-duplicate rows (MinHash/LSH), plus cross-split overlap."
+
+    def run(self, ds: Dataset, ctx: Context) -> CheckResult:
+        threshold = float(ctx.threshold(self.id, "threshold", 0.7))
+        num_perm = int(ctx.threshold(self.id, "num_perm", 64))
+        shingle_size = int(ctx.threshold(self.id, "shingle_size", 3))
+        max_chars = int(ctx.threshold(self.id, "max_chars", 300))
+        max_near_dup_pct = float(ctx.threshold(self.id, "max_near_dup_pct", 2.0))
+
+        total = ds.audited_rows
+        metrics: dict[str, Any] = {
+            "threshold": threshold,
+            "shingle_size": shingle_size,
+            "max_chars": max_chars,
+        }
+        if total == 0:
+            return CheckResult(
+                check_id=self.id,
+                title=self.title,
+                severity=Severity.OK,
+                summary="no rows to check",
+                metrics=metrics,
+            )
+
+        # Long documents are truncated: the opening characters already carry enough
+        # signal, and it keeps the sketch and the exact sets small.
+        texts = [(i, normalize(t)[:max_chars]) for i, t in ds.texts()]
+
+        exact: dict[str, list[int]] = defaultdict(list)
+        for i, norm in texts:
+            if norm:
+                exact[hashlib.sha1(norm.encode()).hexdigest()].append(i)
+        exact_groups = {k: v for k, v in exact.items() if len(v) > 1}
+        exact_rows = sum(len(v) for v in exact_groups.values())
+        metrics["exact_dup_rows"] = exact_rows
+        metrics["exact_dup_pct"] = round(100.0 * exact_rows / total, 3)
+        metrics["exact_dup_groups"] = len(exact_groups)
+
+        sets: dict[int, frozenset[str]] = {}
+        sketches: dict[int, MinHash] = {}
+        lsh = MinHashLSH(threshold=max(0.05, threshold - 0.15), num_perm=num_perm)
+        for i, norm in texts:
+            grams = _shingles(norm, shingle_size)
+            if not grams:
+                continue
+            sets[i] = grams
+            mh = MinHash(num_perm=num_perm)
+            for sh in grams:
+                mh.update(sh.encode())
+            sketches[i] = mh
+            lsh.insert(str(i), mh)
+
+        uf = _UnionFind()
+        candidate_pairs = 0
+        verified_pairs = 0
+        for i, grams in sets.items():
+            for other in lsh.query(sketches[i]):
+                j = int(other)
+                if j <= i:
+                    continue
+                candidate_pairs += 1
+                other_grams = sets[j]
+                inter = len(grams & other_grams)
+                if inter == 0:
+                    continue
+                jaccard = inter / len(grams | other_grams)
+                if jaccard >= threshold:
+                    verified_pairs += 1
+                    uf.union(i, j)
+
+        clusters = {k: v for k, v in uf.clusters().items() if len(v) > 1}
+        near_rows = sum(len(v) for v in clusters.values())
+        metrics["near_dup_rows"] = near_rows
+        metrics["near_dup_pct"] = round(100.0 * near_rows / total, 3)
+        metrics["near_dup_clusters"] = len(clusters)
+        metrics["candidate_pairs"] = candidate_pairs
+        metrics["verified_pairs"] = verified_pairs
+        if clusters:
+            metrics["largest_cluster"] = max(len(v) for v in clusters.values())
+
+        cross_source_rows = 0
+        if len(ds.sources) > 1:
+            for members in clusters.values():
+                if len({ds.source_of(ds.rows[m]) for m in members}) > 1:
+                    cross_source_rows += len(members)
+            metrics["cross_source_dup_rows"] = cross_source_rows
+            metrics["cross_source_dup_pct"] = round(100.0 * cross_source_rows / total, 3)
+
+        findings: list[Finding] = []
+        for members in sorted(clusters.values(), key=len, reverse=True)[:25]:
+            members = sorted(members)
+            srcs = sorted({ds.source_of(ds.rows[m]) for m in members})
+            sample = normalize(ds.text(ds.rows[members[0]]))[:120]
+            findings.append(
+                Finding(
+                    message=f"{len(members)} near-duplicate rows"
+                    + (f" across {', '.join(srcs)}" if len(srcs) > 1 else "")
+                    + f': "{sample}…"',
+                    row=members[0],
+                    detail={"size": len(members), "sources": srcs, "sample_rows": members[:5]},
+                )
+            )
+
+        severity = Severity.OK
+        if metrics["near_dup_pct"] > max_near_dup_pct:
+            severity = Severity.ERROR
+        elif metrics["near_dup_pct"] > max_near_dup_pct / 2:
+            severity = Severity.WARN
+
+        summary = (
+            f"{metrics['near_dup_pct']:.2f}% near-duplicate "
+            f"({len(clusters):,} clusters), {metrics['exact_dup_pct']:.2f}% exact"
+        )
+        if "cross_source_dup_pct" in metrics:
+            summary += f" · {metrics['cross_source_dup_pct']:.2f}% cross-file"
+
+        return CheckResult(
+            check_id=self.id,
+            title=self.title,
+            severity=severity,
+            summary=summary,
+            metrics=metrics,
+            findings=findings,
+        )

datasets_lint/checks/pii.py

@@ -0,0 +1,134 @@
+"""PII detection: high-confidence secrets and warn-level contact data."""
+
+from __future__ import annotations
+
+import re
+from typing import Any
+
+from datasets_lint.checks.base import Check, Context
+from datasets_lint.dataset import Dataset
+from datasets_lint.result import CheckResult, Finding, Severity
+
+# (kind, severity, pattern). Patterns are intentionally conservative: a false
+# positive on a phone number is annoying, a false negative on an API key is an
+# incident. High-confidence secrets default to `error`.
+PATTERNS: list[tuple[str, Severity, re.Pattern[str]]] = [
+    (
+        "api_key",
+        Severity.ERROR,
+        re.compile(
+            r"\b(?:sk-[A-Za-z0-9]{16,}|sk-ant-[A-Za-z0-9_\-]{16,}|"
+            r"AKIA[0-9A-Z]{16}|ghp_[A-Za-z0-9]{36}|gho_[A-Za-z0-9]{36}|"
+            r"xox[baprs]-[A-Za-z0-9\-]{10,}|AIza[0-9A-Za-z_\-]{35})\b"
+        ),
+    ),
+    (
+        "private_key",
+        Severity.ERROR,
+        re.compile(r"-----BEGIN (?:RSA |EC |OPENSSH )?PRIVATE KEY-----"),
+    ),
+    ("iban", Severity.ERROR, re.compile(r"\b[A-Z]{2}\d{2}[A-Z0-9]{11,30}\b")),
+    ("ssn", Severity.ERROR, re.compile(r"(?<!\d)\d{3}-\d{2}-\d{4}(?!\d)")),
+    (
+        "email",
+        Severity.WARN,
+        re.compile(r"\b[A-Za-z0-9._%+\-]+@[A-Za-z0-9.\-]+\.[A-Za-z]{2,}\b"),
+    ),
+    (
+        "phone",
+        Severity.WARN,
+        re.compile(
+            r"(?<![\d.])(?:\+\d{1,3}[\s.\-]?)?(?:\(\d{2,4}\)[\s.\-]?)?\d{3,4}[\s.\-]\d{3,4}(?![\d.])"
+        ),
+    ),
+    (
+        "ipv4",
+        Severity.WARN,
+        re.compile(r"\b(?:(?:25[0-5]|2[0-4]\d|1?\d?\d)\.){3}(?:25[0-5]|2[0-4]\d|1?\d?\d)\b"),
+    ),
+]
+
+# Common false positives for the phone pattern (dates, versions, ranges).
+_PHONE_FALSE_POSITIVES = re.compile(r"^(?:19|20)\d{2}[\s.\-]\d{1,2}[\s.\-]\d{1,2}$|^\d+\.\d+\.\d+$")
+
+
+def redact(value: str) -> str:
+    """Keep the shape, hide the payload: `a***@example.com`, `sk-***`, ..."""
+    if len(value) <= 4:
+        return "*" * len(value)
+    if "@" in value:
+        local, _, domain = value.partition("@")
+        return f"{local[:1]}***@{domain}"
+    if re.fullmatch(r"[A-Za-z0-9_\-]+", value) and len(value) > 12:
+        return f"{value[:4]}***"
+    return f"{value[:2]}***{value[-2:]}"
+
+
+class PiiCheck(Check):
+    id = "pii"
+    title = "PII & secrets"
+    description = "Emails, phones, IPs, IBANs, SSNs and API keys. Findings are redacted."
+
+    def run(self, ds: Dataset, ctx: Context) -> CheckResult:
+        max_error_pct = float(ctx.threshold(self.id, "max_error_findings_pct", 0.1))
+        max_warn_pct = float(ctx.threshold(self.id, "max_warn_findings_pct", 5.0))
+
+        counts: dict[str, int] = {}
+        by_severity: dict[str, int] = {"warn": 0, "error": 0}
+        findings: list[Finding] = []
+        affected_rows = 0
+
+        for i, text in ds.texts():
+            if not text:
+                continue
+            row_hit = False
+            for kind, severity, pattern in PATTERNS:
+                for match in pattern.finditer(text):
+                    value = match.group(0)
+                    if kind == "phone" and _PHONE_FALSE_POSITIVES.match(value.strip()):
+                        continue
+                    if kind == "iban" and not re.search(r"\d", value):
+                        continue
+                    row_hit = True
+                    counts[kind] = counts.get(kind, 0) + 1
+                    by_severity[severity.value] += 1
+                    if len(findings) < 50:
+                        findings.append(
+                            Finding(
+                                message=f"{kind}: {redact(value)}",
+                                row=i,
+                                detail={"kind": kind, "severity": severity.value},
+                            )
+                        )
+            if row_hit:
+                affected_rows += 1
+
+        total = max(ds.audited_rows, 1)
+        affected_pct = round(100.0 * affected_rows / total, 3)
+        metrics: dict[str, Any] = {
+            "rows_with_pii": affected_rows,
+            "rows_with_pii_pct": affected_pct,
+            "findings_by_kind": counts,
+            "findings_by_severity": by_severity,
+        }
+
+        severity = Severity.OK
+        if by_severity["error"] > 0 and affected_pct > max_error_pct:
+            severity = Severity.ERROR
+        elif affected_pct > max_warn_pct or by_severity["error"] > 0:
+            severity = Severity.WARN
+
+        if counts:
+            parts = ", ".join(f"{k} {v}" for k, v in sorted(counts.items(), key=lambda kv: -kv[1]))
+            summary = f"{affected_pct:.2f}% of rows contain PII ({parts})"
+        else:
+            summary = "no PII patterns found"
+
+        return CheckResult(
+            check_id=self.id,
+            title=self.title,
+            severity=severity,
+            summary=summary,
+            metrics=metrics,
+            findings=findings,
+        )

datasets_lint/checks/stats.py

@@ -0,0 +1,96 @@
+"""Distribution sanity: empties, nulls, length outliers, label balance."""
+
+from __future__ import annotations
+
+from collections import Counter
+from typing import Any
+
+from datasets_lint.checks.base import Check, Context
+from datasets_lint.dataset import Dataset, normalize
+from datasets_lint.result import CheckResult, Finding, Severity
+
+
+class StatsCheck(Check):
+    id = "stats"
+    title = "Distribution & shape"
+    description = "Row counts, nulls, empty texts, length percentiles, label balance."
+
+    def run(self, ds: Dataset, ctx: Context) -> CheckResult:
+        findings: list[Finding] = []
+        metrics: dict[str, Any] = {"rows": ds.audited_rows, "columns": len(ds.columns)}
+
+        null_pct: dict[str, float] = {}
+        for col in ds.columns:
+            nulls = sum(1 for r in ds.rows if r.get(col) is None)
+            null_pct[col] = round(100.0 * nulls / max(ds.audited_rows, 1), 3)
+        metrics["null_pct"] = null_pct
+
+        empty_rows = 0
+        lengths: list[int] = []
+        for i, text in ds.texts():
+            norm = normalize(text)
+            if not norm:
+                empty_rows += 1
+                findings.append(Finding(message="empty or whitespace-only text", row=i))
+            lengths.append(len(norm))
+
+        empty_pct = round(100.0 * empty_rows / max(ds.audited_rows, 1), 3)
+        metrics["empty_rows"] = empty_rows
+        metrics["empty_pct"] = empty_pct
+        metrics["text_column"] = ds.text_column
+
+        if lengths:
+            s = sorted(lengths)
+            n = len(s)
+            metrics["text_len"] = {
+                "p50": s[n // 2],
+                "p95": s[min(n - 1, int(n * 0.95))],
+                "max": s[-1],
+                "mean": round(sum(s) / n, 1),
+            }
+            metrics["n_unique_texts"] = len({normalize(t) for _, t in ds.texts()})
+
+        if ds.label_column:
+            counts = Counter(
+                str(r.get(ds.label_column)) for r in ds.rows if r.get(ds.label_column) is not None
+            )
+            total = sum(counts.values()) or 1
+            balance = {k: round(100.0 * v / total, 2) for k, v in counts.most_common(20)}
+            metrics["label_column"] = ds.label_column
+            metrics["label_balance_pct"] = balance
+            if len(balance) > 1:
+                smallest = min(balance.values())
+                if smallest < ctx.threshold(self.id, "min_label_pct", 1.0):
+                    findings.append(
+                        Finding(
+                            message=f"rare label: {smallest:.2f}% of rows "
+                            f"(threshold {ctx.threshold(self.id, 'min_label_pct', 1.0)}%)",
+                            detail={"label_balance_pct": balance},
+                        )
+                    )
+
+        max_empty = float(ctx.threshold(self.id, "max_empty_pct", 1.0))
+        max_null = float(ctx.threshold(self.id, "max_null_pct", 50.0))
+
+        severity = Severity.OK
+        if empty_pct > max_empty:
+            severity = Severity.ERROR
+        elif any(v > max_null for v in null_pct.values()):
+            severity = Severity.WARN
+
+        worst_null = max(null_pct.values()) if null_pct else 0.0
+        summary = (
+            f"{ds.audited_rows:,} rows · {empty_pct:.2f}% empty · "
+            f"worst column {worst_null:.1f}% null"
+        )
+        if "text_len" in metrics:
+            summary += f" · median {metrics['text_len']['p50']:,} chars"
+
+        return CheckResult(
+            check_id=self.id,
+            title=self.title,
+            severity=severity,
+            summary=summary,
+            metrics=metrics,
+            findings=findings[:50],
+        )