| """Search an EmbeddingGemma 2 image index with A1 queries. |
| |
| Build an index from a folder of images with EmbeddingGemma 2 (skip this if you already have one): |
| python example_search.py build --images ./images --index index.npz |
| Search it with Korean or English queries encoded by A1: |
| python example_search.py search --index index.npz --query "월별 매출 추이" --query "login screen" |
| """ |
| import argparse |
| from pathlib import Path |
|
|
| import numpy as np |
|
|
| from a1_query import A1QueryEncoder |
|
|
| EXTS = {".jpg", ".jpeg", ".png", ".webp", ".bmp"} |
|
|
|
|
| def build(images, index, eg2="google/embeddinggemma-2", batch=16): |
| from PIL import Image |
| from sentence_transformers import SentenceTransformer |
| paths = sorted(p for p in Path(images).rglob("*") if p.suffix.lower() in EXTS) |
| st = SentenceTransformer(eg2, device="cpu") |
| vecs = [] |
| for i in range(0, len(paths), batch): |
| ims = [{"image": Image.open(p).convert("RGB")} for p in paths[i:i + batch]] |
| vecs.append(st.encode(ims, batch_size=batch, convert_to_numpy=True, normalize_embeddings=True)) |
| np.savez(index, vectors=np.concatenate(vecs).astype(np.float32), paths=np.array([str(p) for p in paths])) |
| print(f"indexed {len(paths)} images -> {index}") |
|
|
|
|
| def search(index, queries, k=5, model="ThakiCloud/eg2-tiny-query-a1"): |
| d = np.load(index) |
| enc = A1QueryEncoder.from_pretrained(model) |
| q = enc.encode(queries) |
| scores = q @ d["vectors"].T |
| for text, s in zip(queries, scores): |
| print(f"\n{text}") |
| for j in np.argsort(-s)[:k]: |
| print(f" {s[j]:.4f} {d['paths'][j]}") |
| return scores |
|
|
|
|
| if __name__ == "__main__": |
| ap = argparse.ArgumentParser() |
| sub = ap.add_subparsers(dest="cmd", required=True) |
| b = sub.add_parser("build"); b.add_argument("--images", required=True); b.add_argument("--index", required=True) |
| b.add_argument("--eg2", default="google/embeddinggemma-2") |
| s = sub.add_parser("search"); s.add_argument("--index", required=True); s.add_argument("--query", action="append", required=True) |
| s.add_argument("--k", type=int, default=5); s.add_argument("--model", default="ThakiCloud/eg2-tiny-query-a1") |
| a = ap.parse_args() |
| build(a.images, a.index, a.eg2) if a.cmd == "build" else search(a.index, a.query, a.k, a.model) |
|
|
