321 lines
11 KiB
Python
321 lines
11 KiB
Python
#!/usr/bin/env python
|
|
# -*- coding: utf-8 -*-
|
|
"""
|
|
File: ProfessionalNounVector.py
|
|
Date: 2025-05-15
|
|
Author: oyyz
|
|
Description: 专业名词向量化和检索的核心逻辑
|
|
"""
|
|
|
|
import os
|
|
import json
|
|
import shutil
|
|
from typing import List, Dict, Any, Tuple, Optional
|
|
from langchain.embeddings.base import Embeddings
|
|
from langchain_community.vectorstores import FAISS
|
|
from rag2_0.tool.ModelTool import SiliconFlowEmbeddings
|
|
import logging
|
|
|
|
def get_embedding_model(api_key: str = None) -> Embeddings:
|
|
"""
|
|
获取嵌入模型
|
|
|
|
Args:
|
|
api_key: API密钥，如果为None则从环境变量获取
|
|
|
|
Returns:
|
|
嵌入模型实例
|
|
"""
|
|
if not api_key:
|
|
api_key = os.getenv("SILICONFLOW_API_KEY", "sk-ftnofbucchwnscojohyxwmfzgaykdxihafnlphohsinftkbr")
|
|
return SiliconFlowEmbeddings(api_key=api_key)
|
|
|
|
|
|
class ProfessionalNounVectorizer:
|
|
"""专业名词向量化和保存类"""
|
|
|
|
def __init__(self,
|
|
embedding_model: Optional[Embeddings] = None,
|
|
api_key: str = None,
|
|
output_dir: str = None):
|
|
"""
|
|
初始化向量化器
|
|
|
|
Args:
|
|
embedding_model: 嵌入模型，如果为None则使用默认模型
|
|
api_key: SiliconFlow API密钥，仅在embedding_model为None时使用
|
|
|
|
output_dir: 索引输出目录，默认为None使用默认路径
|
|
"""
|
|
# 设置嵌入模型
|
|
if embedding_model:
|
|
self.embedding_model = embedding_model
|
|
else:
|
|
self.embedding_model = get_embedding_model(api_key)
|
|
|
|
|
|
# 设置输出目录
|
|
self.output_dir = output_dir
|
|
if self.output_dir is None:
|
|
current_dir = os.path.dirname(os.path.abspath(__file__))
|
|
self.output_dir = os.path.join(current_dir, "..", "..", "data", "nouns")
|
|
|
|
# 设置索引路径
|
|
self.index_path = os.path.join(self.output_dir, "professional_nouns_index")
|
|
|
|
def _loadfile(self, file_paths: List[str]) -> List[Dict[str, Any]]:
|
|
"""
|
|
加载多个专业术语JSON文件并合并
|
|
|
|
Args:
|
|
file_paths: JSON文件路径列表
|
|
|
|
Returns:
|
|
合并后的术语列表
|
|
"""
|
|
merged_terms = []
|
|
|
|
try:
|
|
for file_path in file_paths:
|
|
if not os.path.exists(file_path):
|
|
logging.warning(f"文件不存在: {file_path}")
|
|
continue
|
|
|
|
with open(file_path, "r", encoding="utf-8") as f:
|
|
terms_data = json.load(f)
|
|
|
|
if isinstance(terms_data, list):
|
|
merged_terms.extend(terms_data)
|
|
logging.info(f"从 {file_path} 加载了 {len(terms_data)} 条专业名词")
|
|
else:
|
|
logging.warning(f"文件格式错误: {file_path}，应为JSON数组")
|
|
|
|
logging.info(f"总共加载了 {len(merged_terms)} 条专业名词")
|
|
return merged_terms
|
|
except Exception as e:
|
|
logging.error(f"加载多个文件失败: {e}")
|
|
return []
|
|
|
|
def vectorize_files_and_save(self, file_paths: List[str]) -> bool:
|
|
"""
|
|
处理多个文件：加载多个术语文件、创建索引并保存
|
|
|
|
Args:
|
|
file_paths: JSON文件路径列表
|
|
|
|
Returns:
|
|
处理成功返回True，否则返回False
|
|
"""
|
|
try:
|
|
# 加载多个文件的术语
|
|
terms = self._loadfile(file_paths)
|
|
|
|
if not terms:
|
|
logging.warning("未找到术语数据，退出处理")
|
|
return False
|
|
|
|
# 根据名称去重
|
|
unique_terms = {}
|
|
for term in terms:
|
|
name = term.get("name", "")
|
|
if name and name not in unique_terms:
|
|
unique_terms[name] = term
|
|
|
|
# 转换回列表
|
|
deduplicated_terms = list(unique_terms.values())
|
|
logging.info(f"去重后剩余 {len(deduplicated_terms)} 条专业名词")
|
|
|
|
# 准备数据
|
|
texts, metadatas = self._prepare_terms_for_faiss(deduplicated_terms)
|
|
|
|
# 创建索引
|
|
faiss_index = self._create_index(texts, metadatas)
|
|
|
|
# 保存索引
|
|
self._save_index(faiss_index)
|
|
|
|
logging.info("完成多文件专业名词向量化和索引创建")
|
|
return True
|
|
except Exception as e:
|
|
logging.error(f"多文件向量化处理失败: {e}")
|
|
return False
|
|
|
|
|
|
def _prepare_terms_for_faiss(self, terms: List[Dict[str, Any]]) -> Tuple[List[str], List[Dict]]:
|
|
"""
|
|
将术语准备为FAISS可用的格式 (内部方法)
|
|
|
|
Args:
|
|
terms: 术语列表
|
|
|
|
Returns:
|
|
格式化的术语文本列表和元数据列表
|
|
"""
|
|
texts = []
|
|
metadatas = []
|
|
|
|
for term in terms:
|
|
name = term["name"]
|
|
texts.append(name.strip())
|
|
synonyms = term.get("synonymous", [])
|
|
description = term.get("description", "")
|
|
# 记录元数据
|
|
metadatas.append({
|
|
"name": name,
|
|
"synonyms": synonyms,
|
|
"description": description
|
|
})
|
|
|
|
if len(synonyms) > 0:
|
|
synonyms_str = ', '.join(synonyms)
|
|
texts.append(synonyms_str.strip())
|
|
metadatas.append({
|
|
"name": name,
|
|
"synonyms": synonyms,
|
|
"description": description
|
|
})
|
|
|
|
if len(description) > 0:
|
|
texts.append(description.strip())
|
|
metadatas.append({
|
|
"name": name,
|
|
"synonyms": synonyms,
|
|
"description": description
|
|
})
|
|
|
|
return texts, metadatas
|
|
|
|
def _create_index(self, texts: List[str], metadatas: List[Dict]) -> FAISS:
|
|
"""
|
|
创建FAISS索引 (内部方法)
|
|
|
|
Args:
|
|
texts: 文本列表
|
|
metadatas: 元数据列表
|
|
|
|
Returns:
|
|
FAISS索引
|
|
"""
|
|
logging.info(f"正在创建FAISS索引，共 {len(texts)} 条数据...")
|
|
return FAISS.from_texts(texts=texts, embedding=self.embedding_model, metadatas=metadatas)
|
|
|
|
def _save_index(self, faiss_index: FAISS) -> None:
|
|
"""
|
|
保存FAISS索引到本地 (内部方法)
|
|
|
|
Args:
|
|
faiss_index: 要保存的FAISS索引
|
|
"""
|
|
try:
|
|
# 确保输出目录存在
|
|
os.makedirs(self.output_dir, exist_ok=True)
|
|
|
|
# 如果索引目录已存在，先删除
|
|
if os.path.exists(self.index_path):
|
|
shutil.rmtree(self.index_path)
|
|
|
|
# 保存FAISS索引
|
|
faiss_index.save_local(folder_path=self.index_path)
|
|
logging.info(f"FAISS索引已保存至 {self.index_path}")
|
|
except Exception as e:
|
|
logging.error(f"保存FAISS索引失败: {e}")
|
|
raise e
|
|
|
|
|
|
class ProfessionalNounRetriever:
|
|
"""专业名词检索类"""
|
|
|
|
def __init__(self,
|
|
embedding_model: Optional[Embeddings] = None,
|
|
api_key: str = None,
|
|
index_dir: str = None):
|
|
"""
|
|
初始化检索器并加载索引
|
|
|
|
Args:
|
|
embedding_model: 嵌入模型，如果为None则使用默认模型
|
|
api_key: SiliconFlow API密钥，仅在embedding_model为None时使用
|
|
index_dir: 索引目录路径，默认为None使用默认路径
|
|
"""
|
|
# 设置嵌入模型
|
|
if embedding_model:
|
|
self.embedding_model = embedding_model
|
|
else:
|
|
self.embedding_model = get_embedding_model(api_key)
|
|
|
|
# 设置索引路径
|
|
self.index_dir = index_dir
|
|
if self.index_dir is None:
|
|
current_dir = os.path.dirname(os.path.abspath(__file__))
|
|
self.index_dir = os.path.join(current_dir, "..", "..", "data", "nouns", "professional_nouns_index")
|
|
|
|
# 在构造函数中加载索引
|
|
self.faiss_index = None
|
|
self._load_index()
|
|
|
|
def _load_index(self) -> None:
|
|
"""
|
|
从本地加载FAISS索引 (内部方法)
|
|
"""
|
|
try:
|
|
# 加载FAISS索引，启用不安全反序列化（仅用于可信数据源）
|
|
self.faiss_index = FAISS.load_local(
|
|
folder_path=self.index_dir,
|
|
embeddings=self.embedding_model,
|
|
allow_dangerous_deserialization=True
|
|
)
|
|
logging.info(f"成功从 {self.index_dir} 加载FAISS索引")
|
|
except Exception as e:
|
|
logging.warning(f"加载FAISS索引失败: {e}")
|
|
self.faiss_index = None
|
|
|
|
def query(self, query_text: str, top_k: int = 5, use_intersection: bool = True) -> List[Dict]:
|
|
"""
|
|
查询FAISS索引，获取最相似的专业名词 (唯一对外接口)
|
|
|
|
Args:
|
|
query_text: 查询文本
|
|
top_k: 返回的结果数量，默认为5
|
|
use_intersection: 是否使用三种检索方式的交集，默认为True
|
|
|
|
Returns:
|
|
相似度最高的专业名词列表
|
|
"""
|
|
try:
|
|
# 检查索引是否已加载
|
|
if self.faiss_index is None:
|
|
logging.warning("FAISS索引未加载，无法执行查询")
|
|
return []
|
|
|
|
# 使用三种检索方式并取交集
|
|
retriever1 = self.faiss_index.as_retriever(search_kwargs={"k": top_k})
|
|
retriever2 = self.faiss_index.as_retriever(
|
|
search_type="mmr",
|
|
search_kwargs={"k": top_k, "fetch_k": 3, "lambda_mult": 0.5}
|
|
)
|
|
retriever3 = self.faiss_index.as_retriever(
|
|
search_type="similarity_score_threshold",
|
|
search_kwargs={"score_threshold": 0.5}
|
|
)
|
|
|
|
# 用json.dumps将dict转为字符串，便于取交集
|
|
set1 = set(json.dumps(i.metadata, sort_keys=True, ensure_ascii=False)
|
|
for i in retriever1.invoke(query_text))
|
|
set2 = set(json.dumps(i.metadata, sort_keys=True, ensure_ascii=False)
|
|
for i in retriever2.invoke(query_text))
|
|
set3 = set(json.dumps(i.metadata, sort_keys=True, ensure_ascii=False)
|
|
for i in retriever3.invoke(query_text))
|
|
|
|
intersection = set1 | set2 | set3
|
|
|
|
# 如果交集为空，使用第一种检索方式的结果
|
|
if not intersection:
|
|
logging.warning("三种检索方式无交集，使用普通检索结果")
|
|
return [json.loads(item) for item in set1]
|
|
|
|
# 转回dict
|
|
return [json.loads(item) for item in intersection]
|
|
|
|
except Exception as e:
|
|
logging.error(f"查询FAISS索引失败: {e}")
|
|
return [] |