Embedding 与相似度

演示文本向量化的核心概念：将文本映射为定长向量，并通过余弦相似度衡量语义接近程度。 为支持离线测试，提供基于哈希的确定性向量器 HashEmbedder；真实场景应使用 OpenAI Embedding API 或开源模型（如 bge-m3）。 对应教程：第1卷《AI 基础》第4章 Embedding。

文件路径: examples/volume-01-ai-fundamentals\embedding_demo.py

---

源代码

"""Embedding 与相似度示例。

演示文本向量化的核心概念：将文本映射为定长向量，并通过余弦相似度衡量语义接近程度。
为支持离线测试，提供基于哈希的确定性向量器 ``HashEmbedder``；真实场景应使用 OpenAI
Embedding API 或开源模型（如 bge-m3）。

对应教程：第1卷《AI 基础》第4章 Embedding。
"""

from __future__ import annotations

import hashlib
import math
import re
from dataclasses import dataclass


def tokenize(text: str) -> list[str]:
    """简易分词：英文按单词、中文按字。

    Args:
        text: 原始文本。

    Returns:
        token 列表（小写化）。
    """
    text = text.lower()
    # 提取英文单词
    words = re.findall(r"[a-z]+", text)
    # 提取中文字符
    chars = re.findall(r"[\u4e00-\u9fff]", text)
    return words + chars


@dataclass
class EmbeddingResult:
    """Embedding 结果。

    Attributes:
        text: 原始文本。
        vector: 向量。
        dim: 向量维度。
    """

    text: str
    vector: list[float]
    dim: int


class HashEmbedder:
    """基于哈希的确定性文本向量器。

    通过分词后对每个 token 哈希到固定维度并累加，构造稀疏向量。
    再做 L2 归一化，使其可直接用于余弦相似度（归一化后点积 = 余弦）。

    特点：
        - 确定性：相同文本始终得到相同向量。
        - 无需训练、无需外部模型。
        - 仅用于教学，语义能力远弱于真实 Embedding 模型。
    """

    def __init__(self, dim: int = 64, seed: int = 0) -> None:
        """初始化向量器。

        Args:
            dim: 输出向量维度。
            seed: 哈希种子，用于产生不同投影。

        Raises:
            ValueError: 维度非正时。
        """
        if dim <= 0:
            raise ValueError("维度必须为正整数")
        self.dim = dim
        self.seed = seed

    def _hash_token(self, token: str) -> tuple[int, float]:
        """将 token 哈希为一个维度索引与一个符号（+1/-1）。

        Args:
            token: 输入 token。

        Returns:
            ``(index, sign)`` 元组。
        """
        digest = hashlib.md5(f"{self.seed}:{token}".encode()).hexdigest()  # noqa: S324
        index = int(digest[:8], 16) % self.dim
        sign = 1.0 if int(digest[8:16], 16) % 2 == 0 else -1.0
        return index, sign

    def embed(self, text: str) -> EmbeddingResult:
        """将文本转换为向量。

        Args:
            text: 原始文本。

        Returns:
            ``EmbeddingResult``，向量已 L2 归一化。
        """
        vec = [0.0] * self.dim
        tokens = tokenize(text)
        for tok in tokens:
            index, sign = self._hash_token(tok)
            vec[index] += sign
        norm = math.sqrt(sum(v * v for v in vec))
        if norm > 0:
            vec = [v / norm for v in vec]
        return EmbeddingResult(text=text, vector=vec, dim=self.dim)

    def embed_batch(self, texts: list[str]) -> list[EmbeddingResult]:
        """批量向量化。

        Args:
            texts: 文本列表。

        Returns:
            Embedding 结果列表。
        """
        return [self.embed(t) for t in texts]


def cosine_similarity(vec_a: list[float], vec_b: list[float]) -> float:
    """计算两个向量的余弦相似度。

    Args:
        vec_a: 向量 A。
        vec_b: 向量 B。

    Returns:
        余弦相似度，范围 ``[-1, 1]``。

    Raises:
        ValueError: 维度不一致或向量为零向量时。
    """
    if len(vec_a) != len(vec_b):
        raise ValueError(f"维度不一致: {len(vec_a)} vs {len(vec_b)}")
    dot = sum(a * b for a, b in zip(vec_a, vec_b, strict=True))
    norm_a = math.sqrt(sum(a * a for a in vec_a))
    norm_b = math.sqrt(sum(b * b for b in vec_b))
    if norm_a == 0 or norm_b == 0:
        raise ValueError("零向量无法计算余弦相似度")
    return dot / (norm_a * norm_b)


def find_most_similar(
    query: str,
    candidates: list[str],
    embedder: HashEmbedder,
) -> list[tuple[str, float]]:
    """在候选文本中检索与 query 最相似的文本（按相似度降序）。

    Args:
        query: 查询文本。
        candidates: 候选文本列表。
        embedder: 向量器实例。

    Returns:
        ``(text, score)`` 列表，按相似度降序排列。
    """
    query_vec = embedder.embed(query).vector
    scored: list[tuple[str, float]] = []
    for cand in candidates:
        cand_vec = embedder.embed(cand).vector
        score = cosine_similarity(query_vec, cand_vec)
        scored.append((cand, score))
    scored.sort(key=lambda x: x[1], reverse=True)
    return scored


def main() -> None:
    """演示 Embedding 与相似度检索。"""
    embedder = HashEmbedder(dim=128)
    candidates = [
        "AI Agent 能够自主决策并执行任务",
        "今天天气不错，适合出门散步",
        "RAG 通过检索增强大模型生成",
        "机器学习是人工智能的一个分支",
    ]
    query = "智能体可以自己做决定"
    print(f"查询: {query}\n")
    for text, score in find_most_similar(query, candidates, embedder):
        print(f"  {score:.4f}  {text}")


if __name__ == "__main__":
    main()