Update app.py

app.py

@@ -61,6 +61,56 @@ def build_index(n=200):
     INDEX["bands"] = bands
     return f"Index built: {len(ids)} images."
 
+def search(text_query, image_query, k=5):
+    if INDEX["feats"] is None:
+        return [], "Build the index first."
+
+    with torch.no_grad():
+        if text_query and text_query.strip():
+            inputs = processor(text=[text_query.strip()], return_tensors="pt")
+            q = model.get_text_features(**inputs)           # [1, 512]
+        elif image_query is not None:
+            pil = image_query.convert("RGB")
+            inputs = processor(images=pil, return_tensors="pt")
+            q = model.get_image_features(**inputs)          # [1, 512]
+        else:
+            return [], "Enter text or upload an image."
+
+        q = F.normalize(q, p=2, dim=-1)[0]                  # [512]
+        sims = (INDEX["feats"] @ q).cpu()                   # [N]
+        topk = torch.topk(sims, k=min(int(k), sims.shape[0]))
+
+        items = []
+        for idx in topk.indices.tolist():
+            cap = f"id: {INDEX['ids'][idx]}  score: {float(sims[idx]):.3f}  band: {INDEX['bands'][idx]}"
+            items.append((INDEX["thumbs"][idx], cap))
+        return items, f"Returned {len(items)} results."
+
+    def search(text_query, image_query, k=5):
+    if INDEX["feats"] is None:
+        return [], "Build the index first."
+
+    with torch.no_grad():
+        if text_query and text_query.strip():
+            inputs = processor(text=[text_query.strip()], return_tensors="pt")
+            q = model.get_text_features(**inputs)          # [1, 512]
+        elif image_query is not None:
+            pil = image_query.convert("RGB")
+            inputs = processor(images=pil, return_tensors="pt")
+            q = model.get_image_features(**inputs)         # [1, 512]
+        else:
+            return [], "Enter text or upload an image."
+
+        q = F.normalize(q, p=2, dim=-1)[0]                 # [512]
+        sims = (INDEX["feats"] @ q).cpu()                  # [N]
+        topk = torch.topk(sims, k=min(int(k), sims.shape[0]))
+
+        items = []
+        for idx in topk.indices.tolist():
+            cap = f"id: {INDEX['ids'][idx]}  score: {float(sims[idx]):.3f}  band: {INDEX['bands'][idx]}"
+            items.append((INDEX["thumbs"][idx], cap))
+        return items, f"Returned {len(items)} results."
+
 # ---------- UI ----------
 with gr.Blocks() as demo:
     gr.Markdown("JWST multimodal search — build the index")