将PDF文档中的内容嵌入到Milvus数据库中。具体步骤如下：
-
加载文档：使用
PyPDFLoader从指定路径加载PDF文件。 -
拆分文档：使用
CharacterTextSplitter将文档拆分成较小的文本块。 -
创建嵌入对象：使用BERT模型。
-
文本块嵌入：对文本块进行嵌入。
-
连接Milvus：通过提供连接信息连接到Milvus数据库。
-
定义集合模式：定义包含ID、嵌入向量和文本字段的集合模式。
-
创建集合：检查集合是否存在，如果不存在则创建。
-
准备数据：准备要插入的ID、嵌入向量和文本内容。
-
插入数据：将数据插入到Milvus中。
-
刷新集合：确保数据写入成功。
import osimport 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 embeddingsexcept 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), # 假设嵌入向量维度为 768FieldSchema(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, # 主键IDdoc_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 成功!")
