LangChain 学习（二）

将PDF文档中的内容嵌入到Milvus数据库中。具体步骤如下：

1. 加载文档：使用PyPDFLoader从指定路径加载PDF文件。

2. 拆分文档：使用CharacterTextSplitter将文档拆分成较小的文本块。

3. 创建嵌入对象：使用BERT模型。

4. 文本块嵌入：对文本块进行嵌入。

5. 连接Milvus：通过提供连接信息连接到Milvus数据库。

6. 定义集合模式：定义包含ID、嵌入向量和文本字段的集合模式。

7. 创建集合：检查集合是否存在，如果不存在则创建。

8. 准备数据：准备要插入的ID、嵌入向量和文本内容。

9. 插入数据：将数据插入到Milvus中。

10. 刷新集合：确保数据写入成功。

import os

import torch
from langchain.text_splitter import CharacterTextSplitter
from langchain_community.document_loaders import PyPDFLoader
from pymilvus import connections, utility, FieldSchema, CollectionSchema, DataType, Collection
from transformers import AutoTokenizer, AutoModel, BertTokenizer

# Step 1: 加载文档目录中PDF文件
loader = PyPDFLoader("pdf/AI大模型.pdf")
documents = loader.load()

# 查看加载的文档
# print(f"加载 {len(documents)} 页文档。")
# total_text_length = sum(len(doc.page_content) for doc in documents)
# print(f"文本总长度: {total_text_length} 字符。")

# Step 2: 拆分文档
# 使用字符文本拆分器，将文档拆分成较小的块
text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=200) # 自定义块大小
spli_docx = text_splitter.split_documents(documents=documents)

# Step 3: 创建 BERT模型 嵌入对象
# model_name = "sentence-transformers/all-MiniLM-L6-v2"  # 可以选择其他模型
# tokenizer = AutoTokenizer.from_pretrained(model_name)
model_name = "bert-base-chinese"
tokenizer = BertTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name)


# 警告是因为 Hugging Face 的 huggingface_hub 在 Windows 系统上无法使用符号链接（symlinks）进行缓存，导致使用缓存时效率较低，并且占用更多磁盘空间。
# 禁用该警告
os.environ["HF_HUB_DISABLE_SYMLINKS_WARNING"] = "1"
# Step 4: 对拆分的文本块进行嵌入
def get_embeddings(text):
    try:
        inputs = tokenizer(text, padding=True, truncation=True, return_tensors="pt")
        with torch.no_grad():
            embeddings = model(**inputs).last_hidden_state[:, 0, :] # 获取[CLS]标记的嵌入
        return embeddings
    except Exception as e:
        print(f"错误处理文本信息: {text[:50]}...")  # 打印部分文本以便调试
        print(f"异常信息: {e}")
        return None  # 或者返回一个默认值

# 对拆分的文本块进行嵌入
doc_embeddings = []
texts = []
ids = []  # ID 列表
for id, doc in enumerate(spli_docx):
    embedding = get_embeddings(doc.page_content)
    if embedding is not None: # 确保只有有效的嵌入被添加
        # 将嵌入向量展平为列表并确保是浮点数
        doc_embeddings.append(embedding.flatten().tolist())  # 将嵌入展平并转为列表
        texts.append(doc.page_content)  # 保存对应的文本
        ids.append(id) # 使用自增的整数作为 ID

# 打印嵌入结果
# print(doc_embeddings)

# Step 5: 连接到 Milvus 数据库
connections.connect(host="127.99.199.88", port="19530", alias="default", user="root", password="Milvus2020")

# Step 6: 定义 Milvus 集合的 schema
fields = [
    FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=False),  # 主键字段
    FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=768),  # 假设嵌入向量维度为 768
    FieldSchema(name="text", dtype=DataType.VARCHAR, max_length=5000)
]
schema = CollectionSchema(fields=fields, description="文档嵌入")

# Step 7: 创建集合（如果不存在）
collection_name = "document_chinese_embeddings"
if not utility.has_collection(collection_name):
    collection = Collection(name=collection_name, schema=schema)
else:
    collection = Collection(name=collection_name)

# Step 8: 准备数据
data = [
    ids,  # 主键ID
    doc_embeddings,  # 确保每个嵌入都是浮点数
    texts # 文本内容
]

# 打印嵌入结果以进行检查
for i, embed in enumerate(doc_embeddings):
    print(f"嵌入向量 {i} 的长度: {len(embed)}, 类型: {type(embed[0])}, 内容: {embed[:5]}")  # 打印前五个值

# Step 9: 将数据插入到 Mivus
collection.insert(data)

# Step 10: 刷新集合以确保数据写入
collection.flush()

print("数据插入 Milvus 成功!")