import hashlib

from typing import List, Optional, Set, Tuple

import numpy as np

from buffdata.engine.client import GeminiClient, LLMClient

from buffdata.models.schemas import DatasetItem

class Deduplicator:

"""Performs exact, lexical (MinHash n-gram), and semantic embedding deduplication."""

def __init__(self, client: Optional[LLMClient] = None):

self.client = client or GeminiClient()

@staticmethod

def _content(item: DatasetItem) -> str:

prompt, response = item.get_prompt_and_response()

combined = f"{prompt} {response}".strip()

return combined or item.get_classification_text()

def deduplicate_exact(self, items: List[DatasetItem]) -> Tuple[List[DatasetItem], List[DatasetItem]]:

"""Fast exact deduplication based on prompt+response xxhash fingerprints."""

seen_hashes: Set[int] = set()

kept: List[DatasetItem] = []

dropped: List[DatasetItem] = []

import xxhash

for it in items:

combined = self._content(it).strip().encode("utf-8")

h = xxhash.xxh64_intdigest(combined)

if h in seen_hashes:

it.metadata["dedup_reason"] = "exact_duplicate"

dropped.append(it)

else:

seen_hashes.add(h)

kept.append(it)

return kept, dropped

def deduplicate_minhash(

self,

items: List[DatasetItem],

threshold: float = 0.85,

shingle_size: int = 3,

) -> Tuple[List[DatasetItem], List[DatasetItem]]:

"""Lexical near-duplicate removal using Jaccard similarity of word shingles.

Each kept item's shingle set is stored as fixed-size 64-bit hash fingerprints

(xxhash), not the raw shingle strings -- for long texts this accumulator is the

dominant memory cost, and a fixed-size int is far cheaper than a variable-length

string. The Jaccard computation uses sorted-list intersection for O(n+m) instead

of set union/intersection which materializes temporary sets.

"""

import xxhash

def get_shingles(text: str) -> Set[int]:

words = text.lower().split()

if len(words) < shingle_size:

return {xxhash.xxh64_intdigest(word.encode("utf-8")) for word in words}

return {xxhash.xxh64_intdigest(" ".join(words[i:i+shingle_size]).encode("utf-8"))

for i in range(len(words) - shingle_size + 1)}

shingle_sets: List[List[int]] = []

kept: List[DatasetItem] = []

dropped: List[DatasetItem] = []

for it in items:

curr_shingles = get_shingles(self._content(it))

curr_sorted = sorted(curr_shingles)

is_dup = False

for prev_sorted in shingle_sets:

# O(n+m) intersection of sorted lists

i = j = 0

intersection = 0

while i < len(curr_sorted) and j < len(prev_sorted):

if curr_sorted[i] == prev_sorted[j]:

intersection += 1

i += 1

j += 1

elif curr_sorted[i] < prev_sorted[j]:

i += 1

else:

j += 1

union = len(curr_sorted) + len(prev_sorted) - intersection

if union > 0:

sim = intersection / union

if sim >= threshold:

is_dup = True

break

if is_dup:

it.metadata["dedup_reason"] = f"minhash_near_dup (threshold {threshold})"

dropped.append(it)

else:

shingle_sets.append(curr_sorted)

kept.append(it)

return kept, dropped

def deduplicate_semantic_local(

self,

items: List[DatasetItem],

threshold: float = 0.85,

batch_size: int = 256,

model_name: str = "all-MiniLM-L6-v2",

) -> Tuple[List[DatasetItem], List[DatasetItem]]:

"""Deduplicate items semantically using batch matrix cosine similarity.

Computes all embeddings at once, then uses vectorized matrix operations

to find near-duplicates in batch instead of sequential per-item comparisons.

Uses incremental blocked comparison to handle large datasets without OOM.

"""

import torch

from sentence_transformers import SentenceTransformer

device = "cuda" if torch.cuda.is_available() else "cpu"

model = SentenceTransformer(model_name, device=device)

texts = [self._content(item) for item in items]

embeddings = model.encode(texts, batch_size=batch_size, convert_to_tensor=True, show_progress_bar=True)

norms = torch.norm(embeddings, dim=1, keepdim=True)

norms[norms == 0] = 1.0

normed = embeddings / norms

kept_indices: List[int] = []

dropped: List[DatasetItem] = []

BLOCK = 4096

for start in range(0, len(items), BLOCK):

end = min(start + BLOCK, len(items))

block = normed[start:end]

if kept_indices:

kept_matrix = normed[kept_indices]

sim_matrix = torch.mm(block, kept_matrix.t())

max_sims, _ = sim_matrix.max(dim=1)

for local_idx in range(block.size(0)):

global_idx = start + local_idx

if max_sims[local_idx].item() >= threshold:

items[global_idx].metadata["dedup_reason"] = f"semantic_local_dup (sim: {max_sims[local_idx].item():.3f})"

dropped.append(items[global_idx])

else:

kept_indices.append(global_idx)

else:

for local_idx in range(block.size(0)):

kept_indices.append(start + local_idx)

return [items[i] for i in kept_indices], dropped

def deduplicate_semantic(

self,

items: List[DatasetItem],

threshold: float = 0.88,

) -> Tuple[List[DatasetItem], List[DatasetItem]]:

"""Semantic vector deduplication using Gemini text embeddings with batch matrix ops."""

if not items:

return [], []

texts = [self._content(item) for item in items]

embeddings = self.client.embed_texts(texts)

vecs = np.array(embeddings, dtype=np.float32)

norms = np.linalg.norm(vecs, axis=1, keepdims=True)

norms[norms == 0] = 1.0

normed = vecs / norms

kept_indices: List[int] = []

dropped_indices: List[int] = []

BLOCK = 4096

for start in range(0, len(items), BLOCK):

end = min(start + BLOCK, len(items))

block = normed[start:end]

if kept_indices:

kept_matrix = normed[kept_indices]

sim_matrix = block @ kept_matrix.T

max_sims = sim_matrix.max(axis=1)

for local_idx in range(block.shape[0]):

global_idx = start + local_idx

if max_sims[local_idx] >= threshold:

items[global_idx].metadata["dedup_reason"] = f"semantic_dup (sim: {max_sims[local_idx]:.3f} >= {threshold})"

dropped_indices.append(global_idx)

else:

kept_indices.append(global_idx)

else:

kept_indices.extend(range(start, end))

kept = [items[i] for i in kept_indices]

dropped = [items[i] for i in dropped_indices]

return kept, dropped