1
# Transformer 简化的前向传播
2
class TransformerEmbedding:
3
    def __init__(self, model_name):
4
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
5
        self.model = AutoModel.from_pretrained(model_name)
6

7
    def encode(self, texts):
8
        # 1. Tokenize
9
        tokens = self.tokenizer(
10
            texts,
11
            padding=True,
12
            truncation=True,
13
            max_length=512
14
        )
15

16
        # 2. Forward
17
        outputs = self.model(**tokens)
18

19
        # 3. Pooling
20
        embeddings = self.mean_pooling(outputs, tokens['attention_mask'])
21
        return F.normalize(embeddings, p=2)

1
# Mean Pooling
2
def mean_pooling(output, attention_mask):
3
    # 掩盖 padding token
4
    input_mask_expanded = attention_mask.unsqueeze(-1).expand(output.size()).float()
5
    return torch.sum(output * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)
6

7
# CLS Pooling（使用 [CLS] 向量）
8
def cls_pooling(output):
9
    return output[:, 0]  # 取第一个 token

1
# 暴力检索
2
def brute_search(query_vector, vectors, top_k):
3
    scores = cosine_similarity(query_vector, vectors)
4
    return top_k_indices(scores, top_k)
5
# 时间复杂度 O(n*d)
6

7
# IVF 索引
8
def ivf_search(query_vector, vectors, top_k):
9
    # 1. 找到最近的聚类中心
10
    center = find_nearest_center(query_vector)
11

12
    # 2. 在中心内搜索
13
    candidates = vectors[center_points[center]]
14
    scores = cosine_similarity(query_vector, candidates)
15
    return top_k(scores, top_k)
16
# 时间复杂度 O(sqrt(n))

1
# HNSW 参数
2
class HNSWIndex:
3
    def __init__(self):
4
        self.ef_construction = 200  # 建图精度
5
        self.M = 16  # 每个节点连接数
6
        self.ef_search = 100  # 搜索精度

1
# Milvus Collection 配置
2
collection_schema:
3
  name: documents
4
  fields:
5
    - name: id
6
      type: int64
7
      is_primary: true
8
    - name: embedding
9
      type: FLOAT32
10
      dim: 1536
11
      index:
12
        type: HNSW
13
        params:
14
          M: 16
15
          efConstruction: 200
16
    - name: text
17
      type: varchar
18
      max_length: 1000

1
# Pinecone SDK
2
import pinecone
3

4
# 连接
5
pinecone.init(api_key="...", environment="us-east1")
6
index = pinecone.Index("documents")
7

8
# 插入
9
index.upsert([
10
    ("1", [0.1] * 1536, {"text": "文档1"}),
11
    ("2", [0.2] * 1536, {"text": "文档2"})
12
])
13

14
# 检索
15
results = index.query(
16
    vector=[0.1] * 1536,
17
    top_k=10,
18
    include_metadata=True
19
)

1
def cosine_similarity(a, b):
2
    """余弦相似度：-1 到 1，1 表示完全相同"""
3
    dot_product = sum(x * y for x, y in zip(a, b))
4
    norm_a = sqrt(sum(x * x for x in a))
5
    norm_b = sqrt(sum(x * x for x in b))
6
    return dot_product / (norm_a * norm_b)

1
def dot_product(a, b):
2
    """点积：值域无限，需要归一化"""
3
    return sum(x * y for x, y in zip(a, b))
4

5
# 归一化后点积等价于余弦相似度
6
normalized_a = a / norm(a)
7
normalized_b = b / norm(b)
8
dot(normalized_a, normalized_b) == cosine_similarity(a, b)

1
def euclidean_distance(a, b):
2
    """欧几里得距离：值域 0 到 +∞"""
3
    return sqrt(sum((x - y) ** 2 for x, y in zip(a, b)))
4

5
def manhattan_distance(a, b):
6
    """曼哈顿距离"""
7
    return sum(abs(x - y) for x, y in zip(a, b))

1
class BatchEmbedding:
2
    def __init__(self, batch_size=32):
3
        self.batch_size = batch_size
4

5
    def encode_batch(self, texts):
6
        results = []
7
        for i in range(0, len(texts), self.batch_size):
8
            batch = texts[i:i+self.batch_size]
9
            embeddings = self.model.encode(batch)
10
            results.extend(embeddings)
11
        return results

1
# 增量添加到向量数据库
2
def add_documents(new_docs):
3
    embeddings = model.encode(new_docs)
4
    index.upsert(
5
        zip(ids, embeddings, metadatas)
6
    )
7

8
# 异步增量索引
9
async def add_async(new_docs):
10
    embeddings = await model.encode_async(new_docs)
11
    await index.upsert_async(embeddings)