Embedding 模型选型与优化实战：从 BGE 到 OpenAI Embedding

import numpy as np

def cosine_similarity(a, b):
    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))

def dot_product(a, b):
    return np.dot(a, b)

def l2_distance(a, b):
    return np.linalg.norm(a - b)

from openai import OpenAI

client = OpenAI()

def embed_texts(texts: list[str], model: str = "text-embedding-3-small") -> list[list[float]]:
    response = client.embeddings.create(
        input=texts,
        model=model,
        # 可以用 dimensions 参数降维，利用 MRL 技术
        # dimensions=512
    )
    return [item.embedding for item in response.data]

# text-embedding-3-small: 1536 维，$0.02/1M tokens
# text-embedding-3-large: 3072 维，$0.13/1M tokens

from FlagEmbedding import BGEM3FlagModel

model = BGEM3FlagModel('BAAI/bge-m3', use_fp16=True)

sentences = ["RAG 系统的检索增强原理", "如何优化向量检索性能"]

# BGE-M3 支持三种检索模式
embeddings = model.encode(
    sentences,
    batch_size=12,
    max_length=8192,
    return_dense=True,    # Dense 向量，用于语义相似度
    return_sparse=True,   # Sparse 权重，类似 BM25
    return_colbert_vecs=True  # ColBERT 多向量，精度最高但存储开销大
)

dense_vecs = embeddings['dense_vecs']    # shape: (2, 1024)
sparse_weights = embeddings['lexical_weights']  # dict: token -> weight
colbert_vecs = embeddings['colbert_vecs']  # shape: (2, seq_len, 128)

import requests

def jina_embed(texts: list[str], task: str = "retrieval.passage") -> list[list[float]]:
    """
    task 可选：
    - retrieval.query：查询侧
    - retrieval.passage：文档侧
    - text-matching：语义相似度
    - classification：分类
    - separation：聚类
    """
    url = "https://api.jina.ai/v1/embeddings"
    headers = {"Authorization": "Bearer YOUR_JINA_API_KEY"}
    payload = {
        "input": texts,
        "model": "jina-embeddings-v3",
        "task": task,
        "dimensions": 1024,
        "late_chunking": False  # 长文档可以开启，在 Embedding 层做分块
    }
    response = requests.post(url, headers=headers, json=payload)
    return [item["embedding"] for item in response.json()["data"]]

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("intfloat/e5-mistral-7b-instruct")

# 注意：e5 系列需要加前缀
query = "Instruct: Retrieve relevant passages for the query\nQuery: 什么是 RAG？"
passage = "passage: RAG（检索增强生成）是一种将向量检索与 LLM 生成相结合的技术..."

query_embedding = model.encode(query, normalize_embeddings=True)
passage_embedding = model.encode(passage, normalize_embeddings=True)

# 用 MTEB 库本地跑评测
import mteb

model = mteb.get_model("BAAI/bge-m3")
tasks = mteb.get_tasks(tasks=["T2Retrieval", "MMarcoRetrieval"], languages=["zho"])
evaluation = mteb.MTEB(tasks=tasks)
results = evaluation.run(model, output_folder="mteb_results")

# text-embedding-3-large 支持指定输出维度
response = client.embeddings.create(
    input=["测试文本"],
    model="text-embedding-3-large",
    dimensions=256  # 从 3072 降到 256，存储节省 12x
)

# 验证精度损失
import numpy as np
full_vec = embed_texts(["测试文本"], dimensions=3072)[0]
small_vec = embed_texts(["测试文本"], dimensions=256)[0]

# MRL 实现原理：直接截取前 N 维后重新归一化
truncated = np.array(full_vec[:256])
truncated = truncated / np.linalg.norm(truncated)

import hashlib
import json
import redis
import numpy as np
from typing import Optional

class EmbeddingCache:
    def __init__(self, redis_url: str, model_name: str, ttl: int = 86400 * 30):
        self.redis = redis.from_url(redis_url)
        self.model_name = model_name
        self.ttl = ttl  # 默认 30 天
    
    def _cache_key(self, text: str) -> str:
        # 包含 model_name 防止不同模型的向量混淆
        content = f"{self.model_name}:{text}"
        return f"emb:{hashlib.sha256(content.encode()).hexdigest()}"
    
    def get(self, text: str) -> Optional[list[float]]:
        key = self._cache_key(text)
        cached = self.redis.get(key)
        if cached:
            return json.loads(cached)
        return None
    
    def set(self, text: str, vector: list[float]) -> None:
        key = self._cache_key(text)
        self.redis.setex(key, self.ttl, json.dumps(vector))
    
    def get_or_embed(self, texts: list[str], embed_fn) -> list[list[float]]:
        results = [None] * len(texts)
        miss_indices = []
        miss_texts = []
        
        # 先查缓存
        for i, text in enumerate(texts):
            cached = self.get(text)
            if cached is not None:
                results[i] = cached
            else:
                miss_indices.append(i)
                miss_texts.append(text)
        
        # 批量 Embed 未命中的
        if miss_texts:
            new_vectors = embed_fn(miss_texts)
            for i, (idx, vec) in enumerate(zip(miss_indices, new_vectors)):
                results[idx] = vec
                self.set(miss_texts[i], vec)
        
        return results


# 使用示例
cache = EmbeddingCache(redis_url="redis://localhost:6379", model_name="text-embedding-3-small")

def embed_with_cache(texts: list[str]) -> list[list[float]]:
    return cache.get_or_embed(texts, lambda t: embed_texts(t))

import asyncio
from openai import AsyncOpenAI

async_client = AsyncOpenAI()

async def batch_embed_async(
    texts: list[str],
    model: str = "text-embedding-3-small",
    batch_size: int = 100,
    max_concurrent: int = 5
) -> list[list[float]]:
    """
    批量异步 Embedding，控制并发数避免触发 rate limit
    """
    semaphore = asyncio.Semaphore(max_concurrent)
    
    async def embed_batch(batch: list[str]) -> list[list[float]]:
        async with semaphore:
            response = await async_client.embeddings.create(
                input=batch,
                model=model
            )
            return [item.embedding for item in response.data]
    
    # 分批
    batches = [texts[i:i+batch_size] for i in range(0, len(texts), batch_size)]
    tasks = [embed_batch(batch) for batch in batches]
    
    # 带重试的执行
    batch_results = await asyncio.gather(*tasks, return_exceptions=True)
    
    # 展平结果
    all_vectors = []
    for result in batch_results:
        if isinstance(result, Exception):
            raise result
        all_vectors.extend(result)
    
    return all_vectors


# 同步入口
def embed_large_corpus(texts: list[str]) -> list[list[float]]:
    return asyncio.run(batch_embed_async(texts))