BBQ 量化如何将 Jina v5 embedding 缩小 29 倍，同时在 Elasticsearch 中保持召回率

from elasticsearch import Elasticsearch, helpers

# 假设 ELASTICSEARCH_URL 和 ELASTICSEARCH_API_KEY 已经定义
es_client = Elasticsearch(
    ELASTICSEARCH_URL,
    api_key=ELASTICSEARCH_API_KEY,
    request_timeout=120
)
es_client.info()

DIMS = 1024
FLOAT_INDEX = "vectors-float32"
BBQ_INDEX = "vectors-bbq"

def create_index(name, index_options):
    if es_client.indices.exists(index=name):
        es_client.indices.delete(index=name)
    es_client.indices.create(
        index=name,
        mappings={
            "properties": {
                "text": {"type": "text"},
                "lang": {"type": "keyword"},
                "embedding": {
                    "type": "dense_vector",
                    "dims": DIMS,
                    "index": True,
                    "similarity": "cosine",
                    "index_options": index_options,
                },
            }
        },
    )

create_index(FLOAT_INDEX, {"type": "hnsw"})    # 原始 float32 基线
create_index(BBQ_INDEX,   {"type": "bbq_hnsw"})  # 1 比特 BBQ

import numpy as np # 导入 numpy 库

INFERENCE_ID = ".jina-embeddings-v5-text-small"

def embed(texts, batch_size=16):
    out = []
    for i in range(0, len(texts), batch_size):
        batch = texts[i : i + batch_size]
        try:
            resp = es_client.inference.text_embedding(
                inference_id=INFERENCE_ID, input=batch
            )
        except AttributeError:  # 针对旧版客户端
            resp = es_client.inference.inference(inference_id=INFERENCE_ID, input=batch)
        out.extend(item["embedding"] for item in resp["text_embedding"])
    return np.array(out, dtype=np.float32)

embed(["hello world"]).shape # 测试

(1, 1024)

from datasets import load_dataset

LANGS = ["en", "de", "ja", "pt", "ru"]
PER_LANG_DOCS = 1000
PER_LANG_QUERIES = 20

docs, queries = [], []
for lang in LANGS:
    ds = load_dataset(
        "hotchpotch/multilingual_cc_news", lang, split="train", streaming=True
    )
    rows = [
        r
        for r in ds.take(PER_LANG_DOCS + PER_LANG_QUERIES)
        if r.get("maintext") and r.get("title")
    ]

    for row in rows[:PER_LANG_DOCS]:
        text = (row["title"] + ". " + row["maintext"]).replace("\n", " ").strip()
        docs.append({"text": text[:1000], "lang": lang})

    for row in rows[PER_LANG_DOCS:]:
        queries.append({"text": row["title"], "lang": lang})  # 标题作为查询

print(f"语料库: {len(docs)} 篇文档 | 查询: {len(queries)} 个")
# RES: 语料库: 3102 篇文档 | 查询: 18 个

doc_vectors = embed([d["text"] for d in docs])
query_vectors = embed([q["text"] for q in queries])

def index_docs(name):
    actions = (
        {
            "_index": name,
            "_id": i,
            "_source": {
                "text": d["text"],
                "lang": d["lang"],
                "embedding": doc_vectors[i].tolist(),
            },
        }
        for i, d in enumerate(docs)
    )
    helpers.bulk(es_client, actions, refresh=True)

for name in (FLOAT_INDEX, BBQ_INDEX):
    index_docs(name)
    es_client.indices.forcemerge(index=name, max_num_segments=1)
    es_client.indices.refresh(index=name)

def vector_disk_bytes(name):
    du = es_client.indices.disk_usage(index=name, run_expensive_tasks=True)
    field = du[name]["fields"]["embedding"]
    knn = field.get("knn_vectors")
    if isinstance(knn, dict):
        return knn["size_in_bytes"]
    return field["knn_vectors_in_bytes"]

float_disk = vector_disk_bytes(FLOAT_INDEX)
bbq_disk = vector_disk_bytes(BBQ_INDEX)

N = len(docs)
float_mem = N * DIMS * 4
bbq_mem = N * (DIMS // 8 + 14)

print(f"磁盘占用 -> float32: {float_disk/1e6:6.2f} MB | BBQ: {bbq_disk/1e6:6.2f} MB")
print(f"内存占用 -> float32: {float_mem/1e6:6.2f} MB | BBQ: {bbq_mem/1e6:6.2f} MB  ({float_mem/bbq_mem:.0f} 倍缩小)")

磁盘占用 -> float32:  12.80 MB | BBQ:  13.25 MB
内存占用 -> float32:  12.71 MB | BBQ:   0.44 MB  (29 倍缩小)

def search_ids(index, qvec, k=10, num_candidates=10):
    resp = es_client.search(
        index=index,
        size=k,
        _source=False,
        knn={
            "field": "embedding",
            "query_vector": qvec.tolist(),
            "k": k,
            "num_candidates": num_candidates,
        },
    )
    return [h["_id"] for h in resp["hits"]["hits"]]

K = 10
# 真实情况（Ground truth）：使用全精度 float32 和广泛的候选列表（近似精确）
ground_truth = [
    set(search_ids(FLOAT_INDEX, qv, k=K, num_candidates=2000)) for qv in query_vectors
]

oversamples = [1, 2, 3, 5, 10]
recalls = []
for f in oversamples:
    num_candidates = max(K * f, K)
    hits = 0
    for gt, qv in zip(ground_truth, query_vectors):
        got = set(search_ids(BBQ_INDEX, qv, k=K, num_candidates=num_candidates))
        hits += len(got & gt)
    recalls.append(hits / (len(query_vectors) * K))
    print(f"过采样 {f:>2}x -> recall@{K} = {recalls[-1]:.3f}")