| """Commerce query encoder (Korean + English) for EmbeddingGemma 2 product-image indexes. |
| |
| Search text -> L2-normalized 768-d vector that you compare (dot product) with product-image |
| vectors produced by google/embeddinggemma-2. The image index does not change. |
| |
| Two INT8 text encoders share one tokenizer: |
| ko/ K4 - trained for Korean product search (77 tokens; long inputs keep the first 57 and last 20 tokens) |
| en/ E0 - English product search (64 tokens, plain truncation) |
| route="hangul" (default) sends a query that contains Hangul to ko/ and every other query to en/. |
| route="ko" or route="en" uses one encoder for everything. |
| |
| Weights are row-wise INT8 (`<name>.__scale` holds one fp16 scale per row), dequantized to float32 at load. |
| """ |
| import json |
| import re |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from safetensors.torch import load_file |
| from transformers import AutoTokenizer, CLIPTextConfig, CLIPTextModel |
|
|
| HANGUL = re.compile(r"[가-힣ㄱ-ㆎ]") |
|
|
|
|
| def _resolve(path_or_repo, revision=None): |
| p = Path(path_or_repo) |
| if p.is_dir(): |
| return p |
| from huggingface_hub import snapshot_download |
| return Path(snapshot_download(str(path_or_repo), revision=revision)) |
|
|
|
|
| def head_tail(ids, max_len, tail): |
| """Keep the first max_len - tail and the last tail token ids (BOS and EOS survive).""" |
| ids = list(ids) |
| return ids if len(ids) <= max_len else ids[:max_len - tail] + ids[-tail:] |
|
|
|
|
| class TinyQueryEncoder(nn.Module): |
| def __init__(self, text_config, out_dim=768, max_tokens=64, tail_tokens=0): |
| super().__init__() |
| self.enc = CLIPTextModel(text_config) |
| self.align = nn.Linear(text_config.hidden_size, out_dim) |
| self.max_tokens = min(max_tokens, 77) |
| self.tail_tokens = tail_tokens |
| self.tok = None |
|
|
| @classmethod |
| def load(cls, folder, tokenizer): |
| folder = Path(folder) |
| meta = json.loads((folder / "config.json").read_text()) |
| base = json.loads((folder / "base_config.json").read_text()) |
| text_cfg = CLIPTextConfig.from_dict(base["text_config"] if base.get("model_type") == "clip" else base) |
| tail = (meta.get("truncation") or {}).get("tail_tokens", 0) |
| m = cls(text_cfg, out_dim=meta["out_dim"], max_tokens=meta["max_tokens"], tail_tokens=tail) |
| raw = load_file(str(folder / "model.safetensors")) |
| sd = {k: (v.float() * raw[f"{k}.__scale"].float()[:, None]) if f"{k}.__scale" in raw else v.float() |
| for k, v in raw.items() if not k.endswith(".__scale")} |
| missing, unexpected = m.load_state_dict(sd, strict=False) |
| ref = m.state_dict() |
| if unexpected or any(ref[k].dtype.is_floating_point for k in missing): |
| raise RuntimeError(f"weight mismatch in {folder}: missing={missing} unexpected={unexpected}") |
| m.tok = tokenizer |
| return m.eval() |
|
|
| def _batch(self, texts): |
| if not self.tail_tokens: |
| return self.tok(texts, padding=True, truncation=True, max_length=self.max_tokens, return_tensors="pt") |
| ids = [head_tail(x, self.max_tokens, self.tail_tokens) for x in self.tok(texts, truncation=False)["input_ids"]] |
| n = max(len(x) for x in ids) |
| a = torch.full((len(ids), n), int(self.tok.pad_token_id or 0), dtype=torch.long) |
| m = torch.zeros(len(ids), n, dtype=torch.long) |
| for i, x in enumerate(ids): |
| a[i, :len(x)] = torch.tensor(x) |
| m[i, :len(x)] = 1 |
| return {"input_ids": a, "attention_mask": m} |
|
|
| @torch.no_grad() |
| def encode(self, texts, batch_size=256): |
| dev = next(self.parameters()).device |
| out = [] |
| for i in range(0, len(texts), batch_size): |
| b = {k: v.to(dev) for k, v in self._batch(list(texts[i:i + batch_size])).items()} |
| out.append(F.normalize(self.align(self.enc(**b).pooler_output).float(), dim=1).cpu().numpy()) |
| return np.concatenate(out) if out else np.zeros((0, 768), np.float32) |
|
|
|
|
| class CommerceQueryEncoder: |
| def __init__(self, ko, en, route="hangul"): |
| self.ko, self.en, self.route = ko, en, route |
|
|
| @classmethod |
| def from_pretrained(cls, path_or_repo="ThakiCloud/eg2-tiny-query-commerce-ko", revision=None, device="cpu", route="hangul"): |
| if route not in ("hangul", "ko", "en"): |
| raise ValueError("route must be 'hangul', 'ko' or 'en'") |
| root = _resolve(path_or_repo, revision) |
| tok = AutoTokenizer.from_pretrained(str(root)) |
| return cls(TinyQueryEncoder.load(root / "ko", tok).to(device), TinyQueryEncoder.load(root / "en", tok).to(device), route) |
|
|
| def encode(self, texts, batch_size=256): |
| """texts: list[str] -> np.ndarray (n, 768), float32, L2-normalized. No query prefix.""" |
| if isinstance(texts, str): |
| texts = [texts] |
| texts = list(texts) |
| if self.route != "hangul": |
| return (self.ko if self.route == "ko" else self.en).encode(texts, batch_size) |
| is_ko = [bool(HANGUL.search(t)) for t in texts] |
| out = np.zeros((len(texts), 768), np.float32) |
| for flag, model in ((True, self.ko), (False, self.en)): |
| idx = [i for i, f in enumerate(is_ko) if f == flag] |
| if idx: |
| out[idx] = model.encode([texts[i] for i in idx], batch_size) |
| return out |
|
|
