deduplicate_dataset.py

# scikitplot/_externals/_sphinx_ext/_sphinx_ai_assistant/_hf_spaces_proxy/deduplicate_dataset.py
#
# flake8: noqa: D213
#
# Authors: The scikit-plots developers
# SPDX-License-Identifier: BSD-3-Clause

r"""
deduplicate_dataset.py
======================
Canonical deduplication script for scikit-plots/ai-assistant-contributions.

Supports schema versions 1 (legacy) and 2 (current).  Records are normalised
to the canonical v2 schema by ``_dataset_schema.normalize_record`` before
deduplication so callers can always expect the full field set.

Usage
-----
    python deduplicate_dataset.py \
        --repo-id scikit-plots/ai-assistant-contributions \
        --output  clean_dataset.jsonl

    # Use a local pre-downloaded snapshot (faster on re-runs):
    python deduplicate_dataset.py \
        --repo-id scikit-plots/ai-assistant-contributions \
        --local-dir /tmp/ai-contributions-snapshot \
        --output clean_dataset.jsonl

Requirements
------------
    huggingface_hub>=0.23,<2
    (optional) hf_transfer for faster downloads
    (optional) _dataset_schema.py (from _hf_spaces_proxy/) for normalization

Notes
-----
* Priority rule: "contribution" beats "feedback" for the same _dedup_key.
* Retraction tombstones (action="retract") are always excluded from the
  clean output even if they win the LWW race.
* Script is idempotent: re-running produces the same output for the same
  dataset state.
* Output records are written with ``sort_keys=True``, so every record's
  keys (including nested objects) appear in a fixed alphabetical order in
  clean_dataset.jsonl.
* Progress and statistics are emitted via the module ``logging`` logger.
  INFO-level records route to stdout; WARNING and ERROR records route to
  stderr — preserving the previous ``print`` /
  ``print(..., file=sys.stderr)`` split so that callers capturing stdout
  see only the NDJSON data.
* When _dataset_schema is importable, records are normalised from v1 to v2
  schema automatically (legacy _sessionId/_page/_model fields mapped to
  conversationId/page/model; editCount/feedbackId/prevFeedbackId back-filled).
  When _dataset_schema is not importable (standalone usage), records are used
  as-is with a warning.
"""  # noqa: D205, D400

from __future__ import annotations

import argparse
import json
import logging
import sys
from pathlib import Path
from typing import Any

logger = logging.getLogger(__name__)

# Optional: import _RedactingFilter from _shared_logic when available so that
# HF token strings embedded in exception messages (e.g. snapshot_download auth
# failures) are scrubbed from CLI log output.  Safe no-op fallback for
# standalone usage where _shared_logic.py is absent.
try:
    from _shared_logic import _RedactingFilter as _REDACTING_FILTER_CLS
except ImportError:
    _REDACTING_FILTER_CLS = None  # type: ignore[assignment,misc]

# Optional: normalize records from v1 to v2 schema when _dataset_schema is
# available alongside this script (standard _hf_spaces_proxy/ deployment).
# Falls back to identity function with a warning for standalone usage.
try:
    from _dataset_schema import normalize_record as _normalize_record

    _SCHEMA_AVAILABLE = True
except ImportError:

    def _normalize_record(raw: dict) -> dict:
        return raw

    _SCHEMA_AVAILABLE = False


# Priority order: lower index = higher priority.
_SOURCE_PRIORITY: dict[str, int] = {
    "contribution": 0,
    "feedback": 1,
}
_DEFAULT_PRIORITY = 99


def _priority(record: dict) -> int:
    return _SOURCE_PRIORITY.get(record.get("_source", ""), _DEFAULT_PRIORITY)


def load_all_records(local_dir: Path) -> list[dict]:
    """Read every *.jsonl file under local_dir into a flat list.

    Parameters
    ----------
    local_dir : pathlib.Path
        Root of the locally downloaded dataset snapshot.

    Returns
    -------
    list[dict]
        All JSON-decoded records, normalised to canonical v2 schema when
        ``_dataset_schema`` is importable.  Malformed lines are skipped with
        a WARNING-level log record.
    """
    records: list[dict] = []
    for jsonl_path in sorted(local_dir.rglob("*.jsonl")):
        with jsonl_path.open(encoding="utf-8") as fh:
            for lineno, line in enumerate(fh, 1):
                line = line.strip()  # noqa: PLW2901
                if not line:
                    continue
                try:
                    raw = json.loads(line)
                except json.JSONDecodeError as exc:
                    logger.warning(
                        "Skipping malformed JSON in %s:%d: %s",
                        jsonl_path,
                        lineno,
                        exc,
                    )
                    continue
                if not isinstance(raw, dict):
                    logger.warning(
                        "%s:%d: expected JSON object, got %s -- skipped",
                        jsonl_path,
                        lineno,
                        type(raw).__name__,
                    )
                    continue
                records.append(_normalize_record(raw))
    return records


def deduplicate(records: list[dict]) -> list[dict]:
    """Deduplicate records by _dedup_key applying the priority rule.

    Parameters
    ----------
    records : list[dict]
        All raw records from both contributions/ and feedback/ folders,
        already normalised to v2 schema by load_all_records.

    Returns
    -------
    list[dict]
        One record per unique _dedup_key.  Records that have no
        _dedup_key (legacy, pre-v1.0 records) are retained unchanged.
        Retraction tombstones are excluded from the output.

    Notes
    -----
    Priority rule
        For the same _dedup_key, the record with the lowest
        _SOURCE_PRIORITY value is kept.  Ties (same source) are broken by
        server-write timestamp (_ts), keeping the most recent.  Deterministic:
        given the same input, the output is always the same.

    Retraction tombstones
        action="retract" records are still used during the LWW loop
        because their later _ts must suppress an earlier rate record
        (correct behaviour).  They are removed in the post-loop filter so they
        cannot leak into clean_dataset.jsonl.

        Degenerate case -- orphaned tombstone wins: silently discarded.
        Net effect: the original rating was explicitly retracted, so no record
        is emitted for that key -- correct for training data quality.

    feedbackId cross-source linkage (v2)
        When both a feedback/ record and a contributions/ record exist
        for the same _dedup_key, the contribution record's feedbackId
        field points directly to the feedback record's feedbackId (1-to-1
        FK).  The winning contribution record therefore carries the complete
        provenance chain without any additional join.
    """
    keyed: dict[str, dict] = {}  # _dedup_key -> winning record
    no_key: list[dict] = []  # legacy records without _dedup_key

    for rec in records:
        dk = rec.get("_dedup_key")
        if dk is None:
            no_key.append(rec)
            continue

        existing = keyed.get(dk)
        if existing is None:
            keyed[dk] = rec
            continue

        # Compare source priorities; lower = better (contribution > feedback).
        new_pri = _priority(rec)
        old_pri = _priority(existing)
        if new_pri < old_pri:
            keyed[dk] = rec
        elif new_pri == old_pri:  # noqa: SIM102
            # Same source: keep the most recently written record (_ts).
            if rec.get("_ts", 0) > existing.get("_ts", 0):
                keyed[dk] = rec

    # Post-loop: discard retraction tombstones from the winning set.
    #
    # Scenario A (normal edit): user rates +1 (_ts=100), edits (tombstone at
    # _ts=200), then rates -1 (_ts=201).  LWW selects -1.  No tombstone. OK
    #
    # Scenario B (orphaned tombstone): +1 at _ts=100, tombstone at _ts=200,
    # but follow-up -1 never reached the server.  LWW selects the tombstone.
    # Without this filter, action="retract" with ratingValue=null would corrupt
    # training.  Filter silently drops it. OK
    clean_keyed = [r for r in keyed.values() if r.get("action") != "retract"]
    return clean_keyed + no_key


def write_output(records: list[dict], output_path: Path) -> None:
    """Write records to output_path as newline-delimited JSON.

    Parameters
    ----------
    records : list[dict]
        Deduplicated records in canonical v2 schema.
    output_path : pathlib.Path
        Destination file.  Parent directories are created if absent.

    Notes
    -----
    Each record is serialised with ``sort_keys=True``, so object keys
    (at every nesting level) are written in a fixed alphabetical order.
    This keeps the output byte-for-byte reproducible across runs and
    makes line-level diffs between dataset snapshots meaningful.
    """
    output_path.parent.mkdir(parents=True, exist_ok=True)
    with output_path.open("w", encoding="utf-8") as fh:
        for rec in records:
            fh.write(json.dumps(rec, ensure_ascii=False, sort_keys=True) + "\n")


def _report_stats(records: list[dict]) -> dict[str, Any]:
    """Return summary statistics for a list of records.

    Parameters
    ----------
    records : list[dict]
        Records to summarise (raw or deduplicated).

    Returns
    -------
    dict
        Counters by source, action, schema version, and FK population.
    """
    by_source: dict[str, int] = {}
    by_action: dict[str, int] = {}
    by_schema: dict[Any, int] = {}
    with_feedback_id = 0
    with_prev_feedback = 0
    tombstones = 0

    for r in records:
        src = r.get("_source", "unknown")
        by_source[src] = by_source.get(src, 0) + 1

        act = r.get("action", "rate")
        by_action[act] = by_action.get(act, 0) + 1

        sv = r.get("schemaVersion", "?")
        by_schema[sv] = by_schema.get(sv, 0) + 1

        if r.get("feedbackId"):
            with_feedback_id += 1
        if r.get("prevFeedbackId"):
            with_prev_feedback += 1
        if act == "retract":
            tombstones += 1

    return {
        "total": len(records),
        "by_source": by_source,
        "by_action": by_action,
        "by_schema": by_schema,
        "with_feedback_id": with_feedback_id,
        "with_prev_feedback_id": with_prev_feedback,
        "tombstones": tombstones,
    }


class _MaxLevelFilter(logging.Filter):
    """Admit only log records whose level is at or below *max_level*.

    Parameters
    ----------
    max_level : int
        Maximum ``logging`` level number (inclusive) to pass through.
        Records with a higher level number are suppressed.  Pass
        ``logging.INFO`` to block WARNING and above.

    Notes
    -----
    Attached to the stdout handler inside ``_configure_logging`` so that
    WARNING / ERROR records are handled exclusively by the stderr handler
    and are not duplicated on stdout.
    """

    def __init__(self, max_level: int) -> None:
        super().__init__()
        self.max_level = max_level

    def filter(self, record: logging.LogRecord) -> bool:  # noqa: A003
        """Return ``True`` if *record.levelno* is at or below *max_level*.

        Parameters
        ----------
        record : logging.LogRecord
            Log record to evaluate.

        Returns
        -------
        bool
            ``True`` to emit the record; ``False`` to suppress it.
        """
        return record.levelno <= self.max_level


def _configure_logging() -> None:
    """Attach stdout and stderr handlers to the root logger for CLI use.

    Routes INFO-level records to stdout with a plain ``%(message)s``
    format, and WARNING / ERROR / CRITICAL records to stderr with a
    ``[%(levelname)s] %(message)s`` format.

    This preserves the stdout / stderr split that the original ``print``
    / ``print(..., file=sys.stderr)`` calls provided:

    * Callers that capture stdout (e.g. downstream JSONL pipelines) see
      only the NDJSON data, never progress lines.
    * Diagnostic warnings and errors still appear on stderr.

    The function overwrites ``logging.root.handlers`` directly, so it is
    idempotent: repeated calls replace handlers rather than stacking
    duplicates.

    Notes
    -----
    This is a CLI-only helper.  Library callers that import the domain
    functions (``load_all_records``, ``deduplicate``, …) should configure
    their own logging handlers; this function is only invoked from
    ``main()``.
    """
    plain_fmt = logging.Formatter("%(message)s")
    level_fmt = logging.Formatter("[%(levelname)s] %(message)s")

    out_handler = logging.StreamHandler(sys.stdout)
    out_handler.setFormatter(plain_fmt)
    out_handler.setLevel(logging.DEBUG)
    out_handler.addFilter(_MaxLevelFilter(logging.INFO))

    err_handler = logging.StreamHandler(sys.stderr)
    err_handler.setFormatter(level_fmt)
    err_handler.setLevel(logging.WARNING)

    # Attach defence-in-depth redaction filter when _shared_logic is available.
    # Scrubs HF token strings that huggingface_hub may embed in auth-error
    # messages before they are emitted to stderr.  No-op when absent.
    if _REDACTING_FILTER_CLS is not None:
        _rf = _REDACTING_FILTER_CLS()
        out_handler.addFilter(_rf)
        err_handler.addFilter(_rf)

    root = logging.getLogger()
    root.handlers = [out_handler, err_handler]
    root.setLevel(logging.DEBUG)


def main(argv: list[str] | None = None) -> int:
    """Run Main."""
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--repo-id",
        required=True,
        help="HuggingFace dataset repo ID, e.g. scikit-plots/ai-assistant-contributions",
    )
    parser.add_argument(
        "--output",
        default="clean_dataset.jsonl",
        help="Output path for the deduplicated NDJSON file (default: clean_dataset.jsonl)",
    )
    parser.add_argument(
        "--local-dir",
        default=None,
        help="Use a pre-downloaded local snapshot instead of downloading.",
    )
    parser.add_argument(
        "--token",
        default=None,
        help="HuggingFace read token (optional; uses cached token if absent).",
    )
    parser.add_argument(
        "--stats-only",
        action="store_true",
        help="Print dataset statistics without writing an output file.",
    )
    args = parser.parse_args(argv)
    _configure_logging()

    if not _SCHEMA_AVAILABLE:
        logger.warning(
            "_dataset_schema.py not found on sys.path.  Records will not be "
            "normalised from v1 to v2 schema.  Copy _dataset_schema.py from "
            "_hf_spaces_proxy/ to the same directory as this script for full "
            "schema normalisation.",
        )

    local_dir: Path
    if args.local_dir:
        local_dir = Path(args.local_dir)
    else:
        try:
            from huggingface_hub import snapshot_download  # noqa: PLC0415
        except ImportError:
            logger.error(
                "huggingface_hub is not installed.  "
                "Run: pip install 'huggingface_hub>=0.23,<2'",
            )
            return 1
        logger.info("Downloading %s ...", args.repo_id)
        try:
            local_dir = Path(
                snapshot_download(
                    repo_id=args.repo_id,
                    repo_type="dataset",
                    token=args.token,
                )
            )
        except Exception as exc:  # noqa: BLE001
            logger.error(
                "Failed to download %s: %s\n"
                "Hint: pass --token <HF_READ_TOKEN> or set HF_TOKEN in your environment.",
                args.repo_id,
                exc,
            )
            return 1

    logger.info("Reading records from %s ...", local_dir)
    all_records = load_all_records(local_dir)

    raw_stats = _report_stats(all_records)
    logger.info("  %d total records read", raw_stats["total"])
    for src, cnt in sorted(raw_stats["by_source"].items()):
        logger.info("    %s: %d", src, cnt)
    for act, cnt in sorted(raw_stats["by_action"].items()):
        logger.info("    action=%r: %d", act, cnt)
    for sv, cnt in raw_stats["by_schema"].items():
        logger.info("    schemaVersion=%s: %d", sv, cnt)
    logger.info("  feedbackId populated:      %d", raw_stats["with_feedback_id"])
    logger.info("  prevFeedbackId populated:  %d", raw_stats["with_prev_feedback_id"])
    if raw_stats["tombstones"]:
        logger.info(
            "  %d retraction tombstone(s) in raw data "
            "(always excluded from clean output)",
            raw_stats["tombstones"],
        )

    if args.stats_only:
        return 0

    clean = deduplicate(all_records)
    duplicates_removed = raw_stats["total"] - raw_stats["tombstones"] - len(clean)
    logger.info("  %d duplicate(s) removed (priority rule applied)", duplicates_removed)
    logger.info("  %d unique records retained", len(clean))

    output_path = Path(args.output)
    write_output(clean, output_path)
    logger.info("Clean dataset written to %s", output_path)
    return 0


if __name__ == "__main__":
    sys.exit(main())