初始化
This commit is contained in:
雅诺狐
592
src/chat/knowledge/embedding_store.py
Normal file
592
src/chat/knowledge/embedding_store.py
Normal file
@@ -0,0 +1,592 @@
|
||||
from dataclasses import dataclass
|
||||
import json
|
||||
import os
|
||||
import math
|
||||
import asyncio
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
from typing import Dict, List, Tuple
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
|
||||
# import tqdm
|
||||
import faiss
|
||||
|
||||
from .utils.hash import get_sha256
|
||||
from .global_logger import logger
|
||||
from rich.traceback import install
|
||||
from rich.progress import (
|
||||
Progress,
|
||||
BarColumn,
|
||||
TimeElapsedColumn,
|
||||
TimeRemainingColumn,
|
||||
TaskProgressColumn,
|
||||
MofNCompleteColumn,
|
||||
SpinnerColumn,
|
||||
TextColumn,
|
||||
)
|
||||
from src.chat.utils.utils import get_embedding
|
||||
from src.config.config import global_config
|
||||
|
||||
|
||||
install(extra_lines=3)
|
||||
|
||||
# 多线程embedding配置常量
|
||||
DEFAULT_MAX_WORKERS = 10 # 默认最大线程数
|
||||
DEFAULT_CHUNK_SIZE = 10 # 默认每个线程处理的数据块大小
|
||||
MIN_CHUNK_SIZE = 1 # 最小分块大小
|
||||
MAX_CHUNK_SIZE = 50 # 最大分块大小
|
||||
MIN_WORKERS = 1 # 最小线程数
|
||||
MAX_WORKERS = 20 # 最大线程数
|
||||
|
||||
ROOT_PATH = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", ".."))
|
||||
EMBEDDING_DATA_DIR = os.path.join(ROOT_PATH, "data", "embedding")
|
||||
EMBEDDING_DATA_DIR_STR = str(EMBEDDING_DATA_DIR).replace("\\", "/")
|
||||
TOTAL_EMBEDDING_TIMES = 3 # 统计嵌入次数
|
||||
|
||||
# 嵌入模型测试字符串，测试模型一致性，来自开发群的聊天记录
|
||||
# 这些字符串的嵌入结果应该是固定的，不能随时间变化
|
||||
EMBEDDING_TEST_STRINGS = [
|
||||
"阿卡伊真的太好玩了，神秘性感大女同等着你",
|
||||
"你怎么知道我arc12.64了",
|
||||
"我是蕾缪乐小姐的狗",
|
||||
"关注Oct谢谢喵",
|
||||
"不是w6我不草",
|
||||
"关注千石可乐谢谢喵",
|
||||
"来玩CLANNAD，AIR，樱之诗，樱之刻谢谢喵",
|
||||
"关注墨梓柒谢谢喵",
|
||||
"Ciallo~",
|
||||
"来玩巧克甜恋谢谢喵",
|
||||
"水印",
|
||||
"我也在纠结晚饭，铁锅炒鸡听着就香！",
|
||||
"test你妈喵",
|
||||
]
|
||||
EMBEDDING_TEST_FILE = os.path.join(ROOT_PATH, "data", "embedding_model_test.json")
|
||||
EMBEDDING_SIM_THRESHOLD = 0.99
|
||||
|
||||
|
||||
def cosine_similarity(a, b):
|
||||
# 计算余弦相似度
|
||||
dot = sum(x * y for x, y in zip(a, b, strict=False))
|
||||
norm_a = math.sqrt(sum(x * x for x in a))
|
||||
norm_b = math.sqrt(sum(x * x for x in b))
|
||||
if norm_a == 0 or norm_b == 0:
|
||||
return 0.0
|
||||
return dot / (norm_a * norm_b)
|
||||
|
||||
|
||||
@dataclass
|
||||
class EmbeddingStoreItem:
|
||||
"""嵌入库中的项"""
|
||||
|
||||
def __init__(self, item_hash: str, embedding: List[float], content: str):
|
||||
self.hash = item_hash
|
||||
self.embedding = embedding
|
||||
self.str = content
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
"""转为dict"""
|
||||
return {
|
||||
"hash": self.hash,
|
||||
"embedding": self.embedding,
|
||||
"str": self.str,
|
||||
}
|
||||
|
||||
|
||||
class EmbeddingStore:
|
||||
def __init__(self, namespace: str, dir_path: str, max_workers: int = DEFAULT_MAX_WORKERS, chunk_size: int = DEFAULT_CHUNK_SIZE):
|
||||
self.namespace = namespace
|
||||
self.dir = dir_path
|
||||
self.embedding_file_path = f"{dir_path}/{namespace}.parquet"
|
||||
self.index_file_path = f"{dir_path}/{namespace}.index"
|
||||
self.idx2hash_file_path = dir_path + "/" + namespace + "_i2h.json"
|
||||
|
||||
# 多线程配置参数验证和设置
|
||||
self.max_workers = max(MIN_WORKERS, min(MAX_WORKERS, max_workers))
|
||||
self.chunk_size = max(MIN_CHUNK_SIZE, min(MAX_CHUNK_SIZE, chunk_size))
|
||||
|
||||
# 如果配置值被调整，记录日志
|
||||
if self.max_workers != max_workers:
|
||||
logger.warning(f"max_workers 已从 {max_workers} 调整为 {self.max_workers} (范围: {MIN_WORKERS}-{MAX_WORKERS})")
|
||||
if self.chunk_size != chunk_size:
|
||||
logger.warning(f"chunk_size 已从 {chunk_size} 调整为 {self.chunk_size} (范围: {MIN_CHUNK_SIZE}-{MAX_CHUNK_SIZE})")
|
||||
|
||||
self.store = {}
|
||||
|
||||
self.faiss_index = None
|
||||
self.idx2hash = None
|
||||
|
||||
def _get_embedding(self, s: str) -> List[float]:
|
||||
"""获取字符串的嵌入向量，处理异步调用"""
|
||||
try:
|
||||
# 尝试获取当前事件循环
|
||||
asyncio.get_running_loop()
|
||||
# 如果在事件循环中，使用线程池执行
|
||||
import concurrent.futures
|
||||
|
||||
def run_in_thread():
|
||||
return asyncio.run(get_embedding(s))
|
||||
|
||||
with concurrent.futures.ThreadPoolExecutor() as executor:
|
||||
future = executor.submit(run_in_thread)
|
||||
result = future.result()
|
||||
if result is None:
|
||||
logger.error(f"获取嵌入失败: {s}")
|
||||
return []
|
||||
return result
|
||||
except RuntimeError:
|
||||
# 没有运行的事件循环，直接运行
|
||||
result = asyncio.run(get_embedding(s))
|
||||
if result is None:
|
||||
logger.error(f"获取嵌入失败: {s}")
|
||||
return []
|
||||
return result
|
||||
|
||||
def _get_embeddings_batch_threaded(self, strs: List[str], chunk_size: int = 10, max_workers: int = 10, progress_callback=None) -> List[Tuple[str, List[float]]]:
|
||||
"""使用多线程批量获取嵌入向量
|
||||
|
||||
Args:
|
||||
strs: 要获取嵌入的字符串列表
|
||||
chunk_size: 每个线程处理的数据块大小
|
||||
max_workers: 最大线程数
|
||||
progress_callback: 进度回调函数，接收一个参数表示完成的数量
|
||||
|
||||
Returns:
|
||||
包含(原始字符串, 嵌入向量)的元组列表，保持与输入顺序一致
|
||||
"""
|
||||
if not strs:
|
||||
return []
|
||||
|
||||
# 分块
|
||||
chunks = []
|
||||
for i in range(0, len(strs), chunk_size):
|
||||
chunk = strs[i:i + chunk_size]
|
||||
chunks.append((i, chunk)) # 保存起始索引以维持顺序
|
||||
|
||||
# 结果存储，使用字典按索引存储以保证顺序
|
||||
results = {}
|
||||
|
||||
def process_chunk(chunk_data):
|
||||
"""处理单个数据块的函数"""
|
||||
start_idx, chunk_strs = chunk_data
|
||||
chunk_results = []
|
||||
|
||||
# 为每个线程创建独立的LLMRequest实例
|
||||
from src.llm_models.utils_model import LLMRequest
|
||||
from src.config.config import model_config
|
||||
|
||||
try:
|
||||
# 创建线程专用的LLM实例
|
||||
llm = LLMRequest(model_set=model_config.model_task_config.embedding, request_type="embedding")
|
||||
|
||||
for i, s in enumerate(chunk_strs):
|
||||
try:
|
||||
# 直接使用异步函数
|
||||
embedding = asyncio.run(llm.get_embedding(s))
|
||||
if embedding and len(embedding) > 0:
|
||||
chunk_results.append((start_idx + i, s, embedding[0])) # embedding[0] 是实际的向量
|
||||
else:
|
||||
logger.error(f"获取嵌入失败: {s}")
|
||||
chunk_results.append((start_idx + i, s, []))
|
||||
|
||||
# 每完成一个嵌入立即更新进度
|
||||
if progress_callback:
|
||||
progress_callback(1)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"获取嵌入时发生异常: {s}, 错误: {e}")
|
||||
chunk_results.append((start_idx + i, s, []))
|
||||
|
||||
# 即使失败也要更新进度
|
||||
if progress_callback:
|
||||
progress_callback(1)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"创建LLM实例失败: {e}")
|
||||
# 如果创建LLM实例失败，返回空结果
|
||||
for i, s in enumerate(chunk_strs):
|
||||
chunk_results.append((start_idx + i, s, []))
|
||||
# 即使失败也要更新进度
|
||||
if progress_callback:
|
||||
progress_callback(1)
|
||||
|
||||
return chunk_results
|
||||
|
||||
# 使用线程池处理
|
||||
with ThreadPoolExecutor(max_workers=max_workers) as executor:
|
||||
# 提交所有任务
|
||||
future_to_chunk = {executor.submit(process_chunk, chunk): chunk for chunk in chunks}
|
||||
|
||||
# 收集结果（进度已在process_chunk中实时更新）
|
||||
for future in as_completed(future_to_chunk):
|
||||
try:
|
||||
chunk_results = future.result()
|
||||
for idx, s, embedding in chunk_results:
|
||||
results[idx] = (s, embedding)
|
||||
except Exception as e:
|
||||
chunk = future_to_chunk[future]
|
||||
logger.error(f"处理数据块时发生异常: {chunk}, 错误: {e}")
|
||||
# 为失败的块添加空结果
|
||||
start_idx, chunk_strs = chunk
|
||||
for i, s in enumerate(chunk_strs):
|
||||
results[start_idx + i] = (s, [])
|
||||
|
||||
# 按原始顺序返回结果
|
||||
ordered_results = []
|
||||
for i in range(len(strs)):
|
||||
if i in results:
|
||||
ordered_results.append(results[i])
|
||||
else:
|
||||
# 防止遗漏
|
||||
ordered_results.append((strs[i], []))
|
||||
|
||||
return ordered_results
|
||||
|
||||
def get_test_file_path(self):
|
||||
return EMBEDDING_TEST_FILE
|
||||
|
||||
def save_embedding_test_vectors(self):
|
||||
"""保存测试字符串的嵌入到本地（使用多线程优化）"""
|
||||
logger.info("开始保存测试字符串的嵌入向量...")
|
||||
|
||||
# 使用多线程批量获取测试字符串的嵌入
|
||||
embedding_results = self._get_embeddings_batch_threaded(
|
||||
EMBEDDING_TEST_STRINGS,
|
||||
chunk_size=min(self.chunk_size, len(EMBEDDING_TEST_STRINGS)),
|
||||
max_workers=min(self.max_workers, len(EMBEDDING_TEST_STRINGS))
|
||||
)
|
||||
|
||||
# 构建测试向量字典
|
||||
test_vectors = {}
|
||||
for idx, (s, embedding) in enumerate(embedding_results):
|
||||
if embedding:
|
||||
test_vectors[str(idx)] = embedding
|
||||
else:
|
||||
logger.error(f"获取测试字符串嵌入失败: {s}")
|
||||
# 使用原始单线程方法作为后备
|
||||
test_vectors[str(idx)] = self._get_embedding(s)
|
||||
|
||||
with open(self.get_test_file_path(), "w", encoding="utf-8") as f:
|
||||
json.dump(test_vectors, f, ensure_ascii=False, indent=2)
|
||||
|
||||
logger.info("测试字符串嵌入向量保存完成")
|
||||
|
||||
def load_embedding_test_vectors(self):
|
||||
"""加载本地保存的测试字符串嵌入"""
|
||||
path = self.get_test_file_path()
|
||||
if not os.path.exists(path):
|
||||
return None
|
||||
with open(path, "r", encoding="utf-8") as f:
|
||||
return json.load(f)
|
||||
|
||||
def check_embedding_model_consistency(self):
|
||||
"""校验当前模型与本地嵌入模型是否一致（使用多线程优化）"""
|
||||
local_vectors = self.load_embedding_test_vectors()
|
||||
if local_vectors is None:
|
||||
logger.warning("未检测到本地嵌入模型测试文件，将保存当前模型的测试嵌入。")
|
||||
self.save_embedding_test_vectors()
|
||||
return True
|
||||
|
||||
# 检查本地向量完整性
|
||||
for idx in range(len(EMBEDDING_TEST_STRINGS)):
|
||||
if local_vectors.get(str(idx)) is None:
|
||||
logger.warning("本地嵌入模型测试文件缺失部分测试字符串，将重新保存。")
|
||||
self.save_embedding_test_vectors()
|
||||
return True
|
||||
|
||||
logger.info("开始检验嵌入模型一致性...")
|
||||
|
||||
# 使用多线程批量获取当前模型的嵌入
|
||||
embedding_results = self._get_embeddings_batch_threaded(
|
||||
EMBEDDING_TEST_STRINGS,
|
||||
chunk_size=min(self.chunk_size, len(EMBEDDING_TEST_STRINGS)),
|
||||
max_workers=min(self.max_workers, len(EMBEDDING_TEST_STRINGS))
|
||||
)
|
||||
|
||||
# 检查一致性
|
||||
for idx, (s, new_emb) in enumerate(embedding_results):
|
||||
local_emb = local_vectors.get(str(idx))
|
||||
if not new_emb:
|
||||
logger.error(f"获取测试字符串嵌入失败: {s}")
|
||||
return False
|
||||
|
||||
sim = cosine_similarity(local_emb, new_emb)
|
||||
if sim < EMBEDDING_SIM_THRESHOLD:
|
||||
logger.error(f"嵌入模型一致性校验失败，字符串: {s}, 相似度: {sim:.4f}")
|
||||
return False
|
||||
|
||||
logger.info("嵌入模型一致性校验通过。")
|
||||
return True
|
||||
|
||||
def batch_insert_strs(self, strs: List[str], times: int) -> None:
|
||||
"""向库中存入字符串（使用多线程优化）"""
|
||||
if not strs:
|
||||
return
|
||||
|
||||
total = len(strs)
|
||||
|
||||
# 过滤已存在的字符串
|
||||
new_strs = []
|
||||
for s in strs:
|
||||
item_hash = self.namespace + "-" + get_sha256(s)
|
||||
if item_hash not in self.store:
|
||||
new_strs.append(s)
|
||||
|
||||
if not new_strs:
|
||||
logger.info(f"所有字符串已存在于{self.namespace}嵌入库中，跳过处理")
|
||||
return
|
||||
|
||||
logger.info(f"需要处理 {len(new_strs)}/{total} 个新字符串")
|
||||
|
||||
with Progress(
|
||||
SpinnerColumn(),
|
||||
TextColumn("[progress.description]{task.description}"),
|
||||
BarColumn(),
|
||||
TaskProgressColumn(),
|
||||
MofNCompleteColumn(),
|
||||
"•",
|
||||
TimeElapsedColumn(),
|
||||
"<",
|
||||
TimeRemainingColumn(),
|
||||
transient=False,
|
||||
) as progress:
|
||||
task = progress.add_task(f"存入嵌入库：({times}/{TOTAL_EMBEDDING_TIMES})", total=total)
|
||||
|
||||
# 首先更新已存在项的进度
|
||||
already_processed = total - len(new_strs)
|
||||
if already_processed > 0:
|
||||
progress.update(task, advance=already_processed)
|
||||
|
||||
if new_strs:
|
||||
# 使用实例配置的参数，智能调整分块和线程数
|
||||
optimal_chunk_size = max(MIN_CHUNK_SIZE, min(self.chunk_size, len(new_strs) // self.max_workers if self.max_workers > 0 else self.chunk_size))
|
||||
optimal_max_workers = min(self.max_workers, max(MIN_WORKERS, len(new_strs) // optimal_chunk_size if optimal_chunk_size > 0 else 1))
|
||||
|
||||
logger.debug(f"使用多线程处理: chunk_size={optimal_chunk_size}, max_workers={optimal_max_workers}")
|
||||
|
||||
# 定义进度更新回调函数
|
||||
def update_progress(count):
|
||||
progress.update(task, advance=count)
|
||||
|
||||
# 批量获取嵌入，并实时更新进度
|
||||
embedding_results = self._get_embeddings_batch_threaded(
|
||||
new_strs,
|
||||
chunk_size=optimal_chunk_size,
|
||||
max_workers=optimal_max_workers,
|
||||
progress_callback=update_progress
|
||||
)
|
||||
|
||||
# 存入结果（不再需要在这里更新进度，因为已经在回调中更新了）
|
||||
for s, embedding in embedding_results:
|
||||
item_hash = self.namespace + "-" + get_sha256(s)
|
||||
if embedding: # 只有成功获取到嵌入才存入
|
||||
self.store[item_hash] = EmbeddingStoreItem(item_hash, embedding, s)
|
||||
else:
|
||||
logger.warning(f"跳过存储失败的嵌入: {s[:50]}...")
|
||||
|
||||
def save_to_file(self) -> None:
|
||||
"""保存到文件"""
|
||||
data = []
|
||||
logger.info(f"正在保存{self.namespace}嵌入库到文件{self.embedding_file_path}")
|
||||
for item in self.store.values():
|
||||
data.append(item.to_dict())
|
||||
data_frame = pd.DataFrame(data)
|
||||
|
||||
if not os.path.exists(self.dir):
|
||||
os.makedirs(self.dir, exist_ok=True)
|
||||
if not os.path.exists(self.embedding_file_path):
|
||||
open(self.embedding_file_path, "w").close()
|
||||
|
||||
data_frame.to_parquet(self.embedding_file_path, engine="pyarrow", index=False)
|
||||
logger.info(f"{self.namespace}嵌入库保存成功")
|
||||
|
||||
if self.faiss_index is not None and self.idx2hash is not None:
|
||||
logger.info(f"正在保存{self.namespace}嵌入库的FaissIndex到文件{self.index_file_path}")
|
||||
faiss.write_index(self.faiss_index, self.index_file_path)
|
||||
logger.info(f"{self.namespace}嵌入库的FaissIndex保存成功")
|
||||
logger.info(f"正在保存{self.namespace}嵌入库的idx2hash映射到文件{self.idx2hash_file_path}")
|
||||
with open(self.idx2hash_file_path, "w", encoding="utf-8") as f:
|
||||
f.write(json.dumps(self.idx2hash, ensure_ascii=False, indent=4))
|
||||
logger.info(f"{self.namespace}嵌入库的idx2hash映射保存成功")
|
||||
|
||||
def load_from_file(self) -> None:
|
||||
"""从文件中加载"""
|
||||
if not os.path.exists(self.embedding_file_path):
|
||||
raise Exception(f"文件{self.embedding_file_path}不存在")
|
||||
logger.info("正在加载嵌入库...")
|
||||
logger.debug(f"正在从文件{self.embedding_file_path}中加载{self.namespace}嵌入库")
|
||||
data_frame = pd.read_parquet(self.embedding_file_path, engine="pyarrow")
|
||||
total = len(data_frame)
|
||||
with Progress(
|
||||
SpinnerColumn(),
|
||||
TextColumn("[progress.description]{task.description}"),
|
||||
BarColumn(),
|
||||
TaskProgressColumn(),
|
||||
MofNCompleteColumn(),
|
||||
"•",
|
||||
TimeElapsedColumn(),
|
||||
"<",
|
||||
TimeRemainingColumn(),
|
||||
transient=False,
|
||||
) as progress:
|
||||
task = progress.add_task("加载嵌入库", total=total)
|
||||
for _, row in data_frame.iterrows():
|
||||
self.store[row["hash"]] = EmbeddingStoreItem(row["hash"], row["embedding"], row["str"])
|
||||
progress.update(task, advance=1)
|
||||
logger.info(f"{self.namespace}嵌入库加载成功")
|
||||
|
||||
try:
|
||||
if os.path.exists(self.index_file_path):
|
||||
logger.info(f"正在加载{self.namespace}嵌入库的FaissIndex...")
|
||||
logger.debug(f"正在从文件{self.index_file_path}中加载{self.namespace}嵌入库的FaissIndex")
|
||||
self.faiss_index = faiss.read_index(self.index_file_path)
|
||||
logger.info(f"{self.namespace}嵌入库的FaissIndex加载成功")
|
||||
else:
|
||||
raise Exception(f"文件{self.index_file_path}不存在")
|
||||
if os.path.exists(self.idx2hash_file_path):
|
||||
logger.info(f"正在加载{self.namespace}嵌入库的idx2hash映射...")
|
||||
logger.debug(f"正在从文件{self.idx2hash_file_path}中加载{self.namespace}嵌入库的idx2hash映射")
|
||||
with open(self.idx2hash_file_path, "r") as f:
|
||||
self.idx2hash = json.load(f)
|
||||
logger.info(f"{self.namespace}嵌入库的idx2hash映射加载成功")
|
||||
else:
|
||||
raise Exception(f"文件{self.idx2hash_file_path}不存在")
|
||||
except Exception as e:
|
||||
logger.error(f"加载{self.namespace}嵌入库的FaissIndex时发生错误：{e}")
|
||||
logger.warning("正在重建Faiss索引")
|
||||
self.build_faiss_index()
|
||||
logger.info(f"{self.namespace}嵌入库的FaissIndex重建成功")
|
||||
self.save_to_file()
|
||||
|
||||
def build_faiss_index(self) -> None:
|
||||
"""重新构建Faiss索引，以余弦相似度为度量"""
|
||||
# 获取所有的embedding
|
||||
array = []
|
||||
self.idx2hash = dict()
|
||||
for key in self.store:
|
||||
array.append(self.store[key].embedding)
|
||||
self.idx2hash[str(len(array) - 1)] = key
|
||||
embeddings = np.array(array, dtype=np.float32)
|
||||
# L2归一化
|
||||
faiss.normalize_L2(embeddings)
|
||||
# 构建索引
|
||||
self.faiss_index = faiss.IndexFlatIP(global_config.lpmm_knowledge.embedding_dimension)
|
||||
self.faiss_index.add(embeddings)
|
||||
|
||||
def search_top_k(self, query: List[float], k: int) -> List[Tuple[str, float]]:
|
||||
"""搜索最相似的k个项，以余弦相似度为度量
|
||||
Args:
|
||||
query: 查询的embedding
|
||||
k: 返回的最相似的k个项
|
||||
Returns:
|
||||
result: 最相似的k个项的(hash, 余弦相似度)列表
|
||||
"""
|
||||
if self.faiss_index is None:
|
||||
logger.debug("FaissIndex尚未构建,返回None")
|
||||
return []
|
||||
if self.idx2hash is None:
|
||||
logger.warning("idx2hash尚未构建,返回None")
|
||||
return []
|
||||
|
||||
# L2归一化
|
||||
faiss.normalize_L2(np.array([query], dtype=np.float32))
|
||||
# 搜索
|
||||
distances, indices = self.faiss_index.search(np.array([query]), k)
|
||||
# 整理结果
|
||||
indices = list(indices.flatten())
|
||||
distances = list(distances.flatten())
|
||||
result = [
|
||||
(self.idx2hash[str(int(idx))], float(sim))
|
||||
for (idx, sim) in zip(indices, distances, strict=False)
|
||||
if idx in range(len(self.idx2hash))
|
||||
]
|
||||
|
||||
return result
|
||||
|
||||
|
||||
class EmbeddingManager:
|
||||
def __init__(self, max_workers: int = DEFAULT_MAX_WORKERS, chunk_size: int = DEFAULT_CHUNK_SIZE):
|
||||
"""
|
||||
初始化EmbeddingManager
|
||||
|
||||
Args:
|
||||
max_workers: 最大线程数
|
||||
chunk_size: 每个线程处理的数据块大小
|
||||
"""
|
||||
self.paragraphs_embedding_store = EmbeddingStore(
|
||||
"paragraph", # type: ignore
|
||||
EMBEDDING_DATA_DIR_STR,
|
||||
max_workers=max_workers,
|
||||
chunk_size=chunk_size,
|
||||
)
|
||||
self.entities_embedding_store = EmbeddingStore(
|
||||
"entity", # type: ignore
|
||||
EMBEDDING_DATA_DIR_STR,
|
||||
max_workers=max_workers,
|
||||
chunk_size=chunk_size,
|
||||
)
|
||||
self.relation_embedding_store = EmbeddingStore(
|
||||
"relation", # type: ignore
|
||||
EMBEDDING_DATA_DIR_STR,
|
||||
max_workers=max_workers,
|
||||
chunk_size=chunk_size,
|
||||
)
|
||||
self.stored_pg_hashes = set()
|
||||
|
||||
def check_all_embedding_model_consistency(self):
|
||||
"""对所有嵌入库做模型一致性校验"""
|
||||
return self.paragraphs_embedding_store.check_embedding_model_consistency()
|
||||
|
||||
def _store_pg_into_embedding(self, raw_paragraphs: Dict[str, str]):
|
||||
"""将段落编码存入Embedding库"""
|
||||
self.paragraphs_embedding_store.batch_insert_strs(list(raw_paragraphs.values()), times=1)
|
||||
|
||||
def _store_ent_into_embedding(self, triple_list_data: Dict[str, List[List[str]]]):
|
||||
"""将实体编码存入Embedding库"""
|
||||
entities = set()
|
||||
for triple_list in triple_list_data.values():
|
||||
for triple in triple_list:
|
||||
entities.add(triple[0])
|
||||
entities.add(triple[2])
|
||||
self.entities_embedding_store.batch_insert_strs(list(entities), times=2)
|
||||
|
||||
def _store_rel_into_embedding(self, triple_list_data: Dict[str, List[List[str]]]):
|
||||
"""将关系编码存入Embedding库"""
|
||||
graph_triples = [] # a list of unique relation triple (in tuple) from all chunks
|
||||
for triples in triple_list_data.values():
|
||||
graph_triples.extend([tuple(t) for t in triples])
|
||||
graph_triples = list(set(graph_triples))
|
||||
self.relation_embedding_store.batch_insert_strs([str(triple) for triple in graph_triples], times=3)
|
||||
|
||||
def load_from_file(self):
|
||||
"""从文件加载"""
|
||||
self.paragraphs_embedding_store.load_from_file()
|
||||
self.entities_embedding_store.load_from_file()
|
||||
self.relation_embedding_store.load_from_file()
|
||||
# 从段落库中获取已存储的hash
|
||||
self.stored_pg_hashes = set(self.paragraphs_embedding_store.store.keys())
|
||||
|
||||
def store_new_data_set(
|
||||
self,
|
||||
raw_paragraphs: Dict[str, str],
|
||||
triple_list_data: Dict[str, List[List[str]]],
|
||||
):
|
||||
if not self.check_all_embedding_model_consistency():
|
||||
raise Exception("嵌入模型与本地存储不一致，请检查模型设置或清空嵌入库后重试。")
|
||||
"""存储新的数据集"""
|
||||
self._store_pg_into_embedding(raw_paragraphs)
|
||||
self._store_ent_into_embedding(triple_list_data)
|
||||
self._store_rel_into_embedding(triple_list_data)
|
||||
self.stored_pg_hashes.update(raw_paragraphs.keys())
|
||||
|
||||
def save_to_file(self):
|
||||
"""保存到文件"""
|
||||
self.paragraphs_embedding_store.save_to_file()
|
||||
self.entities_embedding_store.save_to_file()
|
||||
self.relation_embedding_store.save_to_file()
|
||||
|
||||
def rebuild_faiss_index(self):
|
||||
"""重建Faiss索引（请在添加新数据后调用）"""
|
||||
self.paragraphs_embedding_store.build_faiss_index()
|
||||
self.entities_embedding_store.build_faiss_index()
|
||||
self.relation_embedding_store.build_faiss_index()
|
||||
Reference in New Issue
Block a user