把文本放进向量数据库分几步

向量数据库作为一种新兴的数据存储和检索技术，近年来受到了广泛关注。它能够高效地存储和查询高维向量数据，这使得在文本相似性搜索、推荐系统、图像检索等领域有着广泛的应用前景。

在使用向量数据库进行语义检索之前，文本是怎么一步步变成存储在向量数据库中代表语义信息的向量呢？主要有以下的几个步骤

Chunking (分块)：将长文本分割成更小的、更易于处理的片段。

Tokenization (分词)：将文本片段分解成一个个独立的词或子词（token）。

Embedding (向量化)：将每个 token 转换为一个高维的向量表示，捕捉其语义信息。

Embedding Pooling (向量池化)：将一个文本片段中所有 token 的 embedding 聚合成一个代表该片段的向量。

存储 Embedding 和 Meta 信息：将生成的向量以及相关的元数据存储到向量数据库中。

Chunk (分块)

当处理较长的文本时，直接进行 embedding 可能会遇到以下问题：

超出模型处理长度限制：许多 embedding 模型对输入文本的长度有限制。

语义关联性降低：过长的文本可能包含多个主题，导致生成的 embedding 难以准确表示特定主题。

检索效率降低：对整个长文本进行相似性搜索的效率较低。

因此，我们需要将长文本分割成更小的块（chunk）。一个简单的 chunking 方法是按照固定的字符长度进行分割。

def simple_chunk(text, chunk_size=28):
  """
  简单分块方法：按固定字符长度分割文本。
  """
  chunks = [text[i:i + chunk_size] for i in range(0, len(text), chunk_size)]
  return chunks

text = "This is an example sentence Each sentence is converted"
chunks = simple_chunk(text)
print(chunks)

但是这样会导致切分出来的chunk语义不完整等问题，可以按照下面的方案继续优化：

基于句子的分割: 使用句号、问号、感叹号等标点符号作为分割符，可以更好地保持语义完整性。可以使用 nltk 或 spaCy 等 NLP 工具库进行句子分割。

滑动窗口式分割: 为了避免信息丢失，可以采用滑动窗口的方式，让相邻的 chunk 之间存在一定的重叠。这有助于在检索时保留上下文信息。

语义分块: 更高级的方法是利用语义信息进行分块，例如识别文本中的主题或关键信息，并以此为依据进行分割。这通常需要更复杂的 NLP 技术。

Tokenizer (分词)

Tokenizer 的作用是将文本字符串拆分成一个个单独的 token。

from transformers import AutoTokenizer

sentences = ['This is an example sentence', 'Each sentence is converted']
tokenizer = AutoTokenizer.from_pretrained('sentence-transformers/all-MiniLM-L6-v2')

encoded_input = tokenizer(sentences, padding=True, truncation=True, return_tensors='pt')
print(encoded_input)

输出：

{
	'input_ids': tensor([
					[ 101, 2023, 2003, 2019, 2742, 6251,  102],
	        [ 101, 2169, 6251, 2003, 4991,  102,    0]  // 最后一个0是padding，切分后的长度不够补上的，用以保证向量长度一致
        ]),
	'token_type_ids': tensor([
					[0, 0, 0, 0, 0, 0, 0],
	        [0, 0, 0, 0, 0, 0, 0]
		    ]),
	'attention_mask': tensor([
					[1, 1, 1, 1, 1, 1, 1],
	        [1, 1, 1, 1, 1, 1, 0]   // 最后一个0标记上面对应位置的id其实是补上的padding
        ])
}

Embedding (向量化)

Embedding 步骤将每个 token 转换为一个固定长度的向量，这个向量能够捕捉 token 的语义信息。

from transformers import AutoModel
import torch

model = AutoModel.from_pretrained('sentence-transformers/all-MiniLM-L6-v2')

with torch.no_grad():
    model_output = model(**encoded_input) # tokenizer输出的encoded_input

print(model_output.last_hidden_state.shape)

输出：

torch.Size([2, 7, 384])

这个last_hidden_state就是向量的embedding表示，shape的2，7，384分别代表：

2代表：有两个句子

7代表每个句子有7个token

384代表每个token由一个384维的向量来表示

Embedding Pooling (向量池化)

对于一个文本 chunk，我们通常需要将其所有 token 的 embedding 聚合成一个单独的向量，用于后续的存储和检索。常见的 pooling 方法包括：

Mean Pooling (平均池化): 将所有 token 的 embedding 向量按元素求平均值。

Max Pooling (最大池化): 取每个维度上的最大值。

CLS Token Embedding: 某些 Transformer 模型（如 BERT）在输入序列的开头添加一个特殊的 [CLS] token，其最终的 embedding 可以作为整个序列的表示。

这里我们用 Mean Pooling 作为示例：

import torch.nn.functional as F

def mean_pooling(model_output, attention_mask):
    token_embeddings = model_output[0] #First element of model_output contains all token embeddings
    input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
    return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)

# 池化
sentence_embeddings = mean_pooling(model_output, encoded_input['attention_mask'])

# 归一化
sentence_embeddings = F.normalize(sentence_embeddings, p=2, dim=1)

print(sentence_embeddings.shape)

输出：

torch.Size([2, 384])

这里输出的shape，中

2代表两个句子

384代表每个句子最终的向量表示

优化方案：

根据任务选择 Pooling 方法: 不同的 pooling 方法在不同的任务上可能表现出差异。例如，Mean Pooling 倾向于保留所有 token 的信息，而 Max Pooling 则更关注最重要的特征。

Weighted Pooling: 可以根据 token 的重要性（例如 TF-IDF 值）对其 embedding 进行加权平均。

Attention Pooling: 使用注意力机制来学习不同 token 的重要性，并根据注意力权重进行 pooling。

存储 Embedding 和 Meta 信息

最后一步是将生成的 embedding 向量以及相关的元数据存储到向量数据库中。元数据可以包含原始文本内容、文档 ID、创建时间、来源等信息。

Meta 信息的价值：

过滤和筛选: 在检索时，可以根据元数据对结果进行过滤，例如只搜索特定来源的文档。

上下文信息: 元数据可以提供关于 embedding 的上下文信息，帮助理解检索结果。

溯源和管理: 元数据有助于跟踪和管理存储的文本数据。

# 假设我们已经有了一个向量数据库 client (例如：Weaviate, Pinecone, Chroma)
# 这里仅为示例，不涉及实际数据库操作

vector = pooled_embedding.tolist()
metadata = {
    "text": text,
    "source": "blog post",
    "timestamp": "2023-10-27"
}

# 伪代码：将向量和元数据存储到向量数据库
# vector_db_client.add(vector=vector, metadata=metadata)

print("Embedding 和 Meta 信息已准备好存储。")
print("Vector (部分):", vector[:5])
print("Metadata:", metadata)

在实际应用中，你需要选择一个合适的向量数据库，并使用其提供的 API 将 embedding 向量和元数据添加到数据库中。常见的向量数据库包括：

Pinecone: 云端托管的向量数据库，性能强大。

Weaviate: 开源的向量数据库，支持 GraphQL。

Milvus: 开源的向量数据库，功能丰富。

Chroma: 轻量级的嵌入式向量数据库，易于上手。

完整代码

下面是将以上所有步骤整合在一起的完整代码示例：

from transformers import AutoTokenizer, AutoModel
import torch
import torch.nn.functional as F

def simple_chunk(text, chunk_size=28):
  """
  简单分块方法：按固定字符长度分割文本。
  """
  chunks = [text[i:i + chunk_size] for i in range(0, len(text), chunk_size)]
  return chunks

text = "This is an example sentence Each sentence is converted"
chunks = simple_chunk(text)
print(chunks)


#Mean Pooling - Take attention mask into account for correct averaging
def mean_pooling(model_output, attention_mask):
    token_embeddings = model_output[0] #First element of model_output contains all token embeddings
    input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
    return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)


# Sentences we want sentence embeddings for
sentences = ['This is an example sentence', 'Each sentence is converted']

# Load model from HuggingFace Hub
tokenizer = AutoTokenizer.from_pretrained('sentence-transformers/all-MiniLM-L6-v2')
model = AutoModel.from_pretrained('sentence-transformers/all-MiniLM-L6-v2')

# Tokenize sentences
encoded_input = tokenizer(sentences, padding=True, truncation=True, return_tensors='pt')

# Compute token embeddings
with torch.no_grad():
    model_output = model(**encoded_input)

# Perform pooling
sentence_embeddings = mean_pooling(model_output, encoded_input['attention_mask'])

# Normalize embeddings
sentence_embeddings = F.normalize(sentence_embeddings, p=2, dim=1)
print(sentence_embeddings.shape)