| 12 | |
| 13 | |
| 14 | class EmbeddingIndexManager: |
| 15 | def __init__( |
| 16 | self, |
| 17 | index_path: Path, |
| 18 | labels_path: Path, |
| 19 | max_size: int = 10_000, |
| 20 | dedup_threshold: float = 0.95, |
| 21 | ): |
| 22 | self._index_path = Path(index_path) |
| 23 | self._labels_path = Path(labels_path) |
| 24 | self._max_size = max_size |
| 25 | self._dedup_threshold = dedup_threshold |
| 26 | |
| 27 | if self._index_path.exists() and self._labels_path.exists(): |
| 28 | self._embeddings = np.load(self._index_path) |
| 29 | with open(self._labels_path, encoding="utf-8") as f: |
| 30 | self._labels = json.load(f) |
| 31 | else: |
| 32 | self._embeddings = np.empty((0, 384), dtype=np.float32) |
| 33 | self._labels = [] |
| 34 | |
| 35 | @property |
| 36 | def size(self) -> int: |
| 37 | return len(self._labels) |
| 38 | |
| 39 | def add(self, embedding: np.ndarray, tier_id: int) -> bool: |
| 40 | if self.size >= self._max_size: |
| 41 | self._prune_oldest(self._max_size - 1) |
| 42 | if self.size > 0: |
| 43 | sims = self._embeddings @ embedding |
| 44 | if np.max(sims) >= self._dedup_threshold: |
| 45 | max_idx = int(np.argmax(sims)) |
| 46 | if self._labels[max_idx] == tier_id: |
| 47 | return False |
| 48 | self._embeddings = np.vstack([self._embeddings, embedding.reshape(1, -1)]) |
| 49 | self._labels.append(tier_id) |
| 50 | return True |
| 51 | |
| 52 | def _prune_oldest(self, keep: int) -> None: |
| 53 | if self.size <= keep: |
| 54 | return |
| 55 | self._embeddings = self._embeddings[-keep:] |
| 56 | self._labels = self._labels[-keep:] |
| 57 | |
| 58 | def save(self) -> None: |
| 59 | self._index_path.parent.mkdir(parents=True, exist_ok=True) |
| 60 | np.save(self._index_path, self._embeddings) |
| 61 | with open(self._labels_path, "w", encoding="utf-8") as f: |
| 62 | json.dump(self._labels, f) |
no outgoing calls