Create build_index.py

build_index.py

@@ -0,0 +1,94 @@
+# build_index.py
+import os
+import json
+import numpy as np
+import torch
+import torch.nn.functional as F
+from tqdm import tqdm
+from datasets import concatenate_datasets, load_dataset
+from transformers import AutoProcessor, AutoModel
+from PIL import Image
+
+# ========== CONFIG ==========
+MODEL_ID = "EYEDOL/siglipFULL-agri-finetuned"
+DATASET_NAMES = [f"EYEDOL/AGRILLAVA-image-text{i}" for i in range(1, 16)]
+BATCH_SIZE = 16
+OUT_DIR = "faiss_data"
+INDEX_FILE = os.path.join(OUT_DIR, "texts.faiss")
+METADATA_FILE = os.path.join(OUT_DIR, "texts_meta.json")
+EMBEDS_FILE = os.path.join(OUT_DIR, "text_embeds.npy")
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+USE_FAISS_GPU = False  # set True if you have faiss-gpu and want GPU index building
+# ============================
+
+os.makedirs(OUT_DIR, exist_ok=True)
+
+print("Loading datasets...")
+all_splits = [load_dataset(name)["train"] for name in DATASET_NAMES]
+dataset = concatenate_datasets(all_splits)
+texts = list(dataset["text"])
+print(f"Total texts: {len(texts)}")
+
+print("Loading model & processor...")
+processor = AutoProcessor.from_pretrained(MODEL_ID)
+model = AutoModel.from_pretrained(MODEL_ID).to(DEVICE)
+model.eval()
+
+# Compute text embeddings (in batches) and normalize
+all_embeds = []
+for i in tqdm(range(0, len(texts), BATCH_SIZE), desc="Encoding texts"):
+    batch_texts = texts[i:i + BATCH_SIZE]
+    inputs = processor(text=batch_texts, padding=True, truncation=True, return_tensors="pt").to(DEVICE)
+    with torch.no_grad():
+        text_embeds = model.get_text_features(**inputs)  # shape (bs, dim)
+        text_embeds = text_embeds / text_embeds.norm(p=2, dim=-1, keepdim=True)
+    all_embeds.append(text_embeds.cpu())
+    del inputs, text_embeds
+    if DEVICE == "cuda":
+        torch.cuda.empty_cache()
+
+all_embeds = torch.cat(all_embeds, dim=0).numpy().astype(np.float32)  # (N, D)
+print("Embeddings shape:", all_embeds.shape)
+
+# Save numpy embeddings (optional, useful for re-indexing)
+np.save(EMBEDS_FILE, all_embeds)
+print(f"Saved embeddings to {EMBEDS_FILE}")
+
+# ========== Build FAISS index ==========
+try:
+    import faiss  # type: ignore
+except Exception as e:
+    raise ImportError("Please install faiss (faiss-cpu or faiss-gpu).") from e
+
+d = all_embeds.shape[1]
+# We'll use inner-product on L2-normalized vectors to get cosine similarity.
+index = faiss.IndexFlatIP(d)  # exact index; change if you want IVF/HNSW for large corpora
+
+# If you have faiss-gpu and want to move to GPU:
+if USE_FAISS_GPU:
+    res = faiss.StandardGpuResources()
+    index = faiss.index_cpu_to_gpu(res, 0, index)
+
+print("Adding to index...")
+index.add(all_embeds)  # adds vectors in order, ids 0..N-1
+print("Index ntotal:", index.ntotal)
+
+# Save index
+faiss.write_index(faiss.index_gpu_to_cpu(index) if USE_FAISS_GPU else index, INDEX_FILE)
+print(f"Saved FAISS index to {INDEX_FILE}")
+
+# Save metadata (texts)
+meta = {
+    "count": len(texts),
+    "texts_file": "texts.jsonl"
+}
+with open(METADATA_FILE, "w") as f:
+    json.dump(meta, f, indent=2)
+# Save texts list (one-per-line) for easy lookup
+texts_file = os.path.join(OUT_DIR, "texts.jsonl")
+with open(texts_file, "w", encoding="utf-8") as f:
+    for t in texts:
+        f.write(json.dumps({"text": t}, ensure_ascii=False) + "\n")
+
+print(f"Saved texts metadata to {texts_file}")
+print("Done.")