gardel/Mofox-Core
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Revert "总之就是知识库"

Browse Source 
This reverts commit 0383a999fb.
This commit is contained in:
tt-P607
2025-10-18 17:09:18 +08:00
parent 72ea362e29
commit 373db77768
2 changed files with 153 additions and 114 deletions
Download Patch File Download Diff File Expand all files Collapse all files

91
scripts/lpmm_learning_tool.py
View File

@@ -191,7 +191,11 @@ async def extract_info_async(pg_hash, paragraph, llm_api):
return None, pg_hash
async def extract_information(paragraphs_dict, model_set):
def extract_info_sync(pg_hash, paragraph, llm_api):
return asyncio.run(extract_info_async(pg_hash, paragraph, llm_api))
def extract_information(paragraphs_dict, model_set):
logger.info("--- 步骤 2: 开始信息提取 ---")
os.makedirs(OPENIE_OUTPUT_DIR, exist_ok=True)
os.makedirs(TEMP_DIR, exist_ok=True)
@@ -199,35 +203,32 @@ async def extract_information(paragraphs_dict, model_set):
llm_api = LLMRequest(model_set=model_set)
failed_hashes, open_ie_docs = [], []
tasks = [
extract_info_async(p_hash, p, llm_api)
for p_hash, p in paragraphs_dict.items()
]
with Progress(
SpinnerColumn(),
TextColumn("[progress.description]{task.description}"),
BarColumn(),
TaskProgressColumn(),
MofNCompleteColumn(),
"•",
TimeElapsedColumn(),
"<",
TimeRemainingColumn(),
) as progress:
prog_task = progress.add_task("[cyan]正在提取信息...", total=len(tasks))
for future in asyncio.as_completed(tasks):
doc_item, failed_hash = await future
if failed_hash:
failed_hashes.append(failed_hash)
elif doc_item:
open_ie_docs.append(doc_item)
progress.update(prog_task, advance=1)
with ThreadPoolExecutor(max_workers=5) as executor:
f_to_hash = {
executor.submit(extract_info_sync, p_hash, p, llm_api): p_hash for p_hash, p in paragraphs_dict.items()
}
with Progress(
SpinnerColumn(),
TextColumn("[progress.description]{task.description}"),
BarColumn(),
TaskProgressColumn(),
MofNCompleteColumn(),
"•",
TimeElapsedColumn(),
"<",
TimeRemainingColumn(),
) as progress:
task = progress.add_task("[cyan]正在提取信息...", total=len(paragraphs_dict))
for future in as_completed(f_to_hash):
doc_item, failed_hash = future.result()
if failed_hash:
failed_hashes.append(failed_hash)
elif doc_item:
open_ie_docs.append(doc_item)
progress.update(task, advance=1)
if open_ie_docs:
all_entities = [
e for doc in open_ie_docs for e in doc["extracted_entities"]
]
all_entities = [e for doc in open_ie_docs for e in doc["extracted_entities"]]
num_entities = len(all_entities)
avg_ent_chars = round(sum(len(e) for e in all_entities) / num_entities, 4) if num_entities else 0
avg_ent_words = round(sum(len(e.split()) for e in all_entities) / num_entities, 4) if num_entities else 0
@@ -312,7 +313,7 @@ async def import_data(openie_obj: OpenIE | None = None):
logger.info("--- 数据导入完成 ---")
async def import_from_specific_file():
def import_from_specific_file():
"""从用户指定的 openie.json 文件导入数据"""
file_path = input("请输入 openie.json 文件的完整路径: ").strip()
@@ -327,7 +328,7 @@ async def import_from_specific_file():
try:
logger.info(f"正在从 {file_path} 加载 OpenIE 数据...")
openie_obj = OpenIE.load()
await import_data(openie_obj=openie_obj)
asyncio.run(import_data(openie_obj=openie_obj))
except Exception as e:
logger.error(f"从指定文件导入数据时发生错误: {e}")
@@ -335,20 +336,14 @@ async def import_from_specific_file():
# --- 主函数 ---
async def async_main():
def main():
# 使用 os.path.relpath 创建相对于项目根目录的友好路径
raw_data_relpath = os.path.relpath(
RAW_DATA_PATH, os.path.join(ROOT_PATH, "..")
)
openie_output_relpath = os.path.relpath(
OPENIE_OUTPUT_DIR, os.path.join(ROOT_PATH, "..")
)
raw_data_relpath = os.path.relpath(RAW_DATA_PATH, os.path.join(ROOT_PATH, ".."))
openie_output_relpath = os.path.relpath(OPENIE_OUTPUT_DIR, os.path.join(ROOT_PATH, ".."))
print("=== LPMM 知识库学习工具 ===")
print(f"1. [数据预处理] -> 读取 .txt 文件 (来源: ./{raw_data_relpath}/)")
print(
f"2. [信息提取] -> 提取信息并存为 .json (输出至: ./{openie_output_relpath}/)"
)
print(f"2. [信息提取] -> 提取信息并存为 .json (输出至: ./{openie_output_relpath}/)")
print("3. [数据导入] -> 从 openie 文件夹自动导入最新知识")
print("4. [全流程] -> 按顺序执行 1 -> 2 -> 3")
print("5. [指定导入] -> 从特定的 openie.json 文件导入知识")
@@ -362,20 +357,16 @@ async def async_main():
elif choice == "2":
paragraphs = preprocess_raw_data()
if paragraphs:
await extract_information(
paragraphs, model_config.model_task_config.lpmm_qa
)
extract_information(paragraphs, model_config.model_task_config.lpmm_qa)
elif choice == "3":
await import_data()
asyncio.run(import_data())
elif choice == "4":
paragraphs = preprocess_raw_data()
if paragraphs:
await extract_information(
paragraphs, model_config.model_task_config.lpmm_qa
)
await import_data()
extract_information(paragraphs, model_config.model_task_config.lpmm_qa)
asyncio.run(import_data())
elif choice == "5":
await import_from_specific_file()
import_from_specific_file()
elif choice == "6":
clear_cache()
elif choice == "0":
@@ -385,4 +376,4 @@ async def async_main():
if __name__ == "__main__":
asyncio.run(async_main())
main()

176
src/chat/knowledge/embedding_store.py
View File

@@ -124,60 +124,124 @@ class EmbeddingStore:
self.faiss_index = None
self.idx2hash = None
@staticmethod
def _get_embedding(s: str) -> list[float]:
"""获取字符串的嵌入向量，使用完全同步的方式避免事件循环问题"""
# 创建新的事件循环并在完成后立即关闭
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
try:
# 创建新的LLMRequest实例
from src.config.config import model_config
from src.llm_models.utils_model import LLMRequest
llm = LLMRequest(model_set=model_config.model_task_config.embedding, request_type="embedding")
# 使用新的事件循环运行异步方法
embedding, _ = loop.run_until_complete(llm.get_embedding(s))
if embedding and len(embedding) > 0:
return embedding
else:
logger.error(f"获取嵌入失败: {s}")
return []
except Exception as e:
logger.error(f"获取嵌入时发生异常: {s}, 错误: {e}")
return []
finally:
# 确保事件循环被正确关闭
try:
loop.close()
except Exception:
...
@staticmethod
def _get_embeddings_batch_threaded(
strs: list[str],
main_loop: asyncio.AbstractEventLoop,
chunk_size: int = 10,
max_workers: int = 10,
progress_callback=None,
strs: list[str], chunk_size: int = 10, max_workers: int = 10, progress_callback=None
) -> list[tuple[str, list[float]]]:
"""使用多线程批量获取嵌入向量, 并通过 run_coroutine_threadsafe 在主事件循环中运行异步任务"""
"""使用多线程批量获取嵌入向量
Args:
strs: 要获取嵌入的字符串列表
chunk_size: 每个线程处理的数据块大小
max_workers: 最大线程数
progress_callback: 进度回调函数，接收一个参数表示完成的数量
Returns:
包含(原始字符串, 嵌入向量)的元组列表，保持与输入顺序一致
"""
if not strs:
return []
# 导入必要的模块
from src.config.config import model_config
from src.llm_models.utils_model import LLMRequest
# 在主线程（即主事件循环所在的线程）中创建LLMRequest实例
# 这样可以确保它绑定到正确的事件循环
llm = LLMRequest(model_set=model_config.model_task_config.embedding, request_type="embedding")
# 分块
chunks = [(i, strs[i : i + chunk_size]) for i in range(0, len(strs), chunk_size)]
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 = []
for i, s in enumerate(chunk_strs):
embedding = []
try:
# 将异步的 get_embedding 调用提交到主事件循环
future = asyncio.run_coroutine_threadsafe(llm.get_embedding(s), main_loop)
# 同步等待结果，延长超时时间
embedding_result, _ = future.result(timeout=60)
# 为每个线程创建独立的LLMRequest实例
from src.config.config import model_config
from src.llm_models.utils_model import LLMRequest
if embedding_result and len(embedding_result) > 0:
embedding = embedding_result
else:
logger.error(f"获取嵌入失败（返回为空）: {s}")
try:
# 创建线程专用的LLM实例
llm = LLMRequest(model_set=model_config.model_task_config.embedding, request_type="embedding")
except Exception as e:
logger.error(f"在线程中获取嵌入时发生异常: {s}, 错误: {type(e).__name__}: {e}")
finally:
chunk_results.append((start_idx + i, s, embedding))
for i, s in enumerate(chunk_strs):
try:
# 在线程中创建独立的事件循环
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
try:
embedding = loop.run_until_complete(llm.get_embedding(s))
finally:
loop.close()
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()
@@ -185,14 +249,22 @@ class EmbeddingStore:
results[idx] = (s, embedding)
except Exception as e:
chunk = future_to_chunk[future]
logger.error(f"处理数据块时发生严重异常: {chunk}, 错误: {e}")
logger.error(f"处理数据块时发生异常: {chunk}, 错误: {e}")
# 为失败的块添加空结果
start_idx, chunk_strs = chunk
for i, s_item in enumerate(chunk_strs):
if (start_idx + i) not in results:
results[start_idx + i] = (s_item, [])
for i, s in enumerate(chunk_strs):
results[start_idx + i] = (s, [])
# 按原始顺序返回结果
return [results.get(i, (strs[i], [])) for i in range(len(strs))]
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
@staticmethod
def get_test_file_path():
@@ -202,17 +274,9 @@ class EmbeddingStore:
"""保存测试字符串的嵌入到本地（使用多线程优化）"""
logger.info("开始保存测试字符串的嵌入向量...")
# 获取当前正在运行的事件循环
try:
main_loop = asyncio.get_running_loop()
except RuntimeError:
logger.error("无法获取正在运行的事件循环。请确保在异步上下文中调用此方法。")
return
# 使用多线程批量获取测试字符串的嵌入
embedding_results = self._get_embeddings_batch_threaded(
EMBEDDING_TEST_STRINGS,
main_loop,
chunk_size=min(self.chunk_size, len(EMBEDDING_TEST_STRINGS)),
max_workers=min(self.max_workers, len(EMBEDDING_TEST_STRINGS)),
)
@@ -224,6 +288,8 @@ class EmbeddingStore:
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:
f.write(orjson.dumps(test_vectors, option=orjson.OPT_INDENT_2).decode("utf-8"))
@@ -255,17 +321,9 @@ class EmbeddingStore:
logger.info("开始检验嵌入模型一致性...")
# 获取当前正在运行的事件循环
try:
main_loop = asyncio.get_running_loop()
except RuntimeError:
logger.error("无法获取正在运行的事件循环。请确保在异步上下文中调用此方法。")
return False
# 使用多线程批量获取当前模型的嵌入
embedding_results = self._get_embeddings_batch_threaded(
EMBEDDING_TEST_STRINGS,
main_loop,
chunk_size=min(self.chunk_size, len(EMBEDDING_TEST_STRINGS)),
max_workers=min(self.max_workers, len(EMBEDDING_TEST_STRINGS)),
)
@@ -325,20 +383,11 @@ class EmbeddingStore:
progress.update(task, advance=already_processed)
if new_strs:
try:
main_loop = asyncio.get_running_loop()
except RuntimeError:
logger.error("无法获取正在运行的事件循环。请确保在异步上下文中调用此方法。")
# 更新进度条以反映未处理的项目
progress.update(task, advance=len(new_strs))
return
# 使用实例配置的参数，智能调整分块和线程数
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,
self.chunk_size, len(new_strs) // self.max_workers if self.max_workers > 0 else self.chunk_size
),
)
optimal_max_workers = min(
@@ -355,13 +404,12 @@ class EmbeddingStore:
# 批量获取嵌入，并实时更新进度
embedding_results = self._get_embeddings_batch_threaded(
new_strs,
main_loop,
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: # 只有成功获取到嵌入才存入