Update build_index.py

build_index.py

@@ -1,94 +1,60 @@
-# build_index.py
+# build_index_no_faiss.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 ==========
+# 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")
+OUT_DIR = "faiss_free_data"   # folder to upload into your Space
 EMBEDS_FILE = os.path.join(OUT_DIR, "text_embeds.npy")
+TEXTS_JSONL = os.path.join(OUT_DIR, "texts.jsonl")
 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)
 
+# Load datasets and concat
 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(f"Got {len(texts)} texts.")
 
+# Load model & processor
 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
+# Compute text embeddings in batches and L2-normalize (helps cosine)
 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)
+    batch = texts[i:i+BATCH_SIZE]
+    inputs = processor(text=batch, 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
+        embeds = model.get_text_features(**inputs)  # (bs, dim)
+        # normalize
+        embeds = embeds / embeds.norm(p=2, dim=-1, keepdim=True)
+    all_embeds.append(embeds.cpu().numpy().astype("float32"))
+    del inputs, embeds
     if DEVICE == "cuda":
         torch.cuda.empty_cache()
 
-all_embeds = torch.cat(all_embeds, dim=0).numpy().astype(np.float32)  # (N, D)
+all_embeds = np.concatenate(all_embeds, axis=0)  # shape (N, D)
 print("Embeddings shape:", all_embeds.shape)
 
-# Save numpy embeddings (optional, useful for re-indexing)
+# Save embeddings and texts mapping
 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:
+with open(TEXTS_JSONL, "w", encoding="utf-8") as f:
     for t in texts:
         f.write(json.dumps({"text": t}, ensure_ascii=False) + "\n")
+print(f"Saved texts to {TEXTS_JSONL}")
 
-print(f"Saved texts metadata to {texts_file}")
-print("Done.")
+print("Done. Upload the folder 'faiss_free_data' to your Space repository (git lfs or upload_file).")