文本嵌入

"Go 语言的并发模型"  →  [0.12, -0.34, 0.56, ..., 0.78]  (1536维)
"Golang 的 goroutine" →  [0.11, -0.32, 0.58, ..., 0.75]  (1536维)
"今天天气不错"        →  [0.89, 0.23, -0.45, ..., -0.12]  (1536维)

前两个向量距离很近（语义相关），第三个向量距离较远（语义无关）

from openai import OpenAI
import numpy as np

client = OpenAI(api_key="your-api-key")

def get_embedding(text: str, model: str = "text-embedding-3-small") -> list[float]:
    """获取文本的嵌入向量"""
    response = client.embeddings.create(
        input=text,
        model=model
    )
    return response.data[0].embedding

def get_embeddings_batch(texts: list[str], model: str = "text-embedding-3-small") -> list[list[float]]:
    """批量获取嵌入向量（效率更高）"""
    response = client.embeddings.create(
        input=texts,
        model=model
    )
    return [item.embedding for item in response.data]

# 单条文本嵌入
embedding = get_embedding("RAG 是检索增强生成技术")
print(f"向量维度: {len(embedding)}")  # 输出: 1536

# 批量嵌入
texts = [
    "Docker 容器化部署",
    "Kubernetes 集群管理",
    "今天吃什么",
]
embeddings = get_embeddings_batch(texts)

# 计算相似度
def cosine_similarity(a, b):
    a, b = np.array(a), np.array(b)
    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))

sim_01 = cosine_similarity(embeddings[0], embeddings[1])
sim_02 = cosine_similarity(embeddings[0], embeddings[2])
print(f"Docker vs Kubernetes 相似度: {sim_01:.4f}")  # 较高
print(f"Docker vs 今天吃什么 相似度: {sim_02:.4f}")  # 较低

# text-embedding-3 系列支持通过 dimensions 参数降维
response = client.embeddings.create(
    input="RAG 检索增强生成",
    model="text-embedding-3-large",
    dimensions=512  # 将 3072 维降到 512 维
)
embedding = response.data[0].embedding
print(f"降维后向量维度: {len(embedding)}")  # 输出: 512

from sentence_transformers import SentenceTransformer
import numpy as np

# 加载 BGE 中文模型（首次运行会自动下载）
model = SentenceTransformer("BAAI/bge-base-zh-v1.5")

# 单条文本嵌入
text = "什么是向量数据库？"
embedding = model.encode(text, normalize_embeddings=True)
print(f"向量维度: {embedding.shape}")  # 输出: (768,)

# 批量嵌入
texts = [
    "向量数据库用于存储和检索高维向量",
    "Redis 是一个内存键值数据库",
    "Milvus 是开源的向量检索引擎",
]
embeddings = model.encode(texts, normalize_embeddings=True)

# 计算相似度矩阵
similarity_matrix = np.inner(embeddings, embeddings)
print("相似度矩阵:")
print(np.round(similarity_matrix, 3))

from FlagEmbedding import BGEM3FlagModel

# 加载 BGE-M3 模型
model = BGEM3FlagModel("BAAI/bge-m3", use_fp16=True)

sentences = [
    "什么是检索增强生成？",
    "RAG 的基本原理是什么？",
    "如何做红烧肉？",
]

# 同时获取 dense 和 sparse 向量
embeddings = model.encode(
    sentences,
    return_dense=True,
    return_sparse=True,
)

# Dense 向量用于语义检索
dense_vectors = embeddings["dense_vecs"]  # shape: (3, 1024)

# Sparse 向量用于关键词匹配（类似 BM25）
sparse_vectors = embeddings["lexical_weights"]

def fixed_size_chunk(text: str, chunk_size: int = 500, overlap: int = 100) -> list[str]:
    """固定大小分块"""
    chunks = []
    start = 0
    while start < len(text):
        end = start + chunk_size
        chunks.append(text[start:end])
        start = end - overlap  # 向前重叠
    return chunks

text = "这是一段很长的文档内容..." * 100
chunks = fixed_size_chunk(text, chunk_size=500, overlap=100)
print(f"分块数量: {len(chunks)}")

import re

def sentence_chunk(text: str, max_chunk_size: int = 500) -> list[str]:
    """基于句子的分块"""
    # 按中英文句号、问号、感叹号分句
    sentences = re.split(r'(?<=[。！？.!?])', text)
    sentences = [s.strip() for s in sentences if s.strip()]

    chunks = []
    current_chunk = ""
    for sentence in sentences:
        if len(current_chunk) + len(sentence) <= max_chunk_size:
            current_chunk += sentence
        else:
            if current_chunk:
                chunks.append(current_chunk)
            current_chunk = sentence
    if current_chunk:
        chunks.append(current_chunk)
    return chunks

from sentence_transformers import SentenceTransformer
import numpy as np

def semantic_chunk(text: str, model, threshold: float = 0.5) -> list[str]:
    """基于语义相似度的分块"""
    # 先按句子分割
    sentences = re.split(r'(?<=[。！？.!?])', text)
    sentences = [s.strip() for s in sentences if s.strip()]

    if len(sentences) <= 1:
        return sentences

    # 计算每个句子的嵌入
    embeddings = model.encode(sentences, normalize_embeddings=True)

    # 计算相邻句子的相似度
    chunks = []
    current_chunk = [sentences[0]]
    for i in range(1, len(sentences)):
        similarity = np.dot(embeddings[i-1], embeddings[i])
        if similarity < threshold:
            # 语义变化较大，开始新的分块
            chunks.append("".join(current_chunk))
            current_chunk = [sentences[i]]
        else:
            current_chunk.append(sentences[i])

    if current_chunk:
        chunks.append("".join(current_chunk))
    return chunks