| """Product-image search with the commerce query encoder, alone or fused with BM25. |
| |
| 1) Build an EmbeddingGemma 2 index from a folder of product images (skip if you already have one): |
| python example_search.py build --images ./products --index index.npz |
| 2) Search it with Korean or English queries: |
| python example_search.py search --index index.npz --query "무선 블루투스 이어폰" --query "stainless steel water bottle" |
| 3) Optional: fuse with a BM25 keyword index over your own product titles (one JSON line per image: |
| {"path": "<same path as in the index>", "text": "<title, brand, attributes>"}): |
| python example_search.py search --index index.npz --catalog catalog.jsonl --query "나이키 러닝화 270mm" |
| |
| Fusion is reciprocal-rank fusion (k=60) of the two top-10 lists. In our tests it raised R@10 on some |
| query sets but lowered R@1 whenever BM25 alone was already strong; check it on your own data. |
| """ |
| import argparse |
| import json |
| import math |
| import re |
| from collections import Counter |
| from pathlib import Path |
|
|
| import numpy as np |
|
|
| from commerce_query import CommerceQueryEncoder |
|
|
| EXTS = {".jpg", ".jpeg", ".png", ".webp", ".bmp"} |
| REPO = "ThakiCloud/eg2-tiny-query-commerce-ko" |
|
|
|
|
| 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 tokens(s): |
| return re.findall(r"\w+", (s or "").lower()) |
|
|
|
|
| class BM25: |
| """Okapi BM25 (k1=1.5, b=0.75), whitespace/word tokens. Small and dependency-free.""" |
|
|
| def __init__(self, docs, k1=1.5, b=0.75): |
| self.docs = [Counter(tokens(d)) for d in docs] |
| self.len = np.array([sum(d.values()) for d in self.docs], float) |
| self.avg = float(self.len.mean()) if len(self.docs) else 0.0 |
| df = Counter(t for d in self.docs for t in d) |
| n = len(self.docs) |
| self.idf = {t: math.log((n - f + 0.5) / (f + 0.5) + 1) for t, f in df.items()} |
| self.k1, self.b = k1, b |
|
|
| def scores(self, query): |
| s = np.zeros(len(self.docs)) |
| for t in set(tokens(query)): |
| if t not in self.idf: |
| continue |
| tf = np.array([d.get(t, 0) for d in self.docs], float) |
| s += self.idf[t] * tf * (self.k1 + 1) / (tf + self.k1 * (1 - self.b + self.b * self.len / max(self.avg, 1e-9))) |
| return s |
|
|
|
|
| def rrf(*ranked, k=60, top=10): |
| sc = {} |
| for r in ranked: |
| for pos, d in enumerate(r): |
| sc[int(d)] = sc.get(int(d), 0.0) + 1.0 / (k + pos + 1) |
| return [d for d, _ in sorted(sc.items(), key=lambda x: -x[1])[:top]] |
|
|
|
|
| def search(index, queries, k=5, model=REPO, catalog=None, route="hangul"): |
| d = np.load(index) |
| paths = [str(p) for p in d["paths"]] |
| enc = CommerceQueryEncoder.from_pretrained(model, route=route) |
| q = enc.encode(queries) |
| dense = np.argsort(-(q @ d["vectors"].T), axis=1)[:, :10] |
| bm = None |
| if catalog: |
| text = {} |
| for line in open(catalog, encoding="utf-8"): |
| r = json.loads(line) |
| text[str(r["path"])] = r.get("text", "") |
| bm = BM25([text.get(p, "") for p in paths]) |
| results = [] |
| for i, t in enumerate(queries): |
| top = list(dense[i]) |
| if bm is not None: |
| top = rrf(np.argsort(-bm.scores(t))[:10], top) |
| results.append(top) |
| print(f"\n{t}" + (" [model + BM25]" if bm is not None else " [model]")) |
| for j in top[:k]: |
| print(f" {paths[j]}") |
| return results |
|
|
|
|
| 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=REPO) |
| s.add_argument("--catalog"); s.add_argument("--route", default="hangul", choices=["hangul", "ko", "en"]) |
| 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, a.catalog, a.route) |
|
|
