【大模型】文本向量化与FAISS向量索引构建
·
文本向量化与FAISS向量索引构建
初始化 RAG 系统
加载嵌入模型并初始化 FAISS 索引。模型默认使用 sentence-transformers/all-MiniLM-L6-v2，适用于英文文本。对于中文任务，建议替换为 moka-ai/m3e-base 或 BAAI/bge-large-zh-v1.5。模型会优先尝试从本地缓存加载，若不存在则自动下载
def __init__(self, model_name: str = 'sentence-transformers/all-MiniLM-L6-v2'):
local_model_path = os.path.join(os.getcwd(), "models", model_name.split("/")[-1])
if not os.path.exists(local_model_path):
snapshot_download(repo_id=model_name, local_dir=local_model_path)
self.encoder = SentenceTransformer(local_model_path)
self.dimension = self.encoder.get_sentence_embedding_dimension()
self.index = faiss.IndexFlatL2(self.dimension)
self.documents = []
添加文档到知识库
将文本列表转换为向量并存入 FAISS 索引，同时保留原始文本映射。支持批量添加文档，自动处理文本编码和向量化。
def add_documents(self, texts: List[str]):
self.documents.extend(texts)
embeddings = self.encoder.encode(texts)
embeddings = np.array(embeddings).astype('float32')
self.index.add(embeddings)
执行相似性检索
输入查询文本，返回最相关的 k 个文档。使用 L2 距离进行相似度计算，结果按相关性排序。
def search(self, query: str, k: int = 3) -> List[str]:
query_embedding = np.array(self.encoder.encode([query])).astype('float32')
_, indices = self.index.search(query_embedding, k)
return [self.documents[i] for i in indices[0] if i != -1]
生成增强回答
结合检索结果和大语言模型生成最终回答。需预先配置 dashscope API 密钥环境变量 DASHSCOPE_API_KEY。
def generate_answer(self, query: str, context: List[str]) -> str:
prompt = f"根据以下信息回答问题：\n上下文：{' '.join(context)}\n问题：{query}"
response = dashscope.Generation.call(
model='qwen-turbo',
prompt=prompt
)
return response.output.text if response.output else "生成失败"
中文优化建议
- 替换嵌入模型为中文专用模型如 BAAI/bge-large-zh-v1.5
- 添加中文预处理步骤，如使用 jieba 分词
- 调整 FAISS 索引参数，对于大规模数据建议使用 IndexIVFFlat
rag = SimpleRAG(model_name='BAAI/bge-large-zh-v1.5')
rag.add_documents(["文档1内容", "文档2内容..."])
results = rag.search("用户查询")
answer = rag.generate_answer("用户查询", results)
完整的代码与说明如下：
import os # 导入操作系统接口模块，用于处理文件路径和环境变量
# 设置 Hugging Face 镜像站，解决国内访问问题
# 必须在导入 transformers 或 sentence_transformers 之前设置
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
import re # 导入正则表达式模块，用于文本匹配和处理
import faiss # 导入 FAISS 库，用于高效的向量相似度搜索和聚类
import numpy as np # 导入 NumPy 库，用于进行科学计算和数组操作
from collections import defaultdict # 从 collections 模块导入 defaultdict，用于创建带有默认值的字典
from typing import List, Tuple # 从 typing 模块导入 List 和 Tuple，用于类型注解，提高代码可读性
from sentence_transformers import SentenceTransformer # 导入 SentenceTransformer，用于加载预训练模型并将文本转换为向量
from huggingface_hub import snapshot_download # 导入 snapshot_download，用于手动下载模型
from dotenv import load_dotenv # 导入 load_dotenv，用于从 .env 文件加载环境变量
import jieba # 导入 jieba 库，用于中文分词（虽然本示例主要用 SentenceTransformer，但 jieba 常用于传统 NLP 任务）
import dashscope # 导入 dashscope 库，用于调用阿里云通义千问等大模型 API
# 加载 .env 文件中的环境变量
# 这通常用于安全地管理 API 密钥，而不是直接硬编码在代码中
load_dotenv()
# 设置 dashscope API 密钥
# 从环境变量中获取 'DASHSCOPE_API_KEY'，如果没有设置，需要手动配置或在 .env 文件中添加
api_key = os.getenv("DASHSCOPE_API_KEY")
if api_key:
dashscope.api_key = api_key
else:
print("警告: 未找到 DASHSCOPE_API_KEY 环境变量，请确保已配置。")
class SimpleRAG:
"""
一个简单的 RAG (检索增强生成) 系统类
演示了如何结合向量检索和 LLM 生成
"""
def __init__(self, model_name: str = 'sentence-transformers/all-MiniLM-L6-v2'):
# def __init__(self, model_name: str = 'sentence-transformers/parapnrase-multilingual-MiniLM-L12-v2'):
"""
初始化 RAG 系统
Args:
model_name (str): 嵌入模型的名称。默认使用 'all-MiniLM-L6-v2'，这是一个轻量级且效果不错的模型。
对于中文，推荐使用 'moka-ai/m3e-base' 或 'BAAI/bge-large-zh-v1.5'。
"""
print(f"正在加载嵌入模型: {model_name}...")
# 为了避免 Windows 上的符号链接权限问题，我们将模型下载到本地目录
# 并显式指定 local_dir，这样 huggingface_hub 就不会尝试创建符号链接
if "/" in model_name:
model_dir_name = model_name.split("/")[-1]
else:
model_dir_name = model_name
local_model_path = os.path.join(os.getcwd(), "models", model_dir_name)
try:
if not os.path.exists(local_model_path):
print(f"模型未找到，正在下载到: {local_model_path}")
# snapshot_download 会下载整个仓库到本地目录
snapshot_download(repo_id=model_name, local_dir=local_model_path)
else:
print(f"加载本地模型: {local_model_path}")
# 初始化 SentenceTransformer 模型，用于将文本转化为向量
self.encoder = SentenceTransformer(local_model_path)
except Exception as e:
print(f"模型加载失败: {e}")
print("尝试直接加载模型（可能会遇到符号链接问题）...")
self.encoder = SentenceTransformer(model_name)
# 获取模型输出的向量维度，FAISS 索引需要知道这个维度
self.dimension = self.encoder.get_sentence_embedding_dimension()
# 初始化 FAISS 索引
# IndexFlatL2 是一种基于欧几里得距离（L2）的精确搜索索引
# 适用于数据量不是特别大的情况
self.index = faiss.IndexFlatL2(self.dimension)
# 用于存储原始文档文本的列表
# FAISS 只存储向量，我们需要自己维护向量索引到原始文本的映射
self.documents = []
def add_documents(self, texts: List[str]):
"""
向知识库中添加文档
Args:
texts (List[str]): 文档文本列表
"""
if not texts:
return
print(f"正在处理 {len(texts)} 个文档...")
# 将新文档添加到内部列表中
self.documents.extend(texts)
# 使用编码器将文本转换为向量 (embeddings)
embeddings = self.encoder.encode(texts)
# 将向量转换为 float32 类型的 numpy 数组，这是 FAISS 要求的格式
embeddings = np.array(embeddings).astype('float32')
# 将向量添加到 FAISS 索引中
self.index.add(embeddings)
print(f"成功添加文档，当前索引包含 {self.index.ntotal} 个向量。")
def search(self, query: str, k: int = 3) -> List[str]:
"""
检索与查询最相关的文档
Args:
query (str): 查询文本
k (int): 返回的最相关文档数量，默认为 3
Returns:
List[str]: 检索到的相关文档内容列表
"""
# 将查询文本转换为向量
query_embedding = self.encoder.encode([query])
# 转换为 FAISS 需要的格式
query_embedding = np.array(query_embedding).astype('float32')
# 在索引中搜索最近邻
# D 是距离数组 (distances)，I 是索引数组 (indices)
# k 是我们要查找的最近邻数量
D, I = self.index.search(query_embedding, k)
# 根据索引从 self.documents 中获取对应的文本
# I[0] 是因为我们只查询了一个向量
# if i != -1 是为了过滤掉无效索引（虽然在 IndexFlatL2 中通常不会出现）
results = [self.documents[i] for i in I[0] if i != -1 and i < len(self.documents)]
return results
def generate_answer(self, query: str, context: List[str]) -> str:
"""
根据检索到的上下文生成回答
Args:
query (str): 用户的问题
context (List[str]): 检索到的相关文档内容
Returns:
str: LLM 生成的回答
"""
# 将检索到的文档拼接成一个字符串，作为上下文
context_str = "\n".join(context)
# 构建提示词 (Prompt)
# 这是一个简单的 RAG 提示词模板
prompt = f"""请根据以下参考资料回答问题。如果参考资料中没有答案，请根据你的知识回答，但要优先参考资料。
参考资料：
{context_str}
问题：
{query}
"""
try:
# 调用 dashscope (通义千问) API 生成回答
# 这里使用 qwen-turbo 模型，速度快且成本低
response = dashscope.Generation.call(
model=dashscope.Generation.Models.qwen_turbo,
prompt=prompt
)
# 检查 API 调用是否成功
if response.status_code == 200:
return response.output.text
else:
return f"API 调用失败: {response.code} - {response.message}"
except Exception as e:
return f"生成回答时发生错误: {str(e)}"
# 主程序入口
if __name__ == "__main__":
# 实例化 RAG 系统
# 注意：首次运行会下载模型，可能需要一些时间
rag = SimpleRAG()
# 准备一些示例文档
# 在实际应用中，这些可能来自 PDF、TXT 文件或数据库
sample_docs = [
"RAG (Retrieval-Augmented Generation) 是一种结合了检索和生成的 AI 技术架构。",
"通过检索外部知识库，RAG 可以显著减少大语言模型的幻觉问题，并提供最新的信息。",
"FAISS 是 Facebook AI Research 开发的高效相似性搜索库，能够快速在大规模向量集中进行搜索。",
"SentenceTransformer 是一个基于 PyTorch 和 Transformers 的 Python 库，用于生成句子、文本和图像的嵌入向量。",
"通义千问 (Qwen) 是阿里云推出的大语言模型，具备强大的中文理解和生成能力。",
"Python 是一种广泛使用的高级编程语言，以其简洁的语法和丰富的库生态系统而闻名。"
]
# 将文档添加到 RAG 系统中
rag.add_documents(sample_docs)
# 定义一个测试问题
user_query = "RAG 技术有什么主要优势？"
print(f"\n用户问题: {user_query}")
# 1. 检索阶段：找到相关的文档
retrieved_docs = rag.search(user_query)
print("\n检索到的相关上下文:")
for i, doc in enumerate(retrieved_docs, 1):
print(f"{i}. {doc}")
# 2. 生成阶段：调用 LLM 生成回答
# 只有在设置了 API Key 的情况下才执行
if dashscope.api_key:
print("\n正在生成回答...")
answer = rag.generate_answer(user_query, retrieved_docs)
print(f"\nAI 回答: {answer}")
else:
# 如果没有 API Key，程序不会崩溃，只是跳过生成环节
print("\n--------------------------------------------------")
print("【提示】未检测到 DashScope API Key，跳过 LLM 生成步骤。")
print("当前运行仅演示了 RAG 的检索 (Retrieval) 阶段。")
print("检索功能完全可用，已成功找到相关文档 (如上所示)。")
print("若需体验生成 (Generation) 效果，请在 .env 文件中配置 Key。")
print("--------------------------------------------------")
运行结果示例
更多推荐
2
0- 0
已为社区贡献1条内容
所有评论(0)