embedding-analysis

Files (1)

• SKILL.md

SKILL.md (first lines)

# Embedding Analysis for Dataset Curation

## Overview

Embeddings turn unstructured text/images into a vector space where quality, diversity, and redundancy become measurable. This skill covers embedding-based dataset curation using sentence-transformers, distribution comparison metrics, semantic deduplication, and quality scoring techniques from the 2025-2026 literature.

## When to Use

Use this skill when:
- Detecting near-duplicate or semantically redundant examples.
- Comparing embedding distributions across dataset splits (train vs. val vs. test).
- Measuring dataset diversity or coverage in embedding space.
- Filtering examples based on embedding distance to a "gold" reference set.
- Applying NeMo Curator-style semantic dedup at scale.
- Computing DataRater-style quality signals from embedding neighborhoods.

Do not use for:
- Raw text deduplication (exact match) — use simple hashing.
- Model training — use `transformers` or `trl` skills.
- Vector database persistence — use `chromadb`, `milvus`, or `qdrant` skills.

## Installation

```bash
uv pip install sentence-transformers umap-learn scikit-learn numpy
```

## Core Workflow

### 1. Generate Embeddings

```python
from sentence_transformers import SentenceTransformer
import numpy as np

# Choose model based on domain
model = SentenceTransformer("all-MiniLM-L6-v2")  # fast, general
# model = SentenceTransformer("all-mpnet-base-v2")  # higher quality
# model = SentenceTransformer("BAAI/bge-large-en-v1.5")  # retrieval-optimized
# model = SentenceTransformer("intfloat/multilingual-e5-large")  # multilingual

# Batched encoding (memory-efficient)
embeddings = model.encode(
    texts,
    batch_size=64,
    show_progress_bar=True,
    normalize_embeddings=True,  # cosine similarity via dot product
    convert_to_numpy=True,
)
```

### 2. Semantic Deduplication

Semantic dedup removes examples that are meaning-equivalent but not text-identical. From NeMo Curator (NVIDIA, 2024-2025) and academic literature.

#### Clustering-Based Approach (NeMo Curator SemDedup)

```python
from sklearn.cluster import MiniBatchKMeans
from sklearn.metrics.pairwise import cosine_similarity

def semantic_dedup(embeddings, threshold=0.95, n_clusters=100):
    """
    Cluster embeddings, then remove near-duplicates within each cluster.
    threshold: cosine similarity above which two examples are duplicates.
    Returns: boolean mask (True = keep, False = duplicate).
    """
    n = len(embeddings)
    keep = np.ones(n, dtype=bool)

    # Coarse clustering for efficiency (O(n * k) instead of O(n^2))
    k = min(n_clusters, n // 10)
    clusters = MiniBatchKMeans(n_clusters=k, random_state=42, batch_size=1024).fit_predict(embeddings)

    for c in range(k):
        idx = np.where(clusters == c)[0]
        if len(idx) < 2:
            continue
        sim = cosine_similarity(embeddings[idx])
        # Mark duplicates (lower-index kept, higher-index removed)
        for i in range(len(idx)):
            if not keep[idx[i]]:
                continue
            dupes = np.where(sim[i, i+1:] > threshold)[0]
            for d in dupes:
                keep[idx[i + 1 + d]] = False

    return keep
```

#### Connected Components Approach (LSHBloom, 2025)

```python
# For internet-scale dedup (>100M examples):
# Use LSH for approximate nearest neighbor + union-find
from sklearn.neighbors import NearestNeighbors

def lsh_semantic_dedup(embeddings, threshold=0.90, n_neighbors=10):
    """Finds clusters of semantically similar examples via NN graph."""
    nn = NearestNeighbors(n_neighbors=n_neighbors, metric="cosine", n_jobs=-1)
    nn.fit(embeddings)
    distances, indices = nn.kneighbors(embeddings)

    # Union-find to merge connected components
    parent = np.arange(len(embeddings))
    def find(x):
        while parent[x] != x:
            parent[x] = parent[parent[x]]
            x = parent[x]
        return x
    def union(a, b):
        parent[find(a)] = find(b)

    for i in range(len(embeddings)):
        for j, d in zip(indices[i], distances[i]):
            if i != j and (1 - d) > threshold:  # cosine dist → similarity
                union(i, j)

    # Keep one per component

From mkurman/zorai, MIT.

More from mkurman/zorai