glitchlings 0.4.4__cp313-cp313-win_amd64.whl

glitchlings/dlc/prime.py

@@ -0,0 +1,254 @@
+"""Integration helpers for the optional verifiers prime DLC."""
+
+from __future__ import annotations
+
+from collections.abc import Iterable, Sequence
+from enum import Enum
+from typing import Any, Callable, Protocol, cast
+
+from ..compat import require_datasets, require_jellyfish, require_verifiers
+from ..util.adapters import coerce_gaggle
+from ..zoo import Gaggle, Glitchling, Mim1c, Typogre
+from ._shared import resolve_columns as _resolve_columns_shared
+from ._shared import resolve_environment as _resolve_environment_shared
+
+
+class VerifierEnvironment(Protocol):
+    """Minimal interface for verifiers environments."""
+
+    dataset: Any
+
+
+class VerifierSingleTurnEnv(Protocol):
+    """Minimal interface for single-turn verifier environments."""
+
+    dataset: Any
+    rubric: Any
+
+
+vf = require_verifiers("verifiers is not installed; install glitchlings[prime]")
+_jellyfish = require_jellyfish("jellyfish is not installed; install glitchlings[prime]")
+damerau_levenshtein_distance = _jellyfish.damerau_levenshtein_distance
+
+try:
+    from .huggingface import Dataset as _HuggingFaceDataset
+except ModuleNotFoundError:  # pragma: no cover - optional dependency
+    _HuggingFaceDataset = None
+else:
+    if _HuggingFaceDataset is None:  # pragma: no cover - optional dependency
+        _HuggingFaceDataset = None
+
+Dataset: type[Any]
+if _HuggingFaceDataset is None:
+    Dataset = object
+else:
+    Dataset = _HuggingFaceDataset
+
+
+def _resolve_environment(env: str | VerifierEnvironment) -> VerifierEnvironment:
+    """Return a fully-instantiated verifier environment."""
+    resolved = _resolve_environment_shared(
+        env,
+        loader=vf.load_environment,
+        environment_type=cast(type[Any], vf.Environment),
+    )
+    return cast(VerifierEnvironment, resolved)
+
+
+def _resolve_columns(dataset: Any, columns: Sequence[str] | None) -> list[str]:
+    """Identify which dataset columns should be corrupted."""
+    return _resolve_columns_shared(dataset, columns)
+
+
+class Difficulty(Enum):
+    """Difficulty levels for tutorial environments."""
+
+    Easy = 0.25
+    Normal = 1.0
+    Hard = 1.75
+    Extreme = 3
+    Impossible = 9
+
+
+def tutorial_level(
+    env: VerifierEnvironment | str,
+    seed: int = 151,
+    difficulty: Difficulty = Difficulty.Normal,
+) -> VerifierEnvironment:
+    """Create a low-corruption environment using tuned defaults."""
+    tuned_mim1c = Mim1c(rate=0.01 * difficulty.value)
+    tuned_typogre = Typogre(rate=0.025 * difficulty.value)
+
+    return load_environment(
+        env,
+        glitchlings=[tuned_mim1c, tuned_typogre],
+        seed=seed,
+    )
+
+
+def load_environment(
+    env: str | VerifierEnvironment,
+    glitchlings: Iterable[str | Glitchling] | Glitchling | str | Gaggle | None = None,
+    *,
+    seed: int = 151,
+    columns: Sequence[str] | None = None,
+) -> VerifierEnvironment:
+    """Load an environment and optionally corrupt it with glitchlings."""
+    environment = _resolve_environment(env)
+
+    if glitchlings is None:
+        return environment
+
+    gaggle = coerce_gaggle(glitchlings, seed=seed)
+
+    dataset = environment.dataset
+    corrupt_columns = _resolve_columns(dataset, columns)
+    environment.dataset = gaggle.corrupt_dataset(dataset, corrupt_columns)
+    return environment
+
+
+def _as_gaggle(
+    glitchlings: Iterable[str | Glitchling] | Glitchling | str | Gaggle,
+    *,
+    seed: int,
+) -> Gaggle:
+    """Coerce any supported glitchling specification into a :class:`Gaggle`."""
+    return coerce_gaggle(glitchlings, seed=seed)
+
+
+def _extract_completion_text(completion: Any) -> str:
+    """Normalise a completion payload into a plain string."""
+    if isinstance(completion, str):
+        return completion
+
+    if isinstance(completion, list) and completion:
+        first = completion[0]
+        if isinstance(first, dict) and "content" in first:
+            return str(first["content"])
+        return str(first)
+
+    return str(completion)
+
+
+def symmetric_damerau_levenshtein_similarity(
+    _: Any,
+    completion: Any,
+    answer: str,
+) -> float:
+    """Return ``1 - (distance / max_len)`` using Damerau-Levenshtein distance."""
+    completion_text = _extract_completion_text(completion)
+    target = answer or ""
+    denominator = max(len(completion_text), len(target), 1)
+    distance = cast(int, damerau_levenshtein_distance(completion_text, target))
+    score = 1.0 - (distance / denominator)
+    return max(0.0, min(1.0, score))
+
+
+DEFAULT_CLEANUP_INSTRUCTIONS = (
+    "You are a meticulous copy editor. Restore the provided text to its original form."
+)
+
+
+def echo_chamber(
+    dataset_id: str,
+    column: str,
+    glitchlings: Iterable[str | Glitchling] | Glitchling | str | Gaggle,
+    *,
+    seed: int = 151,
+    instructions: str = DEFAULT_CLEANUP_INSTRUCTIONS,
+    reward_function: Callable[..., float] | None = None,
+    split: str | None = None,
+    **load_dataset_kwargs: Any,
+) -> VerifierSingleTurnEnv:
+    """Create an Echo Chamber Prime environment from a Hugging Face dataset column.
+
+    Args:
+        dataset_id: Identifier of the Hugging Face dataset to load.
+        column: Name of the column whose text should be glitched.
+        glitchlings: Glitchling specifiers that will corrupt the prompts.
+        seed: RNG seed forwarded to :func:`glitchlings.util.adapters.coerce_gaggle`.
+        instructions: System instructions supplied to the environment prompts.
+        reward_function: Optional callable used to score completions. Defaults to
+            :func:`symmetric_damerau_levenshtein_similarity` when omitted.
+        split: Optional dataset split to load.
+        **load_dataset_kwargs: Extra keyword arguments forwarded to
+            :func:`datasets.load_dataset`.
+
+    """
+    datasets_module = require_datasets("datasets is required to build an echo chamber")
+    load_dataset = getattr(datasets_module, "load_dataset", None)
+    if load_dataset is None:  # pragma: no cover - defensive
+        message = "datasets is required to build an echo chamber"
+        raise ModuleNotFoundError(message)
+
+    dataset_dict_cls = getattr(datasets_module, "DatasetDict", dict)
+
+    hf_dataset: Any
+    if split is None:
+        hf_dataset = load_dataset(dataset_id, **load_dataset_kwargs)
+        if isinstance(hf_dataset, dataset_dict_cls):
+            try:
+                hf_dataset = next(iter(hf_dataset.values()))
+            except StopIteration as exc:  # pragma: no cover - defensive
+                raise ValueError("The specified dataset does not contain any splits") from exc
+    else:
+        hf_dataset = load_dataset(dataset_id, split=split, **load_dataset_kwargs)
+
+    if isinstance(hf_dataset, dataset_dict_cls):
+        raise ValueError("Specify which split to use when the dataset loads as a DatasetDict.")
+
+    filtered_dataset = hf_dataset.filter(
+        lambda row: row.get(column) is not None,
+        load_from_cache_file=False,
+    )
+
+    source_column_names = list(filtered_dataset.column_names)
+
+    def _build_prompt(row: dict[str, Any]) -> dict[str, Any]:
+        text = str(row[column])
+        prompt = [
+            {"role": "system", "content": instructions},
+            {"role": "user", "content": f"Corrupted text:\n{text}"},
+        ]
+        return {"prompt": prompt, "answer": text}
+
+    base_dataset = filtered_dataset.map(
+        _build_prompt,
+        remove_columns=source_column_names,
+        load_from_cache_file=False,
+    )
+
+    try:
+        dataset_length = len(base_dataset)
+    except TypeError:
+        preview_rows: list[dict[str, Any]]
+        take_fn = getattr(base_dataset, "take", None)
+        if callable(take_fn):
+            preview_rows = list(take_fn(1))
+        else:
+            iterator = iter(base_dataset)
+            try:
+                first_row = next(iterator)
+            except StopIteration:
+                preview_rows = []
+            else:
+                preview_rows = [first_row]
+        if not preview_rows:
+            raise ValueError(
+                f"Column '{column}' did not yield any textual entries in dataset '{dataset_id}'."
+            )
+    else:
+        if dataset_length == 0:
+            raise ValueError(
+                f"Column '{column}' did not yield any textual entries in dataset '{dataset_id}'."
+            )
+
+    gaggle = _as_gaggle(glitchlings, seed=seed)
+    glitched_dataset = gaggle.corrupt_dataset(base_dataset, ["prompt"])
+
+    rubric_func = reward_function or symmetric_damerau_levenshtein_similarity
+    rubric = vf.Rubric(funcs=[rubric_func], weights=[1.0])
+    return cast(
+        VerifierSingleTurnEnv,
+        vf.SingleTurnEnv(dataset=glitched_dataset, rubric=rubric),
+    )

glitchlings/dlc/pytorch.py

@@ -0,0 +1,166 @@
+"""Integration helpers for PyTorch data loaders."""
+
+from __future__ import annotations
+
+from collections.abc import Iterable, Iterator, Mapping, MutableMapping, Sequence
+from typing import Any, cast
+
+from ..compat import get_torch_dataloader, require_torch
+from ..compat import torch as _torch_dependency
+from ..util.adapters import coerce_gaggle
+from ..zoo import Gaggle, Glitchling
+from ._shared import corrupt_text_value, is_textual_candidate, normalise_column_spec
+
+
+def _apply_to_batch(batch: Any, targets: list[str | int] | None, gaggle: Gaggle) -> Any:
+    """Return ``batch`` with glitchlings applied to the specified ``targets``."""
+    if targets is None:
+        return corrupt_text_value(batch, gaggle)
+
+    if isinstance(batch, Mapping):
+        mutated = cast(MutableMapping[str, Any], dict(batch))
+        for key in targets:
+            if not isinstance(key, str):
+                raise TypeError("Mapping batches require string column names")
+            if key not in mutated:
+                raise ValueError(f"Column '{key}' not found in DataLoader batch")
+            mutated[key] = corrupt_text_value(mutated[key], gaggle)
+        return mutated
+
+    if isinstance(batch, Sequence) and not isinstance(batch, (bytes, bytearray, str)):
+        mutated_sequence = list(batch)
+        for index in targets:
+            if not isinstance(index, int):
+                raise TypeError("Sequence batches require integer column indices")
+            try:
+                mutated_sequence[index] = corrupt_text_value(mutated_sequence[index], gaggle)
+            except IndexError as exc:  # pragma: no cover - defensive
+                raise IndexError("Column index out of range for DataLoader batch") from exc
+        if isinstance(batch, tuple):
+            return tuple(mutated_sequence)
+        return mutated_sequence
+
+    raise TypeError("Unsupported DataLoader batch type for glitching")
+
+
+def _infer_targets(batch: Any) -> list[str | int] | None:
+    """Infer which fields should be glitched from a representative ``batch``."""
+    if isinstance(batch, Mapping):
+        inferred = [key for key, value in batch.items() if is_textual_candidate(value)]
+        if inferred:
+            return inferred
+        raise ValueError("Unable to infer which mapping columns contain text")
+
+    if isinstance(batch, Sequence) and not isinstance(batch, (bytes, bytearray, str)):
+        inferred_indices: list[str | int] = [
+            idx for idx, value in enumerate(batch) if is_textual_candidate(value)
+        ]
+        if inferred_indices:
+            return inferred_indices
+        raise ValueError("Unable to infer which sequence indices contain text")
+
+    if is_textual_candidate(batch):
+        return None
+
+    raise TypeError("Unsupported DataLoader batch type for glitching")
+
+
+class _GlitchedDataLoader(Iterable[Any]):
+    """Wrapper that applies glitchlings lazily to each batch from a data loader."""
+
+    def __init__(
+        self,
+        dataloader: Any,
+        gaggle: Gaggle,
+        *,
+        columns: list[str | int] | None,
+    ) -> None:
+        self._dataloader = dataloader
+        self._gaggle = gaggle
+        self._explicit_columns = columns
+        self._inferred_columns: list[str | int] | None | _Sentinel = _UNINITIALISED
+
+    def __iter__(self) -> Iterator[Any]:
+        # Reset all glitchling RNGs before each fresh pass for determinism.
+        self._gaggle.sort_glitchlings()
+        for batch in self._dataloader:
+            targets = self._resolve_columns(batch)
+            yield _apply_to_batch(batch, targets, self._gaggle)
+
+    def __len__(self) -> int:
+        return len(self._dataloader)
+
+    def __getattr__(self, attribute: str) -> Any:
+        return getattr(self._dataloader, attribute)
+
+    def _resolve_columns(self, batch: Any) -> list[str | int] | None:
+        if self._explicit_columns is not None:
+            return self._explicit_columns
+
+        if self._inferred_columns is _UNINITIALISED:
+            self._inferred_columns = _infer_targets(batch)
+
+        return cast(list[str | int] | None, self._inferred_columns)
+
+
+class _Sentinel:
+    """Sentinel type for deferred column inference."""
+
+
+_UNINITIALISED = _Sentinel()
+
+
+def _ensure_dataloader_class() -> type[Any]:
+    """Return :class:`torch.utils.data.DataLoader` patched with ``.glitch``."""
+    dataloader_cls = get_torch_dataloader()
+    if dataloader_cls is None:
+        require_torch("torch is not installed; install glitchlings[torch]")
+        dataloader_cls = get_torch_dataloader()
+        if dataloader_cls is None:  # pragma: no cover - defensive
+            message = "torch.utils.data.DataLoader is not available"
+            error = _torch_dependency.error
+            if error is not None:
+                raise ModuleNotFoundError(message) from error
+            raise ModuleNotFoundError(message)
+
+    if getattr(dataloader_cls, "glitch", None) is None:
+
+        def glitch(
+            self: Any,
+            glitchlings: Iterable[str | Glitchling] | Glitchling | str | Gaggle,
+            *,
+            columns: str | int | Sequence[str | int] | None = None,
+            seed: int = 151,
+        ) -> _GlitchedDataLoader:
+            """Return a lazily glitched view of the loader's batches."""
+            gaggle = coerce_gaggle(glitchlings, seed=seed)
+            normalised = normalise_column_spec(columns)
+            return _GlitchedDataLoader(self, gaggle, columns=normalised)
+
+        setattr(dataloader_cls, "glitch", glitch)
+
+    return cast(type[Any], dataloader_cls)
+
+
+def _optional_dataloader_class() -> type[Any] | None:
+    """Return the PyTorch :class:`~torch.utils.data.DataLoader` when importable."""
+    dataloader_cls = get_torch_dataloader()
+    if dataloader_cls is None:
+        return None
+    return cast(type[Any], dataloader_cls)
+
+
+def install() -> None:
+    """Monkeypatch PyTorch's :class:`~torch.utils.data.DataLoader` with ``.glitch``."""
+    _ensure_dataloader_class()
+
+
+DataLoader: type[Any] | None
+_DataLoaderAlias = _optional_dataloader_class()
+if _DataLoaderAlias is not None:
+    DataLoader = _ensure_dataloader_class()
+else:  # pragma: no cover - torch is an optional dependency
+    DataLoader = None
+
+
+__all__ = ["DataLoader", "install"]

glitchlings/dlc/pytorch_lightning.py

@@ -0,0 +1,215 @@
+"""Integration helpers for PyTorch Lightning data modules."""
+
+from __future__ import annotations
+
+from collections.abc import Iterable, Mapping, Sequence
+from typing import Any, cast
+
+from ..compat import get_pytorch_lightning_datamodule, require_pytorch_lightning
+from ..util.adapters import coerce_gaggle
+from ..zoo import Gaggle, Glitchling
+from ._shared import corrupt_text_value, normalise_column_spec
+
+
+def _glitch_batch(batch: Any, columns: list[str], gaggle: Gaggle) -> Any:
+    """Apply glitchlings to the configured batch columns."""
+    if not isinstance(batch, Mapping):
+        return batch
+
+    if hasattr(batch, "copy"):
+        mutated = batch.copy()
+    else:
+        mutated = dict(batch)
+
+    missing = [column for column in columns if column not in mutated]
+    if missing:
+        missing_str = ", ".join(sorted(missing))
+        raise ValueError(f"Columns not found in batch: {missing_str}")
+
+    for column in columns:
+        mutated[column] = corrupt_text_value(mutated[column], gaggle)
+
+    return mutated
+
+
+def _wrap_dataloader(dataloader: Any, columns: list[str], gaggle: Gaggle) -> Any:
+    """Wrap a dataloader so yielded batches are corrupted lazily."""
+    if dataloader is None:
+        return None
+
+    if isinstance(dataloader, Mapping):
+        mapping_type = cast(type[Any], dataloader.__class__)
+        return mapping_type(
+            {
+                key: _wrap_dataloader(value, columns, gaggle)
+                for key, value in dataloader.items()
+            }
+        )
+
+    if isinstance(dataloader, list):
+        return [_wrap_dataloader(value, columns, gaggle) for value in dataloader]
+
+    if isinstance(dataloader, tuple):
+        return tuple(_wrap_dataloader(value, columns, gaggle) for value in dataloader)
+
+    if isinstance(dataloader, Sequence) and not isinstance(dataloader, (str, bytes, bytearray)):
+        sequence_type = cast(type[Any], dataloader.__class__)
+        return sequence_type(
+            _wrap_dataloader(value, columns, gaggle) for value in dataloader
+        )
+
+    return _GlitchedDataLoader(dataloader, columns, gaggle)
+
+
+class _GlitchedDataLoader:
+    """Proxy dataloader that glitches batches produced by the wrapped loader."""
+
+    def __init__(self, dataloader: Any, columns: list[str], gaggle: Gaggle) -> None:
+        self._dataloader = dataloader
+        self._columns = columns
+        self._gaggle = gaggle
+
+    def __iter__(self) -> Any:
+        for batch in self._dataloader:
+            yield _glitch_batch(batch, self._columns, self._gaggle)
+
+    def __len__(self) -> int:
+        return len(self._dataloader)
+
+    def __getattr__(self, attribute: str) -> Any:
+        return getattr(self._dataloader, attribute)
+
+
+def _glitch_datamodule(
+    datamodule: Any,
+    glitchlings: Glitchling | Gaggle | str | Iterable[str | Glitchling],
+    column: str | Sequence[str],
+    *,
+    seed: int = 151,
+) -> Any:
+    """Return a proxy that applies glitchlings to batches from the datamodule."""
+
+    columns = normalise_column_spec(column)
+    if columns is None:  # pragma: no cover - defensive
+        raise ValueError("At least one column must be specified")
+    # Lightning datamodules only support string column names (mapping keys)
+    columns_str = cast(list[str], columns)
+    gaggle = coerce_gaggle(glitchlings, seed=seed)
+    return _GlitchedLightningDataModule(datamodule, columns_str, gaggle)
+
+
+class _GlitchedLightningDataModule:
+    """Proxy wrapper around a LightningDataModule applying glitchlings to batches."""
+
+    def __init__(self, base: Any, columns: list[str], gaggle: Gaggle) -> None:
+        object.__setattr__(self, "_glitch_base", base)
+        object.__setattr__(self, "_glitch_columns", columns)
+        object.__setattr__(self, "_glitch_gaggle", gaggle)
+
+    def __getattr__(self, attribute: str) -> Any:
+        return getattr(self._glitch_base, attribute)
+
+    def __setattr__(self, attribute: str, value: Any) -> None:
+        if attribute.startswith("_glitch_"):
+            object.__setattr__(self, attribute, value)
+        else:
+            setattr(self._glitch_base, attribute, value)
+
+    def __delattr__(self, attribute: str) -> None:
+        if attribute.startswith("_glitch_"):
+            object.__delattr__(self, attribute)
+        else:
+            delattr(self._glitch_base, attribute)
+
+    def __dir__(self) -> list[str]:
+        return sorted(set(dir(self.__class__)) | set(dir(self._glitch_base)))
+
+    # LightningDataModule API -------------------------------------------------
+    def prepare_data(self, *args: Any, **kwargs: Any) -> Any:
+        return self._glitch_base.prepare_data(*args, **kwargs)
+
+    def setup(self, *args: Any, **kwargs: Any) -> Any:
+        return self._glitch_base.setup(*args, **kwargs)
+
+    def teardown(self, *args: Any, **kwargs: Any) -> Any:
+        return self._glitch_base.teardown(*args, **kwargs)
+
+    def state_dict(self) -> Mapping[str, Any]:
+        state = self._glitch_base.state_dict()
+        return cast(Mapping[str, Any], state)
+
+    def load_state_dict(self, state_dict: Mapping[str, Any]) -> None:
+        self._glitch_base.load_state_dict(state_dict)
+
+    def transfer_batch_to_device(self, batch: Any, device: Any, dataloader_idx: int) -> Any:
+        return self._glitch_base.transfer_batch_to_device(batch, device, dataloader_idx)
+
+    def on_before_batch_transfer(self, batch: Any, dataloader_idx: int) -> Any:
+        return self._glitch_base.on_before_batch_transfer(batch, dataloader_idx)
+
+    def on_after_batch_transfer(self, batch: Any, dataloader_idx: int) -> Any:
+        return self._glitch_base.on_after_batch_transfer(batch, dataloader_idx)
+
+    def train_dataloader(self, *args: Any, **kwargs: Any) -> Any:
+        loader = self._glitch_base.train_dataloader(*args, **kwargs)
+        return _wrap_dataloader(loader, self._glitch_columns, self._glitch_gaggle)
+
+    def val_dataloader(self, *args: Any, **kwargs: Any) -> Any:
+        loader = self._glitch_base.val_dataloader(*args, **kwargs)
+        return _wrap_dataloader(loader, self._glitch_columns, self._glitch_gaggle)
+
+    def test_dataloader(self, *args: Any, **kwargs: Any) -> Any:
+        loader = self._glitch_base.test_dataloader(*args, **kwargs)
+        return _wrap_dataloader(loader, self._glitch_columns, self._glitch_gaggle)
+
+    def predict_dataloader(self, *args: Any, **kwargs: Any) -> Any:
+        loader = self._glitch_base.predict_dataloader(*args, **kwargs)
+        return _wrap_dataloader(loader, self._glitch_columns, self._glitch_gaggle)
+
+
+def _ensure_datamodule_class() -> Any:
+    """Return the Lightning ``LightningDataModule`` patched with ``.glitch``."""
+
+    datamodule_cls = get_pytorch_lightning_datamodule()
+    if datamodule_cls is None:  # pragma: no cover - dependency is optional
+        module = require_pytorch_lightning("pytorch_lightning is not installed")
+        datamodule_cls = getattr(module, "LightningDataModule", None)
+        if datamodule_cls is None:
+            raise ModuleNotFoundError("pytorch_lightning is not installed")
+
+    if getattr(datamodule_cls, "glitch", None) is None:
+
+        def glitch(
+            self: Any,
+            glitchlings: Glitchling | Gaggle | str | Iterable[str | Glitchling],
+            *,
+            column: str | Sequence[str],
+            seed: int = 151,
+            **_: Any,
+        ) -> Any:
+            return _glitch_datamodule(self, glitchlings, column, seed=seed)
+
+        setattr(datamodule_cls, "glitch", glitch)
+
+    if not issubclass(_GlitchedLightningDataModule, datamodule_cls):
+        _GlitchedLightningDataModule.__bases__ = (datamodule_cls,)
+
+    return datamodule_cls
+
+
+def install() -> None:
+    """Monkeypatch ``LightningDataModule`` with ``.glitch``."""
+
+    _ensure_datamodule_class()
+
+
+LightningDataModule: type[Any] | None
+_LightningDataModuleAlias = get_pytorch_lightning_datamodule()
+if _LightningDataModuleAlias is not None:
+    LightningDataModule = _ensure_datamodule_class()
+else:  # pragma: no cover - optional dependency
+    LightningDataModule = None
+
+
+__all__ = ["LightningDataModule", "install"]
+