Image Similarity Pipeline: Graph Supervision, Contrastive Training, Vector Search¶

import numpy as np
from scipy import stats

def cluster_coherence_filter(edges, embeddings, cluster_ids, z_threshold_drop=-3, z_threshold_flag=-2):
    """
    edges: list of (src_id, dst_id)
    embeddings: {id: np.array} - CSD or CLIP vectors
    cluster_ids: {id: int}
    Returns: (confirmed, flagged, dropped)
    """
    from sklearn.metrics.pairwise import cosine_similarity

    cluster_edges = {}  # cluster_id -> list of (src, dst, similarity)
    for src, dst in edges:
        sim = cosine_similarity([embeddings[src]], [embeddings[dst]])[0][0]
        cid = cluster_ids[src]
        cluster_edges.setdefault(cid, []).append((src, dst, sim))

    confirmed, flagged, dropped = [], [], []
    for cid, edge_list in cluster_edges.items():
        sims = [e[2] for e in edge_list]
        if len(sims) < 10:
            confirmed.extend(edge_list)
            continue
        z_scores = stats.zscore(sims)
        for (src, dst, sim), z in zip(edge_list, z_scores):
            if z < z_threshold_drop:
                dropped.append((src, dst, sim, z))
            elif z < z_threshold_flag:
                flagged.append((src, dst, sim, z))
            else:
                confirmed.append((src, dst, sim, z))

    return confirmed, flagged, dropped

def nn_graph_partition(pin_ids, edges, es_client, k=50):
    """
    Partition edges into: confirmed (in both pin_related AND CSD kNN),
    drop_candidate (in pin_related but NOT in CSD kNN),
    add_candidate (in CSD kNN but NOT in pin_related)
    Typical ratio: 60% confirmed, 25% drop-candidate, 15% add-candidate
    """
    edge_set = set((a, b) for a, b in edges)
    confirmed, drop_candidate, add_candidate = [], [], []

    for pin_id in pin_ids:
        # kNN query to ES for this pin's CSD vector
        knn_results = es_client.knn_search(pin_id, k=k, field="csd_vec")
        knn_set = set(r['_id'] for r in knn_results)

        for neighbor_id in knn_set:
            if (pin_id, neighbor_id) in edge_set:
                confirmed.append((pin_id, neighbor_id))
            else:
                add_candidate.append((pin_id, neighbor_id))

        for src, dst in edge_set:
            if src == pin_id and dst not in knn_set:
                drop_candidate.append((src, dst))

    return confirmed, drop_candidate, add_candidate

def signal_disagreement_rank(edges, clip_vecs, csd_vecs, color_jaccard_fn):
    """Rank edges by variance across similarity signals.
    High variance = model uncertainty = most informative for labeling.
    Target 2-5K labels from the top of this ranking.
    """
    from sklearn.metrics.pairwise import cosine_similarity

    scores = []
    for src, dst in edges:
        csd_sim = cosine_similarity([csd_vecs[src]], [csd_vecs[dst]])[0][0]
        clip_sim = cosine_similarity([clip_vecs[src]], [clip_vecs[dst]])[0][0]
        color_sim = color_jaccard_fn(src, dst)
        variance = np.var([csd_sim, clip_sim, color_sim])
        scores.append((src, dst, variance, csd_sim, clip_sim, color_sim))

    return sorted(scores, key=lambda x: -x[2])  # highest variance first

import torch
import torch.nn as nn
import torch.nn.functional as F

class SimilarityProjectionHead(nn.Module):
    def __init__(self, input_dim=1560, hidden_dim=1024, output_dim=256):
        super().__init__()
        # input_dim = CLIP(768) + CSD(768) + color(24) = 1560
        # Add DINOv3(1024) -> input_dim = 2584
        self.net = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.GELU(),
            nn.Linear(hidden_dim, 512),
            nn.GELU(),
            nn.Linear(512, output_dim),
        )

    def forward(self, x):
        return F.normalize(self.net(x), dim=-1)


def infonce_loss(q, k_pos, k_neg, temperature=0.07):
    """
    q: [B, D], k_pos: [B, D], k_neg: [B, K, D]
    All L2-normalized.
    """
    logits_pos = (q * k_pos).sum(-1, keepdim=True) / temperature
    logits_neg = torch.einsum('bd,bkd->bk', q, k_neg) / temperature
    logits = torch.cat([logits_pos, logits_neg], dim=1)  # [B, 1+K]
    labels = torch.zeros(q.size(0), dtype=torch.long, device=q.device)
    return F.cross_entropy(logits, labels)

# Precompute backbone features first (saves per-epoch re-running)
# Load: clip_vecs[pin_id], csd_vecs[pin_id], dinov3_vecs[pin_id], color_vecs[pin_id]

optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=3)

# Batch: (pin_a_features, pin_b_features) from cleaned edges
# Hard negatives: K=4 samples from same cluster as pin_a, not in edges
# Effective batch: 1024+ (bigger = more in-batch negatives = better)
# Epochs: 3-5
# H100 training time: 2-4 hours
# Expected Recall@10 lift: +5-10 points over weighted-kNN baseline

from peft import LoraConfig, get_peft_model

lora_config = LoraConfig(r=16, lora_alpha=32, target_modules=["q_proj", "v_proj"])
siglip_model = get_peft_model(siglip_model, lora_config)
# Same InfoNCE loss, unfrozen encoder
# Training time: 6-10h H100, cost ~$15

def sample_hard_negatives(batch_ids, cluster_ids, edge_set, all_ids_by_cluster, K=4, margin=0.7):
    """
    For each pin, sample K negatives from same cluster that:
    1. Are NOT in edge_set (true negatives only)
    2. Have similarity < margin (avoid false negatives)
    """
    hard_negs = []
    for pin_id in batch_ids:
        cid = cluster_ids[pin_id]
        candidates = [x for x in all_ids_by_cluster[cid]
                     if (pin_id, x) not in edge_set and x != pin_id]
        hard_negs.append(np.random.choice(candidates, K, replace=False))
    return np.array(hard_negs)

# Sample 200 query pins, stratified by cluster
# For each query: take top 50 candidates from current baseline kNN
# Human-label each (relevant=1, borderline=0.5, irrelevant=0)
# Total: 10,000 judgements, ~10-12h work, save to test_pairs.sqlite
# NEVER include test pins in training edges

# kNN query with num_candidates tuning
body = {
    "knn": {
        "field": "projection_vec",
        "query_vector": query_embedding.tolist(),
        "k": 10,
        "num_candidates": 200  # never below k*10
    }
}

# ES Reciprocal Rank Fusion: dense visual + sparse tag signals
body = {
    "retriever": {
        "rrf": {
            "retrievers": [
                {"knn": {"field": "projection_vec", "query_vector": vec, "k": 50}},
                {"standard": {"query": {"terms": {"color_tags": query_tags}}}},
                {"standard": {"query": {"terms": {"mood_tags": query_moods}}}}
            ],
            "rank_constant": 60,
            "rank_window_size": 100
        }
    }
}
# Expected lift: 5-15% NDCG over dense-only for visual similarity