refactor

Author: b8zhong

python-threatexchange/threatexchange/signal_type/pdq/pdq_index2.py

@@ -24,116 +24,129 @@
 
 PDQIndexMatch = IndexMatchUntyped[SignalSimilarityInfoWithIntDistance, IndexT]
 
+def pick_n_centroids(n: int) -> int:
+    """Pick number of centroids based on dataset size."""
+    return max(int(n ** 0.5), 1)
+
+class _PDQHashIndex:
+    """
+    A wrapper around the faiss index for pickle serialization
+    """
+
+    def __init__(self, faiss_index: faiss.Index) -> None:
+        self.faiss_index = faiss_index
+
+    def search(
+        self,
+        queries: t.Sequence[str],
+        threshold: int,
+    ) -> t.List[t.List[t.Tuple[int, float]]]:
+        """
+        Search method that returns a mapping from query_str => (id, distance)
+        """
+        qs = convert_pdq_strings_to_ndarray(queries)
+        limits, D, I = self.faiss_index.range_search(qs, threshold + 1)
+
+        results = []
+        for i in range(len(queries)):
+            matches = [result.item() for result in I[limits[i] : limits[i + 1]]]
+            distances = [dist for dist in D[limits[i] : limits[i + 1]]]
+            results.append(list(zip(matches, distances)))
+        return results
+
+    def add(self, vectors: np.ndarray) -> None:
+        """Add vectors to the index."""
+        self.faiss_index.add(vectors)
+
+    def train(self, vectors: np.ndarray) -> None:
+        """Train the index if needed."""
+        if isinstance(self.faiss_index, faiss.IndexIVFFlat):
+            self.faiss_index.train(vectors)
+
+    def __getstate__(self):
+        return faiss.serialize_index(self.faiss_index)
+
+    def __setstate__(self, data):
+        self.faiss_index = faiss.deserialize_index(data)
+
 
 class PDQIndex2(SignalTypeIndex[IndexT]):
     """
     Indexing and querying PDQ signals using Faiss for approximate nearest neighbor search.
-
-    This is a redo of the existing PDQ index,
-    designed to be simpler and fix hard-to-squash bugs in the existing implementation.
-    Purpose of this class: to replace the original index in pytx 2.0
     """
 
     def __init__(
         self,
-        index: t.Optional[faiss.Index] = None,
-        entries: t.Iterable[t.Tuple[str, IndexT]] = (),
-        *,
         threshold: int = PDQ_CONFIDENT_MATCH_THRESHOLD,
+        faiss_index: t.Optional[faiss.Index] = None,
     ) -> None:
         super().__init__()
+        if faiss_index is None:
+            faiss_index = faiss.IndexFlatL2(BITS_IN_PDQ)
+        self.faiss_index = _PDQHashIndex(faiss_index)
         self.threshold = threshold
-
-        if index is None:
-            index = faiss.IndexFlatL2(BITS_IN_PDQ)
-        self._index = _PDQFaissIndex(index)
-
-        # Matches hash to Faiss index
         self._deduper: t.Dict[str, int] = {}
-        # Entry mapping: Each list[entries]'s index is its hash's index
         self._idx_to_entries: t.List[t.List[IndexT]] = []
 
-        self.add_all(entries=entries)
-
-    def __len__(self) -> int:
-        return len(self._idx_to_entries)
-
-    def query(self, hash: str) -> t.Sequence[PDQIndexMatch[IndexT]]:
-        """
-        Look up entries against the index, up to the threshold.
-        """
-        results: t.List[PDQIndexMatch[IndexT]] = []
-        matches_list: t.List[t.Tuple[int, int]] = self._index.search(
-            queries=[hash], threshold=self.threshold
-        )
-
-        for match, distance in matches_list:
-            entries = self._idx_to_entries[match]
-            # Create match objects for each entry
-            results.extend(
-                PDQIndexMatch(
-                    SignalSimilarityInfoWithIntDistance(distance=distance),
-                    entry,
-                )
-                for entry in entries
+    def query(self, query: str) -> t.Sequence[PDQIndexMatch[IndexT]]:
+        """Look up entries against the index."""
+        results = self.faiss_index.search([query], self.threshold)
+        return [
+            PDQIndexMatch(
+                SignalSimilarityInfoWithIntDistance(distance=int(distf)),
+                entry
             )
-        return results
+            for idx, distf in results[0]
+            for entry in self._idx_to_entries[idx]
+        ]
+
+    def _dedupe_and_add(self, signal_str: str, entry: IndexT, add_to_faiss: bool = True) -> None:
+        """Helper to handle deduplication and adding entries."""
+        existing_id = self._deduper.get(signal_str)
+        if existing_id is None:
+            next_id = len(self._deduper)
+            self._deduper[signal_str] = next_id
+            self._idx_to_entries.append([entry])
+            if add_to_faiss:
+                self.faiss_index.add(convert_pdq_strings_to_ndarray([signal_str]))
+        else:
+            self._idx_to_entries[existing_id].append(entry)
 
     def add(self, signal_str: str, entry: IndexT) -> None:
-        self.add_all(((signal_str, entry),))
+        """Add an entry to the index."""
+        self._dedupe_and_add(signal_str, entry)
 
     def add_all(self, entries: t.Iterable[t.Tuple[str, IndexT]]) -> None:
-        for h, i in entries:
-            existing_faiss_id = self._deduper.get(h)
-            if existing_faiss_id is None:
-                self._index.add([h])
-                self._idx_to_entries.append([i])
-                next_id = len(self._deduper)  # Because faiss index starts from 0 up
-                self._deduper[h] = next_id
-            else:
-                # Since this already exists, we don't add it to Faiss because Faiss cannot handle duplication
-                self._idx_to_entries[existing_faiss_id].append(i)
-
-
-class _PDQFaissIndex:
-    """
-    A wrapper around the faiss index for pickle serialization
-    """
-
-    def __init__(self, faiss_index: faiss.Index) -> None:
-        self.faiss_index = faiss_index
-
-    def add(self, pdq_strings: t.Sequence[str]) -> None:
-        """
-        Add PDQ hashes to the FAISS index.
-        """
-        vectors = convert_pdq_strings_to_ndarray(pdq_strings)
-        self.faiss_index.add(vectors)
+        """Add multiple entries to the index."""
+        for signal_str, entry in entries:
+            self._dedupe_and_add(signal_str, entry)
 
-    def search(
-        self, queries: t.Sequence[str], threshold: int
-    ) -> t.List[t.Tuple[int, int]]:
+    @classmethod
+    def build(cls: t.Type[Self], entries: t.Iterable[t.Tuple[str, IndexT]]) -> Self:
         """
-        Search the FAISS index for matches to the given PDQ queries.
+        Faiss has many potential options that we can choose based on the size of the index.
         """
-        query_array: np.ndarray = convert_pdq_strings_to_ndarray(queries)
-        limits, distances, indices = self.faiss_index.range_search(
-            query_array, threshold + 1
-        )
-
-        results: t.List[t.Tuple[int, int]] = []
-        for i in range(len(queries)):
-            matches = [idx.item() for idx in indices[limits[i] : limits[i + 1]]]
-            dists = [dist for dist in distances[limits[i] : limits[i + 1]]]
-            for j in range(len(matches)):
-                results.append((matches[j], dists[j]))
-        return results
-
-    def __getstate__(self):
-        return faiss.serialize_index(self.faiss_index)
+        entry_list = list(entries)
+        xn = len(entry_list)
+        if xn < 1024:  # If small enough, just use brute force
+            return cls(entries=entry_list)
+        
+        centroids = pick_n_centroids(xn)
+        index = faiss.IndexIVFFlat(faiss.IndexFlatL2(BITS_IN_PDQ), BITS_IN_PDQ, centroids)
+        index.nprobe = 16  # 16-64 should be high enough accuracy for 1-10M
+        index.cp.min_points_per_centroid = 1  # Allow small clusters
+        
+        ret = cls(faiss_index=index)
+        for signal_str, entry in entry_list:
+            ret._dedupe_and_add(signal_str, entry, add_to_faiss=False)
+            
+        xb = convert_pdq_strings_to_ndarray([s for s in ret._deduper])
+        ret.faiss_index.train(xb)
+        ret.faiss_index.add(xb)
+        return ret
 
-    def __setstate__(self, data):
-        self.faiss_index = faiss.deserialize_index(data)
+    def __len__(self) -> int:
+        return len(self._idx_to_entries)
 
 
 class PDQSignalTypeIndex2(SignalTypeIndex[IndexT]):