glitchlings 0.4.2__cp311-cp311-macosx_11_0_universal2.whl → 0.4.3__cp311-cp311-macosx_11_0_universal2.whl

glitchlings/dlc/pytorch_lightning.py

@@ -0,0 +1,233 @@
+"""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 ..zoo.core import _is_transcript
+
+
+def _normalise_columns(column: str | Sequence[str]) -> list[str]:
+    """Normalise a column specification to a list."""
+    if isinstance(column, str):
+        return [column]
+
+    normalised = list(column)
+    if not normalised:
+        raise ValueError("At least one column must be specified")
+    return normalised
+
+
+def _glitch_value(value: Any, gaggle: Gaggle) -> Any:
+    """Apply glitchlings to a value when it contains textual content."""
+    if isinstance(value, str) or _is_transcript(value, allow_empty=False, require_all_content=True):
+        return gaggle.corrupt(value)
+
+    if isinstance(value, Sequence) and value and all(isinstance(item, str) for item in value):
+        return [gaggle.corrupt(item) for item in value]
+
+    return value
+
+
+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] = _glitch_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_columns(column)
+    gaggle = coerce_gaggle(glitchlings, seed=seed)
+    return _GlitchedLightningDataModule(datamodule, columns, 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"]
+

glitchlings/lexicon/__init__.py

@@ -91,7 +91,6 @@ class LexiconBackend(Lexicon):
         """Persist the backend cache to ``path`` and return the destination."""
 
 
-from .graph import GraphLexicon  # noqa: E402
 from .metrics import (  # noqa: E402
     compare_lexicons,
     coverage_ratio,
@@ -100,10 +99,13 @@ from .metrics import (  # noqa: E402
 )
 from .vector import VectorLexicon, build_vector_cache  # noqa: E402
 
+_WordNetLexicon: type[LexiconBackend] | None
 try:  # pragma: no cover - optional dependency
-    from .wordnet import WordNetLexicon
+    from .wordnet import WordNetLexicon as _WordNetLexicon
 except Exception:  # pragma: no cover - triggered when nltk unavailable
-    WordNetLexicon = None  # type: ignore[assignment]
+    _WordNetLexicon = None
+
+WordNetLexicon: type[LexiconBackend] | None = _WordNetLexicon
 
 
 _BACKEND_FACTORIES: dict[str, Callable[[int | None], Lexicon | None]] = {}
@@ -135,16 +137,6 @@ def _vector_backend(seed: int | None) -> Lexicon | None:
     return VectorLexicon(cache_path=cache_path, seed=seed)
 
 
-def _graph_backend(seed: int | None) -> Lexicon | None:
-    config = get_config()
-    cache_path = config.lexicon.graph_cache
-    if cache_path is None:
-        return None
-    if not cache_path.exists():
-        return None
-    return GraphLexicon(cache_path=cache_path, seed=seed)
-
-
 def _wordnet_backend(seed: int | None) -> Lexicon | None:  # pragma: no cover - optional
     if WordNetLexicon is None:
         return None
@@ -156,7 +148,6 @@ def _wordnet_backend(seed: int | None) -> Lexicon | None:  # pragma: no cover -
 
 
 register_backend("vector", _vector_backend)
-register_backend("graph", _graph_backend)
 register_backend("wordnet", _wordnet_backend)
 
 
@@ -184,7 +175,6 @@ __all__ = [
     "Lexicon",
     "LexiconBackend",
     "VectorLexicon",
-    "GraphLexicon",
     "WordNetLexicon",
     "build_vector_cache",
     "compare_lexicons",

glitchlings/lexicon/_cache.py

@@ -6,7 +6,7 @@ import json
 from dataclasses import dataclass
 from hashlib import blake2s
 from pathlib import Path
-from typing import Mapping, Sequence
+from typing import Mapping, Sequence, cast
 
 CacheEntries = dict[str, list[str]]
 
@@ -19,7 +19,7 @@ class CacheSnapshot:
     checksum: str | None = None
 
 
-def _normalise_entries(payload: Mapping[str, Sequence[str]]) -> CacheEntries:
+def _normalise_entries(payload: Mapping[str, object]) -> CacheEntries:
     """Convert raw cache payloads into canonical mapping form."""
     entries: CacheEntries = {}
     for key, values in payload.items():
@@ -49,27 +49,31 @@ def load_cache(path: Path) -> CacheSnapshot:
         return CacheSnapshot(entries={}, checksum=None)
 
     with path.open("r", encoding="utf8") as handle:
-        payload = json.load(handle)
+        payload_obj = json.load(handle)
 
     checksum: str | None = None
-    entries_payload: Mapping[str, Sequence[str]]
+    entries_payload: Mapping[str, object]
 
-    if isinstance(payload, Mapping) and "__meta__" in payload and "entries" in payload:
-        meta = payload["__meta__"]
-        entries_payload = payload["entries"]  # type: ignore[assignment]
-        if not isinstance(entries_payload, Mapping):
+    if not isinstance(payload_obj, Mapping):
+        raise RuntimeError("Synonym cache payload must be a mapping of strings to lists.")
+
+    payload = cast(Mapping[str, object], payload_obj)
+
+    if "__meta__" in payload and "entries" in payload:
+        meta_obj = payload["__meta__"]
+        entries_obj = payload["entries"]
+        if not isinstance(entries_obj, Mapping):
             raise RuntimeError("Synonym cache entries must be stored as a mapping.")
-        if isinstance(meta, Mapping):
-            raw_checksum = meta.get("checksum")
+        entries_payload = cast(Mapping[str, object], entries_obj)
+        if isinstance(meta_obj, Mapping):
+            raw_checksum = meta_obj.get("checksum")
             if raw_checksum is not None and not isinstance(raw_checksum, str):
                 raise RuntimeError("Synonym cache checksum must be a string when provided.")
-            checksum = raw_checksum
+            checksum = raw_checksum if isinstance(raw_checksum, str) else None
         else:
             raise RuntimeError("Synonym cache metadata must be a mapping.")
-    elif isinstance(payload, Mapping):
-        entries_payload = payload  # legacy format without metadata
     else:
-        raise RuntimeError("Synonym cache payload must be a mapping of strings to lists.")
+        entries_payload = payload  # legacy format without metadata
 
     entries = _normalise_entries(entries_payload)
     if checksum is not None:
@@ -84,7 +88,9 @@ def load_cache(path: Path) -> CacheSnapshot:
 
 def write_cache(path: Path, entries: Mapping[str, Sequence[str]]) -> CacheSnapshot:
     """Persist ``entries`` to ``path`` with checksum metadata."""
-    serialisable = {key: list(values) for key, values in sorted(entries.items())}
+    serialisable: CacheEntries = {
+        key: list(values) for key, values in sorted(entries.items())
+    }
     checksum = compute_checksum(serialisable)
     payload = {
         "__meta__": {

glitchlings/lexicon/data/default_vector_cache.json

@@ -1,16 +1,82 @@
 {
-  "sing": ["croon", "warble", "chant", "serenade"],
-  "happy": ["cheerful", "joyful", "contented", "gleeful"],
-  "songs": ["tunes", "melodies", "ballads", "airs"],
-  "quickly": ["rapidly", "swiftly", "speedily", "promptly"],
-  "text": ["passage", "excerpt", "phrase", "content"],
-  "words": ["terms", "phrases", "lexemes", "expressions"],
-  "alpha": ["beta", "gamma", "delta"],
-  "beta": ["alpha", "gamma", "delta"],
-  "gamma": ["alpha", "beta", "delta"],
-  "delta": ["alpha", "beta", "gamma"],
-  "they": ["these people", "those folks", "those individuals"],
-  "quick": ["rapid", "swift", "brisk", "prompt"],
-  "fast": ["rapid", "swift", "quick", "speedy"],
-  "slow": ["sluggish", "lethargic", "unhurried", "deliberate"]
+  "alpha": [
+    "beta",
+    "gamma",
+    "delta"
+  ],
+  "beta": [
+    "alpha",
+    "gamma",
+    "delta"
+  ],
+  "delta": [
+    "alpha",
+    "beta",
+    "gamma"
+  ],
+  "fast": [
+    "rapid",
+    "swift",
+    "speedy",
+    "brisk"
+  ],
+  "gamma": [
+    "alpha",
+    "beta",
+    "delta"
+  ],
+  "happy": [
+    "glad",
+    "joyful",
+    "content",
+    "upbeat"
+  ],
+  "quick": [
+    "swift",
+    "rapid",
+    "speedy",
+    "nimble"
+  ],
+  "quickly": [
+    "swiftly",
+    "rapidly",
+    "promptly",
+    "speedily"
+  ],
+  "sing": [
+    "croon",
+    "serenade",
+    "vocalize",
+    "perform"
+  ],
+  "slow": [
+    "sluggish",
+    "leisurely",
+    "unhurried",
+    "gradual"
+  ],
+  "songs": [
+    "tracks",
+    "melodies",
+    "ballads",
+    "tunes"
+  ],
+  "text": [
+    "passage",
+    "copy",
+    "script",
+    "narrative"
+  ],
+  "they": [
+    "those people",
+    "those individuals",
+    "the group",
+    "those folks"
+  ],
+  "words": [
+    "terms",
+    "phrases",
+    "lexicon",
+    "vocabulary"
+  ]
 }

glitchlings/lexicon/vector.py

@@ -4,6 +4,7 @@ from __future__ import annotations
 
 import argparse
 import importlib
+import importlib.util
 import json
 import math
 import sys
@@ -188,6 +189,58 @@ def _load_spacy_language(model_name: str) -> Any:
     return spacy_module.load(model_name)
 
 
+def _load_sentence_transformer(model_name: str) -> Any:
+    """Return a ``SentenceTransformer`` instance for ``model_name``."""
+
+    if importlib.util.find_spec("sentence_transformers") is None:
+        raise RuntimeError(
+            "sentence-transformers is required for this source; install the 'st' extra."
+        )
+
+    module = importlib.import_module("sentence_transformers")
+    try:
+        model_cls = getattr(module, "SentenceTransformer")
+    except AttributeError as exc:  # pragma: no cover - defensive
+        raise RuntimeError("sentence-transformers does not expose SentenceTransformer") from exc
+
+    return model_cls(model_name)
+
+
+def _build_sentence_transformer_embeddings(
+    model_name: str, tokens: Sequence[str]
+) -> Mapping[str, Sequence[float]]:
+    """Return embeddings for ``tokens`` using ``model_name``."""
+
+    if not tokens:
+        return {}
+
+    model = _load_sentence_transformer(model_name)
+
+    unique_tokens: list[str] = []
+    seen: set[str] = set()
+    for token in tokens:
+        normalized = token.strip()
+        if not normalized or normalized in seen:
+            continue
+        unique_tokens.append(normalized)
+        seen.add(normalized)
+
+    if not unique_tokens:
+        return {}
+
+    embeddings = model.encode(
+        unique_tokens,
+        batch_size=64,
+        normalize_embeddings=True,
+        convert_to_numpy=True,
+    )
+
+    return {
+        token: [float(value) for value in vector]
+        for token, vector in zip(unique_tokens, embeddings, strict=True)
+    }
+
+
 def _resolve_source(source: Any | None) -> _Adapter | None:
     """Return an adapter instance for ``source`` if possible."""
     if source is None:
@@ -248,6 +301,7 @@ class VectorLexicon(LexiconBackend):
         case_sensitive: bool = False,
         seed: int | None = None,
     ) -> None:
+        """Initialise the lexicon with an embedding ``source`` and optional cache."""
         super().__init__(seed=seed)
         self._adapter = _resolve_source(source)
         self._max_neighbors = max(1, max_neighbors)
@@ -350,6 +404,7 @@ class VectorLexicon(LexiconBackend):
         return synonyms
 
     def get_synonyms(self, word: str, pos: str | None = None, n: int = 5) -> list[str]:
+        """Return up to ``n`` deterministic synonyms drawn from the embedding cache."""
         normalized = self._normalize_for_lookup(word)
         synonyms = self._ensure_cached(original=word, normalized=normalized)
         return self._deterministic_sample(synonyms, limit=n, word=word, pos=pos)
@@ -390,6 +445,7 @@ class VectorLexicon(LexiconBackend):
         return target
 
     def supports_pos(self, pos: str | None) -> bool:
+        """Always return ``True`` because vector sources do not encode POS metadata."""
         return True
 
     def __repr__(self) -> str:  # pragma: no cover - debug helper
@@ -452,7 +508,8 @@ def _parse_cli(argv: Sequence[str] | None = None) -> argparse.Namespace:
         "--source",
         required=True,
         help=(
-            "Vector source specification. Use 'spacy:<model>' for spaCy pipelines "
+            "Vector source specification. Use 'spacy:<model>' for spaCy pipelines, "
+            "'sentence-transformers:<model>' for HuggingFace checkpoints (requires --tokens), "
             "or provide a path to a gensim KeyedVectors/word2vec file."
         ),
     )
@@ -534,22 +591,44 @@ def main(argv: Sequence[str] | None = None) -> int:
 
         normalizer = _identity
 
-    source = load_vector_source(args.source)
+    tokens_from_file: list[str] | None = None
     if args.tokens is not None:
-        token_iter: Iterable[str] = _iter_tokens_from_file(args.tokens)
+        tokens_from_file = list(_iter_tokens_from_file(args.tokens))
+        if args.limit is not None:
+            tokens_from_file = tokens_from_file[: args.limit]
+
+    source_spec = args.source
+    token_iter: Iterable[str]
+    if source_spec.startswith("sentence-transformers:"):
+        model_name = source_spec.split(":", 1)[1].strip()
+        if not model_name:
+            model_name = "sentence-transformers/all-mpnet-base-v2"
+        if tokens_from_file is None:
+            raise SystemExit(
+                "Sentence-transformers sources require --tokens to supply a vocabulary."
+            )
+        source = _build_sentence_transformer_embeddings(model_name, tokens_from_file)
+        token_iter = tokens_from_file
     else:
-        lexicon = VectorLexicon(
-            source=source,
-            max_neighbors=args.max_neighbors,
-            min_similarity=args.min_similarity,
-            case_sensitive=args.case_sensitive,
-            normalizer=normalizer,
-            seed=args.seed,
-        )
-        token_iter = lexicon.iter_vocabulary()
-
-    if args.limit is not None:
-        token_iter = (token for index, token in enumerate(token_iter) if index < args.limit)
+        source = load_vector_source(source_spec)
+        if tokens_from_file is not None:
+            token_iter = tokens_from_file
+        else:
+            lexicon = VectorLexicon(
+                source=source,
+                max_neighbors=args.max_neighbors,
+                min_similarity=args.min_similarity,
+                case_sensitive=args.case_sensitive,
+                normalizer=normalizer,
+                seed=args.seed,
+            )
+            iterator = lexicon.iter_vocabulary()
+            if args.limit is not None:
+                token_iter = (
+                    token for index, token in enumerate(iterator) if index < args.limit
+                )
+            else:
+                token_iter = iterator
 
     build_vector_cache(
         source=source,