File size: 4,693 Bytes
d992912 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 | import logging
from typing import List, Optional, Tuple
import numpy as np
import faiss
from backend.app.config import SearchConfig
logger = logging.getLogger(__name__)
class DualFAISSIndex:
"""
Two parallel FAISS indices (image + text) fused via Reciprocal Rank Fusion.
"""
def __init__(self, dim: int, config: SearchConfig):
self.dim = dim
self.config = config
self.image_index = None
self.text_index = None
def _create_index(self, n_vectors: int) -> faiss.Index:
if n_vectors < 5000:
logger.info(f"Using IndexFlatIP (exact, n={n_vectors:,})")
return faiss.IndexFlatIP(self.dim)
n_clusters = min(self.config.n_clusters, max(16, n_vectors // 40))
logger.info(f"Using IndexIVFFlat (n={n_vectors:,}, clusters={n_clusters})")
quantizer = faiss.IndexFlatIP(self.dim)
index = faiss.IndexIVFFlat(
quantizer, self.dim, n_clusters, faiss.METRIC_INNER_PRODUCT
)
return index
def build(self, image_embeddings: np.ndarray, text_embeddings: np.ndarray):
image_embeddings = image_embeddings.astype(np.float32)
text_embeddings = text_embeddings.astype(np.float32)
assert image_embeddings.shape == text_embeddings.shape, (
f"Shape mismatch: images {image_embeddings.shape} vs text {text_embeddings.shape}"
)
n = image_embeddings.shape[0]
logger.info("Building image FAISS index...")
self.image_index = self._create_index(n)
if hasattr(self.image_index, 'train'):
try:
self.image_index.train(image_embeddings)
except Exception:
self.image_index = faiss.IndexFlatIP(self.dim)
self.image_index.add(image_embeddings)
logger.info("Building text FAISS index...")
self.text_index = self._create_index(n)
if hasattr(self.text_index, 'train'):
try:
self.text_index.train(text_embeddings)
except Exception:
self.text_index = faiss.IndexFlatIP(self.dim)
self.text_index.add(text_embeddings)
logger.info(
f"Dual index built: {self.image_index.ntotal:,} image, "
f"{self.text_index.ntotal:,} text vectors"
)
def search_image_index(self, query: np.ndarray, top_k: int):
q = query.astype(np.float32).reshape(1, -1)
if hasattr(self.image_index, 'nprobe'):
self.image_index.nprobe = self.config.n_probe
return self.image_index.search(q, top_k)
def search_text_index(self, query: np.ndarray, top_k: int):
q = query.astype(np.float32).reshape(1, -1)
if hasattr(self.text_index, 'nprobe'):
self.text_index.nprobe = self.config.n_probe
return self.text_index.search(q, top_k)
def search_fused(
self, query: np.ndarray, top_k: int,
image_weight: Optional[float] = None,
text_weight: Optional[float] = None,
) -> Tuple[List[int], List[float]]:
iw = image_weight or self.config.image_index_weight
tw = text_weight or self.config.text_index_weight
rrf_k = self.config.rrf_k
broad_k = min(top_k * 3, self.image_index.ntotal)
_, img_ids = self.search_image_index(query, broad_k)
_, txt_ids = self.search_text_index(query, broad_k)
img_ids = img_ids[0]
txt_ids = txt_ids[0]
img_rank = {int(idx): rank + 1 for rank, idx in enumerate(img_ids) if idx >= 0}
txt_rank = {int(idx): rank + 1 for rank, idx in enumerate(txt_ids) if idx >= 0}
all_candidates = set(img_rank.keys()) | set(txt_rank.keys())
scores = {}
for idx in all_candidates:
score = 0.0
if idx in img_rank:
score += iw / (rrf_k + img_rank[idx])
if idx in txt_rank:
score += tw / (rrf_k + txt_rank[idx])
scores[idx] = score
ranked = sorted(scores.items(), key=lambda x: -x[1])[:top_k]
return [r[0] for r in ranked], [r[1] for r in ranked]
def save(self, image_path: str, text_path: str):
faiss.write_index(self.image_index, image_path)
faiss.write_index(self.text_index, text_path)
logger.info(f"Saved dual index to {image_path} and {text_path}")
def load(self, image_path: str, text_path: str):
self.image_index = faiss.read_index(image_path)
self.text_index = faiss.read_index(text_path)
logger.info(
f"Loaded dual index: {self.image_index.ntotal:,} image, "
f"{self.text_index.ntotal:,} text vectors"
)
__all__ = ["DualFAISSIndex"]
|