embedding-condensation 0.1.0__tar.gz

embedding_condensation-0.1.0/.gitignore

@@ -0,0 +1,31 @@
+.vscode
+data
+**/__pycache__
+**/.DS_Store
+**/results/
+
+# Models
+**/*.pt
+**/*.pkl
+**/*.pth
+**/*.pth.tar
+
+# Data
+**/*.npy
+**/*.npz
+
+# local env configs
+env.sh
+
+slurm_out/
+*.csv
+*.old
+wandb/
+*.out
+*.png
+old/
+*.arrow
+*.json
+eval_log.txt
+model_cache/
+results_old/

embedding_condensation-0.1.0/LICENSE

@@ -0,0 +1 @@
+Non-Commercial License — see ../LICENSE.md in the LM-Dispersion repository.

embedding_condensation-0.1.0/PKG-INFO

@@ -0,0 +1,71 @@
+Metadata-Version: 2.4
+Name: embedding-condensation
+Version: 0.1.0
+Summary: Measure layer-wise token embedding cosine similarity (embedding condensation diagnostic).
+Project-URL: Homepage, https://chenliu-1996.github.io/projects/LM-Dispersion/
+Project-URL: Repository, https://github.com/ChenLiu-1996/LM-Dispersion
+Author: Chen Liu
+License: Non-Commercial License — see ../LICENSE.md in the LM-Dispersion repository.
+License-File: LICENSE
+Keywords: embeddings,language-models,representation-geometry,transformers
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.9
+Requires-Dist: datasets>=2.14
+Requires-Dist: matplotlib>=3.7
+Requires-Dist: nltk>=3.8
+Requires-Dist: numpy>=1.24
+Requires-Dist: torch>=2.0
+Requires-Dist: tqdm>=4.65
+Requires-Dist: transformers<4.48,>=4.40
+Provides-Extra: test
+Requires-Dist: pytest>=7.0; extra == 'test'
+Description-Content-Type: text/markdown
+
+# embedding-condensation
+
+Minimal library for the **embedding condensation** diagnostic from [LM-Dispersion](https://github.com/ChenLiu-1996/LM-Dispersion): layer-wise token cosine-similarity matrices and optional heatmaps.
+
+## Install
+
+```bash
+cd pypi
+pip install -e ".[test]"
+```
+
+## Usage
+
+```python
+from transformers import AutoModel, AutoTokenizer
+from embedding_condensation import measure_embedding_condensation
+
+model = AutoModel.from_pretrained("gpt2").eval()
+tokenizer = AutoTokenizer.from_pretrained("gpt2")
+
+result = measure_embedding_condensation(
+    model,
+    tokenizer,
+    texts=["Your long input text here. " * 200],
+    repetitions=1,
+    plot=False,
+)
+print(result.mean_cossim_by_layer)
+```
+
+## PyPI upload
+
+```bash
+cd pypi
+pip install build twine
+python -m build
+twine upload dist/*
+```
+
+## Test
+
+```bash
+cd pypi
+pytest
+```

embedding_condensation-0.1.0/README.md

@@ -0,0 +1,45 @@
+# embedding-condensation
+
+Minimal library for the **embedding condensation** diagnostic from [LM-Dispersion](https://github.com/ChenLiu-1996/LM-Dispersion): layer-wise token cosine-similarity matrices and optional heatmaps.
+
+## Install
+
+```bash
+cd pypi
+pip install -e ".[test]"
+```
+
+## Usage
+
+```python
+from transformers import AutoModel, AutoTokenizer
+from embedding_condensation import measure_embedding_condensation
+
+model = AutoModel.from_pretrained("gpt2").eval()
+tokenizer = AutoTokenizer.from_pretrained("gpt2")
+
+result = measure_embedding_condensation(
+    model,
+    tokenizer,
+    texts=["Your long input text here. " * 200],
+    repetitions=1,
+    plot=False,
+)
+print(result.mean_cossim_by_layer)
+```
+
+## PyPI upload
+
+```bash
+cd pypi
+pip install build twine
+python -m build
+twine upload dist/*
+```
+
+## Test
+
+```bash
+cd pypi
+pytest
+```

embedding_condensation-0.1.0/embedding_condensation/__init__.py

@@ -0,0 +1,235 @@
+"""Embedding condensation measurement for Hugging Face transformer models."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import List, Optional, Sequence
+
+import os
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from datasets import load_dataset
+from nltk.tokenize import word_tokenize
+from tqdm import tqdm
+
+try:
+    word_tokenize('Arbitrary sentence.')
+except:
+    import nltk
+    nltk.download('punkt_tab')
+
+
+__all__ = [
+    "CondensationResult",
+    "measure_embedding_condensation",
+    "plot_similarity_heatmap",
+]
+
+
+@dataclass
+class CondensationResult:
+    """Per-layer token cosine-similarity matrices (L, S, S)."""
+    cossim_by_layer: np.ndarray
+
+
+def get_random_long_text(
+    dataset_name: str,
+    min_word_count: int = 1024,
+    max_word_count: int = 1280,
+    split: str = "train",
+    random_seed: int = 0,
+) -> str:
+    if dataset_name == "wikipedia":
+        dataset = load_dataset("wikitext", "wikitext-103-v1")
+        key = "text"
+    elif dataset_name == "pubmed":
+        dataset = load_dataset("pubmed_qa", "pqa_labeled")
+        key = "long_answer"
+    elif dataset_name == "imdb":
+        dataset = load_dataset("imdb")
+        key = "text"
+    elif dataset_name == "squad":
+        dataset = load_dataset("squad")
+        key = "context"
+    else:
+        raise ValueError(
+            f"Unknown dataset {dataset_name!r}. "
+            "Use wikipedia, pubmed, imdb, squad, or pass `texts=` directly."
+        )
+
+    text = ""
+    rng = np.random.default_rng(seed=random_seed)
+    idx = rng.integers(0, int(len(dataset["train"]) * 0.95)).item()
+    while len(word_tokenize(text)) < min_word_count:
+        text += dataset[split][idx][key]
+        idx += 1
+        if len(word_tokenize(text)) > max_word_count:
+            break
+    return text
+
+
+def _normalize(x: np.ndarray, p: int = 2, axis: int = 1, eps: float = 1e-3) -> np.ndarray:
+    norm = np.linalg.norm(x, ord=p, axis=axis, keepdims=True)
+    return x / np.maximum(norm, eps)
+
+
+def organize_embeddings(embeddings: Sequence[torch.Tensor]) -> List[np.ndarray]:
+    return [z.squeeze(0).float().cpu().numpy() for z in embeddings]
+
+
+def compute_cosine_similarities(embeddings: Sequence[np.ndarray]) -> List[np.ndarray]:
+    out = []
+    for z in embeddings:
+        z = _normalize(z, axis=1)
+        out.append(np.matmul(z, z.T).clip(-1, 1))
+    return out
+
+
+def plot_similarity_heatmap(
+    cossim_matrix_by_layer: Sequence[np.ndarray],
+    save_path: Optional[str] = None,
+    step: int = 1,
+    bins: int = 128,
+):
+    n_layers = len(cossim_matrix_by_layer)
+    denom = max(n_layers - 1, 1)
+    selected = [(i, data) for i, data in enumerate(cossim_matrix_by_layer) if i % step == 0]
+    layer_fractions, hist_data = [], []
+    for layer_idx, cossim_matrix in selected:
+        hist, _ = np.histogram(cossim_matrix.flatten(), bins=bins, density=True, range=(-1, 1))
+        hist_data.append(hist)
+        layer_fractions.append(layer_idx / denom)
+    hist_matrix = np.array(hist_data)
+
+    plt.rcParams["font.family"] = "sans-serif"
+    fig = plt.figure(figsize=(12, 10))
+    ax = fig.add_subplot(1, 1, 1)
+    ax.spines["top"].set_visible(False)
+    ax.spines["right"].set_visible(False)
+    im = ax.imshow(
+        hist_matrix.T,
+        aspect="auto",
+        origin="lower",
+        cmap="Reds",
+        extent=[0, layer_fractions[-1], -1, 1],
+        vmin=0,
+        vmax=10,
+    )
+    ax.tick_params(axis="both", which="major", labelsize=26)
+    ax.set_xlabel("Layer Fraction", fontsize=36)
+    ax.set_xticks([0, 0.2, 0.4, 0.6, 0.8, 1])
+    ax.set_xticklabels([0, 0.2, 0.4, 0.6, 0.8, 1])
+    ax.set_ylabel("Cosine Similarity", fontsize=36)
+    cbar = fig.colorbar(im, ax=ax)
+    cbar.ax.tick_params(axis="both", which="major", labelsize=26)
+    cbar.ax.set_title("Probability\nDensity", fontsize=20, pad=20)
+    fig.tight_layout(pad=2)
+
+    os.makedirs(os.path.dirname(save_path) or ".", exist_ok=True)
+    fig.savefig(save_path, dpi=300)
+    plt.close(fig)
+    return fig
+
+
+def _mean_cossim_per_layer(cossim_by_layer: Sequence[np.ndarray]) -> np.ndarray:
+    return np.array([float(m[np.triu_indices(m.shape[0], k=1)].mean()) for m in cossim_by_layer])
+
+
+def measure_embedding_condensation(
+    model: torch.nn.Module,
+    tokenizer,
+    *,
+    texts: Optional[Sequence[str]] = None,
+    dataset: str = "wikipedia",
+    repetitions: int = 100,
+    max_length: int = 512,
+    min_word_count: int = 1024,
+    max_word_count: int = 1280,
+    include_logits_layer: bool = False,
+    plot: bool = True,
+    save_path: Optional[str] = "./test_embedding_condensation.png",
+    show_progress: bool = True,
+) -> CondensationResult:
+    """
+    Run the LM-Dispersion embedding condensation measurement pipeline.
+
+    Pass `texts` to use fixed inputs. Otherwise samples
+    random long text from a Hugging Face dataset each repetition.
+    """
+    model.eval()
+    device = next(model.parameters()).device
+    stacked: Optional[List[np.ndarray]] = None
+
+    if texts is not None:
+        if len(texts) == 0:
+            raise ValueError("`texts` must be non-empty when provided.")
+        rep_iter = range(repetitions)
+        if show_progress:
+            rep_iter = tqdm(rep_iter, desc="condensation")
+        for r in rep_iter:
+            torch.manual_seed(r)
+            text = texts[r % len(texts)]
+            curr = _forward_cossim(
+                model, tokenizer, text, device, max_length, include_logits_layer
+            )
+            stacked = _stack_repetition(stacked, curr)
+    else:
+        rep_iter = range(repetitions)
+        if show_progress:
+            rep_iter = tqdm(rep_iter, desc="condensation")
+        for random_seed in rep_iter:
+            torch.manual_seed(random_seed)
+            text = get_random_long_text(
+                dataset,
+                random_seed=random_seed,
+                min_word_count=min_word_count,
+                max_word_count=max_word_count,
+            )
+            curr = _forward_cossim(
+                model, tokenizer, text, device, max_length, include_logits_layer
+            )
+            stacked = _stack_repetition(stacked, curr)
+
+    assert stacked is not None
+    averaged = [m.mean(axis=0) for m in stacked]
+    cossim_arr = np.stack(averaged, axis=0)
+    result = CondensationResult(
+        cossim_by_layer=cossim_arr,
+    )
+
+    if plot or save_path:
+        plot_similarity_heatmap(averaged, save_path=save_path)
+    return result
+
+
+def _forward_cossim(
+    model: torch.nn.Module,
+    tokenizer,
+    text: str,
+    device: torch.device,
+    max_length: int,
+    include_logits_layer: bool,
+) -> List[np.ndarray]:
+    tokens = tokenizer(text, return_tensors="pt", truncation=True, max_length=max_length)
+    tokens = {k: v.to(device) for k, v in tokens.items()}
+    with torch.no_grad():
+        output = model(**tokens, output_hidden_states=True)
+        curr = compute_cosine_similarities(organize_embeddings(output.hidden_states))
+        if include_logits_layer and hasattr(output, "logits"):
+            logits = torch.nn.functional.normalize(output.logits.squeeze(0).float(), dim=1)
+            curr.append(torch.matmul(logits, logits.T).clamp(-1, 1).cpu().numpy())
+    return curr
+
+
+def _stack_repetition(
+    stacked: Optional[List[np.ndarray]],
+    curr: List[np.ndarray],
+) -> List[np.ndarray]:
+    clipped = [m.clip(-1, 1) for m in curr]
+    if stacked is None:
+        return [m[None, ...] for m in clipped]
+    return [
+        np.concatenate((stacked[i], clipped[i][None, ...]), axis=0)
+        for i in range(len(stacked))
+    ]

embedding_condensation-0.1.0/pyproject.toml

@@ -0,0 +1,42 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "embedding-condensation"
+version = "0.1.0"
+description = "Measure layer-wise token embedding cosine similarity (embedding condensation diagnostic)."
+readme = "README.md"
+requires-python = ">=3.9"
+license = { file = "LICENSE" }
+authors = [{ name = "Chen Liu" }]
+keywords = ["language-models", "embeddings", "transformers", "representation-geometry"]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Science/Research",
+    "Programming Language :: Python :: 3",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+dependencies = [
+    "numpy>=1.24",
+    "torch>=2.0",
+    "transformers>=4.40,<4.48",
+    "datasets>=2.14",
+    "nltk>=3.8",
+    "matplotlib>=3.7",
+    "tqdm>=4.65",
+]
+
+[project.optional-dependencies]
+test = ["pytest>=7.0"]
+
+[project.urls]
+Homepage = "https://chenliu-1996.github.io/projects/LM-Dispersion/"
+Repository = "https://github.com/ChenLiu-1996/LM-Dispersion"
+
+[tool.hatch.build.targets.wheel]
+packages = ["embedding_condensation"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+addopts = "-q"