需求文档是测试工程师日常工作的核心工具,如何快速检索需求文档中的关键信息(文本、表格、图片等),并将其转化为可供 AI 查询的知识库,是提升工作效率的重要手段。本文将通过对 需求文档(docx 格式) 的处理,详细讲解如何构建一个 安全的本地知识库,并通过代码实现具体操作,确保每一步都可落地。
在构建本地知识库时,安全性是首要考虑的因素,尤其是对于需求文档这样的敏感数据。以下是本地知识库的安全性与连接方案:
以下所有代码均基于本地运行,适配 Windows、Mac 和 Linux 环境。
针对需求文档(docx 格式),我们需要将文档中的 文本、表格、图片 等内容提取并向量化。以下是完整流程图:
+-------------------------------------------+
| 需求文档加载 |
| 使用 python-docx 提取 docx 文件内容 |
| 支持文本、表格、图片 |
+-------------------------------------------+
↓
+-------------------------------------------+
| 数据清洗 |
| 移除空行、无效字符,提取关键信息 |
+-------------------------------------------+
↓
+-------------------------------------------+
| 文本分割 |
| 将文本分割为小段(500字左右) |
| 表格转化为结构化文本 |
+-------------------------------------------+
↓
+-------------------------------------------+
| 文本向量化 |
| 使用中文预训练模型生成向量 |
+-------------------------------------------+
↓
+-------------------------------------------+
| 图片向量化 |
| 使用 OpenAI CLIP 模型生成图片向量 |
+-------------------------------------------+
↓
+-------------------------------------------+
| 存储到向量数据库 |
| 使用 Chroma 将向量和原始数据存储 |
+-------------------------------------------+
↓
+-------------------------------------------+
| 本地知识库检索 |
| 根据用户输入,检索最相关的文档片段 |
+-------------------------------------------+
安装以下 Python 库:
pip install python-docx chromadb sentence-transformers PIL
使用 python-docx
提取 docx 文件中的内容(包括文本与表格)。对于图片,则使用 PIL
进行处理。
以下是提取需求文档的代码:
from docx import Document
from PIL import Image
from io import BytesIO
import os
# 提取文本和表格
def extract_text_and_tables(docx_path):
document = Document(docx_path)
texts = []
tables = []
# 提取段落文本
for paragraph in document.paragraphs:
if paragraph.text.strip(): # 去掉空行
texts.append(paragraph.text.strip())
# 提取表格内容
for table in document.tables:
table_data = []
for row in table.rows:
row_data = [cell.text.strip() for cell in row.cells]
table_data.append(row_data)
tables.append(table_data)
return texts, tables
# 提取图片
def extract_images(docx_path, output_folder):
document = Document(docx_path)
os.makedirs(output_folder, exist_ok=True)
image_paths = []
for rel in document.part.rels.values():
if "image" in rel.target_ref:
image_data = rel.target_part.blob
image = Image.open(BytesIO(image_data))
image_path = os.path.join(output_folder, f"image_{len(image_paths)+1}.png")
image.save(image_path)
image_paths.append(image_path)
return image_paths
# 示例:提取需求文档内容
docx_path = "requirement_document.docx"
output_folder = "extracted_images"
texts, tables = extract_text_and_tables(docx_path)
image_paths = extract_images(docx_path, output_folder)
print(f"提取文本段落:{len(texts)} 段")
print(f"提取表格:{len(tables)} 个")
print(f"提取图片:{len(image_paths)} 张")
对提取的文本和表格进行清洗,移除无效字符和空白内容:
import re
# 清洗文本
def clean_text(text):
text = re.sub(r"\s+", " ", text) # 移除多余空白
text = re.sub(r"[^\u4e00-\u9fa5a-zA-Z0-9,。!?;:]", "", text) # 移除特殊字符
return text
cleaned_texts = [clean_text(text) for text in texts]
# 转换表格为结构化文本
def table_to_text(table):
return "\n".join([" | ".join(row) for row in table])
cleaned_tables = [table_to_text(table) for table in tables]
print("清洗完成!")
将文本和表格分割为适合向量化的小段(每段 500 字左右):
from langchain.text_splitter import CharacterTextSplitter
# 分割文本
text_splitter = CharacterTextSplitter(chunk_size=500, chunk_overlap=50)
documents = text_splitter.split_text("\n".join(cleaned_texts + cleaned_tables))
print(f"分割为 {len(documents)} 个片段")
使用 sentence-transformers
的中文模型生成文本向量:
from sentence_transformers import SentenceTransformer
# 加载中文预训练模型
model = SentenceTransformer("damo/nlp_corom_sentence-embedding_chinese-base")
# 向量化文本片段
text_embeddings = model.encode(documents)
print(f"生成了 {len(text_embeddings)} 个文本向量")
使用 OpenAI 的 CLIP 模型生成图片向量:
pip install open_clip_torch
import open_clip
import torch
from PIL import Image
# 加载 CLIP 模型
model, preprocess = open_clip.create_model_and_transforms('ViT-B-32', pretrained='openai')
tokenizer = open_clip.get_tokenizer('ViT-B-32')
# 向量化图片
image_embeddings = []
for image_path in image_paths:
image = preprocess(Image.open(image_path)).unsqueeze(0) # 预处理图片
with torch.no_grad():
image_embedding = model.encode_image(image)
image_embeddings.append(image_embedding.numpy())
print(f"生成了 {len(image_embeddings)} 个图片向量")
将文本与图片的向量及原始数据存储到 Chroma 中:
import chromadb
# 初始化数据库
client = chromadb.PersistentClient(path="./chroma_knowledge_base")
# 创建集合
collection = client.get_or_create_collection(name="requirement_docs")
# 存储文本向量
collection.add(
documents=documents, # 原始文本
embeddings=text_embeddings, # 文本向量
ids=[f"text_{i}" for i in range(len(documents))]
)
# 存储图片向量
collection.add(
documents=image_paths, # 图片路径
embeddings=image_embeddings, # 图片向量
ids=[f"image_{i}" for i in range(len(image_paths))]
)
print("成功存储文本与图片向量")
根据用户输入的问题,检索相关文本或图片:
# 用户输入问题
query = "登录界面的设计要求是什么?"
query_embedding = model.encode([query])
# 检索最相关的结果
results = collection.query(
query_embeddings=query_embedding,
n_results=5
)
# 打印检索结果
for i, doc in enumerate(results["documents"][0]):
print(f"相关结果 {i+1}:{doc}")
本文介绍了如何针对需求文档(docx 格式)构建本地知识库的完整流程,包括 文本、表格、图片的提取与向量化,并结合 Chroma 数据库实现了安全的本地存储与检索。以下是本地知识库的关键特点:
通过这种方式,测试工程师可以快速构建起属于自己的安全本地知识库,大幅提升工作效率!