emb-diversity 0.0.3__py3-none-any.whl

emb_diversity/__init__.py

@@ -0,0 +1,63 @@
+"""emb-diversity – embedding-based diversity measures for text and vector data."""
+
+from __future__ import annotations
+
+### Distance-Based Diversity Measures
+from .measures.mean_pw_dist import mean_pw_dist
+from .measures.dist_dispersion import dist_dispersion
+from .measures.hamdiv import hamdiv
+from .measures.diameter import diameter
+from .measures.bottleneck import bottleneck
+from .measures.sum_bottleneck import sum_bottleneck
+from .measures.sum_diameter import sum_diameter
+from .measures.energy import energy
+from .measures.cluster_inertia import cluster_inertia
+from .measures.span_centroid import span_centroid
+from .measures.chamfer_dist import chamfer_dist
+
+### Volume-Based Diversity Measures
+from .measures.convex_hull_volume_2d import convex_hull_volume_2d
+from .measures.radius import radius
+from .measures.span_medoid import span_medoid
+
+### Distribution-Based Diversity Measures
+from .measures.vendi_score import vendi_score
+from .measures.renyi_entropy import renyi_entropy
+from .measures.dcscore import dcscore
+from .measures.log_determinant import log_determinant
+from .measures.bins_entropy import bins_entropy
+
+### Graph-Based Diversity Measures
+from .measures.graph_entropy import graph_entropy
+from .measures.mst_dispersion import mst_dispersion
+
+### Registries
+from .axes_registry import axes
+from .measures_registry import measures
+
+### Embedding helper
+from .embed import embed_texts
+
+### Main entry point
+from .convenience import measure_diversity
+
+### Caching utilities
+from .compute_pairwise import compute_pairwise_distances, clear_distance_cache, distance_cache_info
+
+
+__all__ = [
+    # Main entry point
+    "measure_diversity",
+    # Individual measures
+    "mean_pw_dist", "dist_dispersion", "hamdiv", "diameter", "bottleneck",
+    "sum_bottleneck", "sum_diameter", "energy", "cluster_inertia", "span_centroid", "chamfer_dist",
+    "convex_hull_volume_2d", "radius", "span_medoid", "vendi_score",
+    "renyi_entropy", "dcscore", "log_determinant", "bins_entropy",
+    "graph_entropy", "mst_dispersion",
+    # Helpers
+    "embed_texts",
+    # Registries
+    "axes", "measures",
+    # Pairwise distance caching
+    "compute_pairwise_distances", "clear_distance_cache", "distance_cache_info",
+]

emb_diversity/_accepts_text.py

@@ -0,0 +1,104 @@
+"""Decorator that lets measure functions accept raw text.
+
+When a measure function is decorated with ``@accepts_text``, it gains
+the ability to receive a list of strings instead of embeddings.  The
+decorator detects text input, embeds it via :func:`emb_diversity.embed.embed_texts`,
+and passes the resulting vectors to the original measure function.
+
+Two keyword arguments are added to the decorated function:
+
+- ``diversity_axis`` (default ``"semantic"``): which axis to use for embedding
+- ``embedding_model``: explicit model name, overrides the axis
+
+If the input is already numeric (e.g. a numpy array), the decorator
+passes it through unchanged.
+"""
+
+from __future__ import annotations
+
+import functools
+import inspect
+from typing import Sequence
+
+from .embed import embed_texts
+
+
+# ── Input detection ──────────────────────────────────────────────────
+
+
+def _is_text_input(data):
+    """Return True if *data* looks like a list of strings."""
+    return len(data) > 0 and isinstance(data[0], str)
+
+
+# ── Decorator ────────────────────────────────────────────────────────
+
+
+def accepts_text(func):
+    """Wrap a measure function so it can receive raw text.
+
+    The decorated function gains two optional keyword arguments:
+
+    * ``diversity_axis``  – registered axis name (default ``"semantic"``)
+    * ``embedding_model`` – explicit model id (overrides axis)
+
+    When the first positional argument (*data*) is a list of strings the
+    decorator embeds them before calling the original measure.  When *data*
+    is already numeric the decorator passes it through unchanged.
+    """
+
+    @functools.wraps(func)
+    def wrapper(data, *args, diversity_axis="semantic", embedding_model=None, **kwargs):
+        if _is_text_input(data):
+            data = embed_texts(
+                data,
+                diversity_axis=diversity_axis,
+                embedding_model=embedding_model,
+            )
+        return func(data, *args, **kwargs)
+
+    # Preserve the original signature but add the new params for docs/IDE
+    _patch_signature(wrapper, func)
+    return wrapper
+
+
+def _patch_signature(wrapper, func):
+    """Fix the wrapper's signature so help() and IDEs show useful parameter names.
+
+    This is purely cosmetic — it only affects what help() and IDEs display,
+    not how the function actually runs.
+
+    Example for mean_pw_dist:
+        Before: mean_pw_dist(data, *args, diversity_axis="semantic", embedding_model=None, **kwargs)
+        After:  mean_pw_dist(data, metric="cosine", diversity_axis="semantic", embedding_model=None, **metric_kwargs)
+    """
+    # The two parameters that the decorator adds
+    diversity_axis_param = inspect.Parameter(
+        "diversity_axis", inspect.Parameter.KEYWORD_ONLY, default="semantic",
+    )
+    embedding_model_param = inspect.Parameter(
+        "embedding_model", inspect.Parameter.KEYWORD_ONLY, default=None,
+    )
+    new_params = [diversity_axis_param, embedding_model_param]
+
+    # Read the original function's parameters
+    # e.g. for mean_pw_dist: [data, metric="cosine", **metric_kwargs]
+    original_sig = inspect.signature(func)
+    original_params = list(original_sig.parameters.values())
+
+    # Find the **kwargs parameter if it exists (e.g. **metric_kwargs)
+    kwargs_params = [p for p in original_params if p.kind == inspect.Parameter.VAR_KEYWORD]
+    has_kwargs = len(kwargs_params) > 0
+
+    # Build the combined parameter list:
+    #   original params + new params + **kwargs at the end
+    # We insert before **kwargs so the signature reads naturally:
+    #   (data, metric, diversity_axis, embedding_model, **metric_kwargs)
+    if has_kwargs:
+        kwargs_position = original_params.index(kwargs_params[0])
+        params_before_kwargs = original_params[:kwargs_position]
+        combined = params_before_kwargs + new_params + [kwargs_params[0]]
+    else:
+        combined = original_params + new_params
+
+    wrapper.__signature__ = original_sig.replace(parameters=combined)

emb_diversity/_registry.py

@@ -0,0 +1,51 @@
+"""Generic registry pattern for O(1) key-value lookup.
+
+Based on https://dev.to/dentedlogic/stop-writing-giant-if-else-chains-master-the-python-registry-pattern-ldm
+"""
+
+from __future__ import annotations
+
+class Registry:
+    """A simple key-value store with register, get, and list operations.
+
+    Used by both the axes registry and the measures registry to avoid
+    repeated dictionary boilerplate.
+    """
+
+    def __init__(self):
+        self._store: dict[str, object] = {}
+
+    def register(self, key: str, value: object) -> None:
+        """Add an entry. Overwrites if key already exists."""
+        self._store[key] = value
+
+    def get(self, key: str) -> object:
+        """Look up by key.
+
+        Raises:
+            KeyError: If the key has not been registered.
+        """
+        if key not in self._store:
+            registered = ", ".join(sorted(self._store)) or "(none)"
+            raise KeyError(f"Unknown key {key!r}. Registered: {registered}")
+        return self._store[key]
+
+    def list_all(self) -> list[object]:
+        """Return all values sorted by key."""
+        return [self._store[k] for k in sorted(self._store)]
+
+    def keys(self) -> list[str]:
+        """Return all registered keys."""
+        return list(self._store.keys())
+
+    def __contains__(self, key: str) -> bool:
+        return key in self._store
+
+    def __len__(self) -> int:
+        return len(self._store)
+
+    def __getitem__(self, key: str) -> object:
+        return self.get(key)
+
+    def __iter__(self):
+        return iter(self._store)

emb_diversity/axes_registry.py

@@ -0,0 +1,54 @@
+"""Diversity axes registry.
+
+A diversity axis maps a concept (e.g. "semantic", "style") to a default
+embedding model and optional alternatives.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+from ._registry import Registry
+
+
+@dataclass
+class DiversityAxis:
+    """Configuration for a diversity axis.
+
+    Attributes:
+        name: Short identifier (e.g. ``"semantic"``).
+        default_model: HuggingFace model id used by default for this axis.
+        alternative_models: Other models that work well for this axis.
+        description: Human-readable explanation shown in docs and CLI.
+    """
+
+    name: str
+    default_model: str
+    alternative_models: list[str] = field(default_factory=list)
+    description: str = ""
+
+
+# Module-level registry instance
+axes = Registry()
+
+# ── Built-in axes ────────────────────────────────────────────────────
+
+axes.register(
+    "semantic",
+    DiversityAxis(
+        name="semantic",
+        default_model="all-mpnet-base-v2",
+        alternative_models=["all-MiniLM-L6-v2"],
+        description="Meaning-based diversity using semantic similarity",
+    ),
+)
+
+axes.register(
+    "style",
+    DiversityAxis(
+        name="style",
+        default_model="AnnaWegmann/Style-Embedding",
+        alternative_models=["StyleDistance/styledistance", "rrivera1849/LUAR-MUD", "AIDA-UPM/star"],
+        description="Writing style diversity",
+    ),
+)

emb_diversity/cli.py

@@ -0,0 +1,212 @@
+"""Command-line interface for emb-diversity."""
+
+from __future__ import annotations
+
+import csv
+import json
+import sys
+from pathlib import Path
+from typing import Optional
+
+import typer
+
+app = typer.Typer(
+    name="emb-diversity",
+    help="Measure embedding-based diversity of text data.\n\nRun directly: emb-diversity texts.txt",
+    no_args_is_help=True,
+)
+
+
+@app.command(hidden=True)
+def default(
+    input_file: Path = typer.Argument(
+        ..., help="Path to input file (.txt, .csv, or .tsv)"
+    ),
+    measures: Optional[list[str]] = typer.Option(
+        None,
+        "--measure",
+        "-m",
+        help=(
+            "Measure to run. Use -m name for one, repeat for several: "
+            "-m mean_pw_dist -m diameter. "
+            "Use -m variety|balance|disparity for a named set, -m all for all. "
+            "Default: graph_entropy."
+        ),
+    ),
+    axis: str = typer.Option(
+        "semantic", "--axis", "-a", help="Diversity axis (e.g. semantic, style)."
+    ),
+    model: Optional[str] = typer.Option(
+        None, "--model", help="Explicit embedding model; overrides --axis."
+    ),
+    column: str = typer.Option(
+        "text", "--column", "-c", help="Column name containing text (for CSV/TSV)."
+    ),
+    output_format: str = typer.Option(
+        "table", "--format", "-f", help="Output format: table, json, csv."
+    ),
+) -> None:
+    """Measure diversity from a text file or CSV/TSV."""
+    _run_measure(input_file, measures, axis, model, column, output_format)
+
+
+@app.command("measure")
+def measure_cmd(
+    input_file: Path = typer.Argument(
+        ..., help="Path to input file (.txt, .csv, or .tsv)"
+    ),
+    measures: Optional[list[str]] = typer.Option(
+        None,
+        "--measure",
+        "-m",
+        help=(
+            "Measure to run. Use -m name for one, repeat for several: "
+            "-m mean_pw_dist -m diameter. "
+            "Use -m variety|balance|disparity for a named set, -m all for all. "
+            "Default: graph_entropy."
+        ),
+    ),
+    axis: str = typer.Option(
+        "semantic", "--axis", "-a", help="Diversity axis (e.g. semantic, style)."
+    ),
+    model: Optional[str] = typer.Option(
+        None, "--model", help="Explicit embedding model; overrides --axis."
+    ),
+    column: str = typer.Option(
+        "text", "--column", "-c", help="Column name containing text (for CSV/TSV)."
+    ),
+    output_format: str = typer.Option(
+        "table", "--format", "-f", help="Output format: table, json, csv."
+    ),
+) -> None:
+    """Measure diversity from a text file or CSV/TSV."""
+    _run_measure(input_file, measures, axis, model, column, output_format)
+
+
+def _run_measure(input_file, measures, axis, model, column, output_format):
+    """Shared logic for the measure command."""
+    from .convenience import measure_diversity
+
+    # ── Read texts ───────────────────────────────────────────────
+    texts = _read_texts(input_file, column)
+    if len(texts) < 2:
+        typer.echo("Error: need at least 2 texts to measure diversity.", err=True)
+        raise typer.Exit(code=1)
+
+    # ── Convert Typer's list format to measure_diversity()'s format ──
+    # Typer always gives a list (e.g. ["all"]), but measure_diversity()
+    # expects a plain string for "all" and named-set shortcuts.
+    from .measures_registry import MEASURE_SETS
+
+    if measures is None:
+        measure_arg = None
+    elif len(measures) == 1 and measures[0] in ("all", *MEASURE_SETS):
+        measure_arg = measures[0]
+    else:
+        measure_arg = measures
+
+    # ── Compute ──────────────────────────────────────────────────
+    typer.echo(f"Measuring diversity of {len(texts)} texts...", err=True)
+    try:
+        results = measure_diversity(
+            texts, measure=measure_arg, diversity_axis=axis, embedding_model=model,
+        )
+    except KeyError as exc:
+        # measure_diversity() raises KeyError for unknown measure names
+        typer.echo(f"Error: {exc}. Run 'emb-diversity list-measures' to see available measures.", err=True)
+        raise typer.Exit(code=1)
+
+    # ── Output ───────────────────────────────────────────────────
+    _print_results(results, output_format)
+
+
+@app.command("list-measures")
+def list_measures_cmd() -> None:
+    """List all available diversity measures."""
+    from .measures_registry import DEFAULT_MEASURE, MEASURE_SETS, measures
+
+    for name in sorted(measures):
+        tags = []
+        if name in DEFAULT_MEASURE:
+            tags.append("default")
+        for set_name, members in MEASURE_SETS.items():
+            if name in members:
+                tags.append(set_name)
+        suffix = f"  [{', '.join(tags)}]" if tags else ""
+        typer.echo(f"  {name}{suffix}")
+
+
+@app.command("list-axes")
+def list_axes_cmd() -> None:
+    """List registered diversity axes and their models."""
+    from .axes_registry import axes
+
+    for ax in axes.list_all():
+        typer.echo(f"  {ax.name}")
+        typer.echo(f"    default model: {ax.default_model}")
+        if ax.alternative_models:
+            typer.echo(f"    alternatives:  {', '.join(ax.alternative_models)}")
+        if ax.description:
+            typer.echo(f"    {ax.description}")
+        typer.echo()
+
+
+# ── Helpers ──────────────────────────────────────────────────────
+
+
+def _read_texts(path: Path, column: str) -> list[str]:
+    """Read texts from a file, returning a list of non-empty strings.
+
+    Supported formats:
+        - .txt: one text per line, empty lines are skipped
+        - .csv: comma-separated, reads the column specified by `column`
+        - .tsv: tab-separated, reads the column specified by `column`
+
+    For CSV/TSV, rows with empty values in the text column are dropped.
+
+    Args:
+        path: Path to the input file.
+        column: Column name to read from CSV/TSV files (default "text").
+
+    Returns:
+        List of text strings, stripped of leading/trailing whitespace.
+    """
+    import pandas as pd
+
+    suffix = path.suffix.lower()
+    if suffix == ".txt":
+        # Read each line as one text, skip empty lines
+        return [line.strip() for line in path.read_text().splitlines() if line.strip()]
+    elif suffix in (".csv", ".tsv"):
+        separator = "\t" if suffix == ".tsv" else ","
+        df = pd.read_csv(path, sep=separator)
+        if column not in df.columns:
+            typer.echo(
+                f"Error: column {column!r} not found. "
+                f"Available: {list(df.columns)}",
+                err=True,
+            )
+            raise typer.Exit(code=1)
+        return df[column].astype(str).tolist()
+    else:
+        typer.echo(f"Error: unsupported file extension {suffix!r}. Use .txt, .csv, or .tsv.", err=True)
+        raise typer.Exit(code=1)
+
+
+def _print_results(results: dict[str, float], fmt: str) -> None:
+    """Format and print results."""
+    if fmt == "json":
+        typer.echo(json.dumps(results, indent=2))
+    elif fmt == "csv":
+        writer = csv.writer(sys.stdout)
+        writer.writerow(["measure", "score"])
+        for name, score in results.items():
+            writer.writerow([name, score])
+    else:  # table
+        max_name = max(len(n) for n in results) if results else 0
+        for name, score in results.items():
+            typer.echo(f"  {name:<{max_name}}  {score:.6f}")
+
+
+if __name__ == "__main__":
+    app()

emb_diversity/compute_pairwise.py

@@ -0,0 +1,144 @@
+"""
+Two-level cached pairwise distance computation.
+
+scipy.pdist is the bottleneck when several measures are run on the same
+embedding matrix. This module wraps it with:
+
+  Level 1 — in-process memory: a small bounded LRU dict keyed by full
+    content fingerprint of the matrix plus the metric / kwargs. Up to
+    _MEMORY_MAX entries are kept; oldest is evicted on overflow.
+
+  Level 2 — disk: condensed distance array stored under .cache/pdist/
+    as safetensors, keyed the same way. Survives across processes — a
+    SLURM job that finished yesterday leaves a cache that today's job
+    can pick up.
+
+  Level 3 — compute: scipy.pdist + write through both layers.
+
+The cache key folds in the metric and any metric_kwargs, so different
+metrics on the same data do not collide.
+"""
+from collections import OrderedDict
+from pathlib import Path
+from typing import Any, Callable, Sequence, Union
+
+import numpy as np
+import xxhash
+from scipy.spatial.distance import pdist
+from safetensors.numpy import save_file, load_file
+
+DISTANCE_METRIC = Union[str, Callable[[np.ndarray, np.ndarray], float]]
+DEFAULT_CACHE_DIR = Path(".cache/pdist")
+# how many chunks we feed into the hash function at a time, to keep memory
+# usage constant regardless of input size
+_HASH_CHUNK = 1_000_000
+# how many distance matrices to keep in memory before evicting the oldest one (LRU)
+_MEMORY_MAX = 4
+
+# in-memory cache (LRU)
+_memory: "OrderedDict[str, np.ndarray]" = OrderedDict()
+
+
+def _fingerprint(X: np.ndarray) -> str:
+    """Full-content xxhash of an array, chunked to keep memory constant."""
+    h = xxhash.xxh64()
+    h.update(str(X.shape).encode())
+    h.update(str(X.dtype).encode())
+    flat = X.ravel()
+    for i in range(0, len(flat), _HASH_CHUNK):
+        h.update(flat[i:i + _HASH_CHUNK].tobytes())
+    return h.hexdigest()
+
+
+def _metric_key(metric: DISTANCE_METRIC, metric_kwargs: dict) -> str:
+    """Stable, filesystem-safe key for metric + kwargs."""
+    if not metric_kwargs and isinstance(metric, str):
+        return metric
+    parts = [str(metric)]
+    for k in sorted(metric_kwargs):
+        parts.append(f"{k}={metric_kwargs[k]!r}")
+    return xxhash.xxh64("|".join(parts).encode()).hexdigest()
+
+
+def _store_memory(key: str, result: np.ndarray) -> None:
+    if _MEMORY_MAX <= 0:
+        return
+    if key in _memory:
+        _memory.move_to_end(key)
+        return
+    if len(_memory) >= _MEMORY_MAX:
+        _memory.popitem(last=False)
+    _memory[key] = result
+
+
+def compute_pairwise_distances(
+    data: Sequence[Sequence[float]],
+    metric: DISTANCE_METRIC = "cosine",
+    cache_dir: Path = DEFAULT_CACHE_DIR,
+    **metric_kwargs: Any,
+) -> np.ndarray:
+    """
+    Compute pairwise distances with two-level (memory + disk) caching.
+
+    Args:
+        data: 2D array-like of shape (n_samples, n_features).
+        metric: Distance metric name (e.g. "cosine", "euclidean") or callable.
+        cache_dir: Root directory for the disk cache.
+        **metric_kwargs: Extra keyword arguments forwarded to scipy.pdist.
+            Included in the cache key so different kwargs do not collide.
+
+    Returns:
+        Condensed distance array (upper triangle from scipy.pdist).
+
+    Raises:
+        ValueError: If data is empty or single row.
+    """
+    X = np.asarray(data, dtype=float)
+    n = X.shape[0]
+    if n == 0:
+        raise ValueError("Cannot compute distances for empty data")
+    if n == 1:
+        raise ValueError("Cannot compute distances for single data point")
+
+    metric_id = _metric_key(metric, metric_kwargs)
+    fp = _fingerprint(X)
+    key = f"{fp}|{metric_id}"
+
+    # Level 1: in-memory match by content fingerprint
+    if key in _memory:
+        _memory.move_to_end(key)
+        return _memory[key]
+
+    # Level 2: disk
+    cache_dir.mkdir(parents=True, exist_ok=True)
+    path = cache_dir / f"{fp}_{metric_id}.safetensors"
+    if path.exists():
+        result = load_file(path)["distances"]
+        _store_memory(key, result)
+        return result
+
+    # Level 3: compute, populate both layers
+    result = pdist(X, metric=metric, **metric_kwargs)
+    _store_memory(key, result)
+    save_file({"distances": result}, path)
+    return result
+
+
+def clear_distance_cache(cache_dir: Path = DEFAULT_CACHE_DIR) -> None:
+    """Clear both memory and disk caches."""
+    import shutil
+    _memory.clear()
+    if cache_dir.exists():
+        shutil.rmtree(cache_dir)
+
+
+def distance_cache_info(cache_dir: Path = DEFAULT_CACHE_DIR) -> dict:
+    """Return memory and disk cache statistics."""
+    disk_files = list(cache_dir.glob("*.safetensors")) if cache_dir.exists() else []
+    return {
+        "memory_entries": len(_memory),
+        "memory_mb": round(sum(v.nbytes for v in _memory.values()) / 1024 / 1024, 2),
+        "memory_max": _MEMORY_MAX,
+        "disk_files": len(disk_files),
+        "disk_mb": round(sum(f.stat().st_size for f in disk_files) / 1024 / 1024, 2),
+    }