Spaces:
Paused
Paused
| import json | |
| import sys | |
| from pathlib import Path | |
| import numpy as np | |
| import open_clip | |
| import torch | |
| # --- Settings (must match build_index.py) --- | |
| EMB_DIR = Path("embeddings") | |
| MODEL_NAME = "ViT-B-32" | |
| PRETRAINED = "laion2b_s34b_b79k" | |
| # --- Load everything once, when the module is imported --- | |
| model, _, _ = open_clip.create_model_and_transforms(MODEL_NAME, pretrained=PRETRAINED) | |
| model.eval() | |
| tokenizer = open_clip.get_tokenizer(MODEL_NAME) | |
| embeddings = np.load(EMB_DIR / "image_embeddings.npy") # (1000, 512) | |
| filenames = json.loads((EMB_DIR / "filenames.json").read_text()) # 1000 names | |
| def embed_text(query: str) -> np.ndarray: | |
| """Turn a text query into a normalized 512-number vector.""" | |
| tokens = tokenizer([query]) # batch of 1 query | |
| with torch.no_grad(): | |
| vec = model.encode_text(tokens) # (1, 512) | |
| vec = vec / vec.norm(dim=-1, keepdim=True) | |
| return vec[0].cpu().numpy() # (512,) | |
| def search(query: str, k: int = 5) -> list[tuple[str, float]]: | |
| """Return the top-k (filename, score) pairs, best first.""" | |
| q = embed_text(query) | |
| scores = embeddings @ q # (1000, 512) @ (512,) -> (1000,) | |
| top = np.argsort(scores)[::-1][:k] # positions of the k highest scores | |
| return [(filenames[i], float(scores[i])) for i in top] | |
| if __name__ == "__main__": | |
| query = " ".join(sys.argv[1:]) or "a dog running on the beach" | |
| for rank, (name, score) in enumerate(search(query, k=5), start=1): | |
| print(f"{rank}. {name} score={score:.3f}") |