| """Embedding backends: OpenAI API or local sentence-transformers, with the |
| query/document prefixes each model family expects.""" |
| import numpy as np |
|
|
| PREFIX_RULES = [ |
| ("bge-", {"query": "Represent this sentence for searching relevant passages: ", |
| "doc": ""}), |
| ("e5-", {"query": "query: ", "doc": "passage: "}), |
| ("multilingual-e5", {"query": "query: ", "doc": "passage: "}), |
| ("Qwen3-Embedding", {"query": "Instruct: Given a figure search query, retrieve the " |
| "figure caption that matches it\nQuery: ", |
| "doc": ""}), |
| ] |
|
|
|
|
| def prefixes_for(model_name): |
| for needle, rules in PREFIX_RULES: |
| if needle.lower() in model_name.lower(): |
| return rules |
| return {"query": "", "doc": ""} |
|
|
|
|
| class OpenAIEmbedder: |
| name = "openai" |
|
|
| def __init__(self, model="text-embedding-3-small", dim=512): |
| import os |
| from openai import OpenAI |
| self.client = OpenAI(api_key=os.environ["OPENAI_API_KEY"], |
| max_retries=0, timeout=60) |
| self.model = model |
| self.dim = dim |
|
|
| def embed(self, texts, is_query=False, batch_size=128): |
| import time |
| out = [] |
| for i in range(0, len(texts), batch_size): |
| batch = texts[i:i + batch_size] |
| for attempt in range(8): |
| try: |
| resp = self.client.embeddings.create( |
| model=self.model, input=batch, dimensions=self.dim) |
| out.extend(d.embedding for d in resp.data) |
| break |
| except Exception as e: |
| if "429" not in str(e) or attempt == 7: |
| raise |
| time.sleep(10 * (attempt + 1)) |
| X = np.array(out, dtype=np.float32) |
| X /= np.linalg.norm(X, axis=1, keepdims=True) |
| return X |
|
|
| def info(self): |
| return {"backend": "openai", "model": self.model, "dim": self.dim} |
|
|
|
|
| class LocalEmbedder: |
| name = "local" |
|
|
| def __init__(self, model="BAAI/bge-base-en-v1.5", device=None): |
| from sentence_transformers import SentenceTransformer |
| self.model_name = model |
| self.model = SentenceTransformer(model, device=device) |
| self.prefixes = prefixes_for(model) |
| self.dim = self.model.get_sentence_embedding_dimension() |
|
|
| def embed(self, texts, is_query=False, batch_size=64): |
| prefix = self.prefixes["query" if is_query else "doc"] |
| inputs = [prefix + t for t in texts] |
| X = self.model.encode(inputs, batch_size=batch_size, |
| normalize_embeddings=True, |
| show_progress_bar=len(texts) > 1000) |
| return np.asarray(X, dtype=np.float32) |
|
|
| def info(self): |
| return {"backend": "local", "model": self.model_name, "dim": self.dim} |
|
|
|
|
| def make_embedder(backend, model=None, dim=512): |
| if backend == "openai": |
| return OpenAIEmbedder(model=model or "text-embedding-3-small", dim=dim) |
| if backend == "local": |
| return LocalEmbedder(model=model or "BAAI/bge-base-en-v1.5") |
| raise ValueError(f"unknown backend {backend}") |
|
|
|
|
| def embedder_from_info(info): |
| if info["backend"] == "openai": |
| return OpenAIEmbedder(model=info["model"], dim=info["dim"]) |
| return LocalEmbedder(model=info["model"]) |
|
|
