gradiend 0.1.0__py3-none-any.whl

gradiend/__init__.py

@@ -0,0 +1,103 @@
+"""
+GRADIEND: Feature Learning within Neural Networks
+
+GRADIEND is a method for learning features within neural networks
+by training an encoder-decoder architecture on gradients.
+
+Public API (from gradiend):
+    - Model: GradiendModel, ParamMappedGradiendModel, ModelWithGradiend
+    - Data: TextFilterConfig, TextPredictionDataCreator, TextPreprocessConfig,
+      SpacyTagSpec, preprocess_texts, resolve_base_data
+    - Trainer: TextPredictionTrainer, TextPredictionConfig, TrainingArguments,
+      load_training_stats, GradientTrainingDataset, TextGradientTrainingDataset,
+      create_model_with_gradiend
+    - Logging: setup_logging, get_logger
+
+Sub-packages (use when you need modality-specific or internal APIs):
+    - gradiend.trainer: Trainer, PrePruneConfig, PostPruneConfig, callbacks, etc.
+    - gradiend.trainer.text: TextModelWithGradiend, TextBatchedDataset, etc.
+    - gradiend.evaluator: EncoderEvaluator, DecoderEvaluator, Evaluator
+    - gradiend.visualizer: Visualizer, plot_encoder_distributions, plot_topk_overlap_heatmap, plot_topk_overlap_venn, etc.
+    - gradiend.data: same as top-level data API (alternative import path)
+
+For experimental features (analysis, plotting, LaTeX export), install with:
+    pip install gradiend[recommended]
+"""
+
+from importlib.metadata import PackageNotFoundError, version
+
+try:
+    __version__ = version("gradiend")
+except PackageNotFoundError:
+    __version__ = "0.1.0"  # editable install before package is installed
+
+# Core model classes
+from gradiend.model import GradiendModel, ParamMappedGradiendModel, ModelWithGradiend
+
+# High-level data API (filter config, data creator, preprocess)
+from gradiend.data import (
+    TextFilterConfig,
+    TextPredictionDataCreator,
+    TextPreprocessConfig,
+    SpacyTagSpec,
+    preprocess_texts,
+    resolve_base_data,
+)
+
+# Text prediction trainer (high-level)
+from gradiend.trainer.text.prediction.trainer import TextPredictionTrainer, TextPredictionConfig
+
+# Logging
+from gradiend.util.logging import setup_logging, get_logger
+
+# Visualization (top-k overlap plots)
+from gradiend.visualizer import (
+    plot_topk_overlap_heatmap,
+    plot_topk_overlap_venn,
+)
+
+# Training
+from gradiend.trainer import (
+    load_training_stats,
+    set_seed,
+    TrainingArguments,
+    TrainerConfig,
+    GradientTrainingDataset,
+    TextGradientTrainingDataset,
+    create_model_with_gradiend,
+    PrePruneConfig,
+    PostPruneConfig
+)
+
+__all__ = [
+    # Core model
+    "GradiendModel",
+    "ParamMappedGradiendModel",
+    "ModelWithGradiend",
+    # Data (high-level)
+    "TextFilterConfig",
+    "TextPredictionDataCreator",
+    "TextPreprocessConfig",
+    "SpacyTagSpec",
+    "preprocess_texts",
+    "resolve_base_data",
+    # Trainers and configs
+    "TextPredictionTrainer",
+    "TextPredictionConfig",
+    "TrainerConfig",
+    "PrePruneConfig",
+    "PostPruneConfig",
+    # Logging
+    "setup_logging",
+    "get_logger",
+    # Training
+    "load_training_stats",
+    "set_seed",
+    "TrainingArguments",
+    "GradientTrainingDataset",
+    "TextGradientTrainingDataset",
+    "create_model_with_gradiend",
+    # Visualization
+    "plot_topk_overlap_heatmap",
+    "plot_topk_overlap_venn",
+]

gradiend/data/__init__.py

@@ -0,0 +1,23 @@
+"""
+Data module: create training and neutral datasets from base corpora.
+
+Use TextPredictionDataCreator to filter and mask base texts based on TextFilterConfig.
+"""
+
+from gradiend.data.core import resolve_base_data
+from gradiend.data.text import (
+    SpacyTagSpec,
+    TextFilterConfig,
+    TextPreprocessConfig,
+    preprocess_texts,
+)
+from gradiend.data.text.prediction import TextPredictionDataCreator
+
+__all__ = [
+    "SpacyTagSpec",
+    "TextFilterConfig",
+    "TextPreprocessConfig",
+    "TextPredictionDataCreator",
+    "preprocess_texts",
+    "resolve_base_data",
+]

gradiend/data/core/__init__.py

@@ -0,0 +1,7 @@
+"""Core data loading (modality-agnostic)."""
+
+from gradiend.data.core.base_loader import resolve_base_data
+
+__all__ = [
+    "resolve_base_data",
+]

gradiend/data/core/base_loader.py

@@ -0,0 +1,121 @@
+"""
+Base data loader: resolve from HF, DataFrame, CSV, or list of strings.
+
+Modality-agnostic; used by both text prediction and (future) classification.
+"""
+
+from __future__ import annotations
+
+import random
+from pathlib import Path
+from typing import List, Optional, Union
+
+import pandas as pd
+
+from gradiend.util.logging import get_logger
+
+logger = get_logger(__name__)
+
+
+def resolve_base_data(
+    source: Union[str, pd.DataFrame, List[str]],
+    text_column: str = "text",
+    max_size: Optional[int] = None,
+    split: str = "train",
+    seed: int = 42,
+    hf_config: Optional[str] = None,
+    trust_remote_code: bool = False,
+) -> List[str]:
+    """Resolve input to a list of text strings.
+
+    Supports HuggingFace dataset ID, pandas DataFrame, CSV path, or list of strings.
+    Data is shuffled with `seed` before applying `max_size` to avoid bias from
+    ordered sources (e.g. chronological, single-author).
+
+    Args:
+        source: HF dataset ID (str), CSV path (str ending in .csv or path exists),
+            pandas DataFrame, or List[str].
+        text_column: Column name for text (default "text"). Ignored for List[str].
+        max_size: Cap on number of items (applied after shuffle). None = no cap.
+        split: Dataset split for HF (default "train").
+        seed: Random seed for shuffle.
+        hf_config: HuggingFace dataset config/subset (e.g. "20220301.en" for wikipedia).
+            Only used when source is an HF dataset ID.
+        trust_remote_code: Passed to load_dataset when loading from HF. Default False.
+
+    Returns:
+        List of strings (texts).
+    """
+    if not isinstance(text_column, str):
+        raise TypeError(f"text_column must be str, got {type(text_column).__name__}")
+    if max_size is not None and not isinstance(max_size, int):
+        raise TypeError(f"max_size must be int or None, got {type(max_size).__name__}")
+    if not isinstance(split, str):
+        raise TypeError(f"split must be str, got {type(split).__name__}")
+    if not isinstance(seed, int):
+        raise TypeError(f"seed must be int, got {type(seed).__name__}")
+    if hf_config is not None and not isinstance(hf_config, str):
+        raise TypeError(f"hf_config must be str or None, got {type(hf_config).__name__}")
+    if not isinstance(trust_remote_code, bool):
+        raise TypeError(f"trust_remote_code must be bool, got {type(trust_remote_code).__name__}")
+
+    texts: List[str]
+    if isinstance(source, list):
+        if not all(isinstance(x, str) for x in source):
+            raise TypeError("source as list must contain only strings")
+        texts = [str(x).strip() for x in source if str(x).strip()]
+    elif isinstance(source, pd.DataFrame):
+        if text_column not in source.columns:
+            raise ValueError(f"DataFrame missing column '{text_column}'. Columns: {list(source.columns)}")
+        texts = source[text_column].dropna().astype(str).str.strip().tolist()
+        texts = [x for x in texts if x]
+    elif isinstance(source, str):
+        texts = _load_from_string_source(
+            source, text_column, split, hf_config, trust_remote_code
+        )
+    else:
+        raise TypeError(f"source must be str, DataFrame, or List[str]; got {type(source)}")
+
+    rng = random.Random(seed)
+    rng.shuffle(texts)
+    if max_size is not None and len(texts) > max_size:
+        texts = texts[:max_size]
+    logger.debug(f"resolve_base_data: {len(texts)} texts")
+    return texts
+
+
+def _load_from_string_source(
+    source: str,
+    text_column: str,
+    split: str,
+    hf_config: Optional[str] = None,
+    trust_remote_code: bool = False,
+) -> List[str]:
+    """Load from HF dataset or CSV path."""
+    path = Path(source)
+    if path.suffix.lower() == ".csv" and path.exists():
+        df = pd.read_csv(source)
+        if text_column not in df.columns:
+            raise ValueError(f"CSV missing column '{text_column}'. Columns: {list(df.columns)}")
+        texts = df[text_column].dropna().astype(str).str.strip().tolist()
+        return [x for x in texts if x]
+
+    try:
+        from datasets import load_dataset
+    except ImportError as e:
+        raise ImportError("HuggingFace datasets required for HF source. pip install datasets") from e
+
+    if hf_config is not None:
+        ds = load_dataset(
+            source, hf_config, split=split, trust_remote_code=trust_remote_code
+        )
+    else:
+        ds = load_dataset(source, split=split, trust_remote_code=trust_remote_code)
+    if hasattr(ds, "to_pandas"):
+        df = ds.to_pandas()
+    else:
+        df = pd.DataFrame(list(ds))
+    if text_column not in df.columns:
+        raise ValueError(f"HF dataset missing column '{text_column}'. Columns: {list(df.columns)}")
+    texts = df[text_column].dropna().astype(str).str.strip().tolist()
+    return [x for x in texts if x]

gradiend/data/core/spacy_util.py

@@ -0,0 +1,63 @@
+"""
+Spacy model loading with clearer errors and optional auto-download.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+# Spacy model-not-found error code
+_SPACY_E050 = "E050"
+
+
+def load_spacy_model(
+    model_name: str,
+    download_if_missing: bool = False,
+) -> Any:
+    """Load a spacy model with improved error handling.
+
+    Args:
+        model_name: Spacy model name (e.g. "de_core_news_sm").
+        download_if_missing: If True and the model is not found, attempt to
+            download it via `python -m spacy download`. Default False.
+
+    Returns:
+        Loaded spacy Language model.
+
+    Raises:
+        ImportError: If spacy is not installed.
+        OSError: If the model is not found, with a message including the
+            install command: "python -m spacy download {model_name}".
+    """
+    try:
+        import spacy
+    except ImportError as e:
+        raise ImportError(
+            "spacy is required for morphological filtering. "
+            "Install with: pip install spacy "
+            "Or: pip install gradiend[data]"
+        ) from e
+
+    try:
+        return spacy.load(model_name)
+    except OSError as e:
+        err_str = str(e)
+        is_model_not_found = _SPACY_E050 in err_str or "Can't find model" in err_str
+
+        if is_model_not_found and download_if_missing:
+            try:
+                from spacy.cli import download
+                download(model_name)
+                return spacy.load(model_name)
+            except Exception as download_err:
+                raise OSError(
+                    f"Spacy model '{model_name}' not found and auto-download failed: {download_err}. "
+                    f"Install manually with: python -m spacy download {model_name}"
+                ) from e
+
+        if is_model_not_found:
+            raise OSError(
+                f"Spacy model '{model_name}' not found. "
+                f"Install it with: python -m spacy download {model_name}"
+            ) from e
+        raise

gradiend/data/text/__init__.py

@@ -0,0 +1,16 @@
+"""Text-specific data modules."""
+
+from gradiend.data.text.filter_config import SpacyTagSpec, TextFilterConfig
+from gradiend.data.text.preprocess import (
+    TextPreprocessConfig,
+    iter_sentences_from_texts,
+    preprocess_texts,
+)
+
+__all__ = [
+    "SpacyTagSpec",
+    "TextFilterConfig",
+    "TextPreprocessConfig",
+    "iter_sentences_from_texts",
+    "preprocess_texts",
+]

gradiend/data/text/filter_config.py

@@ -0,0 +1,93 @@
+"""
+Filter configuration for text: targets, spacy tags, lemma matching.
+
+Not prediction-specific; used by text filtering in general.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional, Union
+
+# Type alias for spacy morph/tag constraints (e.g. {"pos": "DET", "Case": "Nom"})
+SpacyTagSpec = Dict[str, str]
+
+
+def _merge_spacy_tags(parent: Optional[SpacyTagSpec], child: Optional[SpacyTagSpec]) -> Optional[SpacyTagSpec]:
+    """Merge child tags into parent; child overrides parent for same keys."""
+    if parent is None and child is None:
+        return None
+    if parent is None:
+        return dict(child) if child else None
+    if child is None:
+        return dict(parent)
+    out = dict(parent)
+    out.update(child)
+    return out
+
+
+def _first_target_string(targets: List[Union["TextFilterConfig", str]]) -> Optional[str]:
+    """Extract first target string from targets for default id."""
+    for t in targets:
+        if isinstance(t, str):
+            return t
+        if isinstance(t, TextFilterConfig):
+            s = _first_target_string(t.targets)
+            if s is not None:
+                return s
+    return None
+
+
+@dataclass
+class TextFilterConfig:
+    """Configuration for filtering and masking target tokens in text.
+
+    Simple use: TextFilterConfig(target="das") or TextFilterConfig(targets=["der", "die"], spacy_tags={...})
+    Advanced: TextFilterConfig(targets=[TextFilterConfig(["das"], spacy_tags={...}), "der"], spacy_tags={...})
+    Nested entries inherit parent spacy_tags and may add/override their own.
+
+    When used in TextPredictionDataCreator.feature_targets (a list), id names the output class.
+    If id is None, it is inferred from the first target string.
+    """
+
+    targets: Optional[List[Union["TextFilterConfig", str]]] = None
+    target: Optional[str] = None
+    spacy_tags: Optional[SpacyTagSpec] = None
+    use_lemma: bool = False
+    id: Optional[str] = None
+    mask: str = "[MASK]"
+    weight: float = 1.0
+
+    def __post_init__(self) -> None:
+        if self.target is not None and self.targets is not None:
+            raise ValueError("TextFilterConfig: provide either target or targets, not both")
+        if self.target is not None:
+            self.targets = [self.target]
+        elif isinstance(self.targets, str):
+            self.targets = [self.targets]
+        if self.targets is None:
+            self.targets = []
+        if not self.targets:
+            raise ValueError("TextFilterConfig.targets must be non-empty")
+        if self.id is None:
+            self.id = _first_target_string(self.targets)
+
+    def get_effective_tags(self, parent_tags: Optional[SpacyTagSpec] = None) -> Optional[SpacyTagSpec]:
+        """Get merged spacy_tags (self + parent)."""
+        return _merge_spacy_tags(parent_tags, self.spacy_tags)
+
+    def flatten_targets_with_tags(
+        self,
+        parent_tags: Optional[SpacyTagSpec] = None,
+    ) -> List[tuple[str, Optional[SpacyTagSpec]]]:
+        """Flatten nested config to list of (target_str, effective_spacy_tags)."""
+        effective = self.get_effective_tags(parent_tags)
+        out: List[tuple[str, Optional[SpacyTagSpec]]] = []
+        for t in self.targets:
+            if isinstance(t, str):
+                out.append((t, effective))
+            elif isinstance(t, TextFilterConfig):
+                out.extend(t.flatten_targets_with_tags(effective))
+            else:
+                raise TypeError(f"target must be str or TextFilterConfig; got {type(t)}")
+        return out

gradiend/data/text/prediction/__init__.py

@@ -0,0 +1,10 @@
+"""Text prediction data: filter engine, TextPredictionDataCreator."""
+
+from gradiend.data.text.prediction.creator import TextPredictionDataCreator
+from gradiend.data.text.prediction.filter_engine import filter_sentences, mask_sentence
+
+__all__ = [
+    "filter_sentences",
+    "mask_sentence",
+    "TextPredictionDataCreator",
+]