From 17e52bc315ed0e352cd10cf2d4725e33f83f53f8 Mon Sep 17 00:00:00 2001 From: HuangHai <10402852@qq.com> Date: Tue, 19 Aug 2025 08:00:57 +0800 Subject: [PATCH] 'commit' --- dsSchoolBuddy/ElasticSearch/T2_BgeM3.py | 68 +++++++++++++++++++++---- 1 file changed, 57 insertions(+), 11 deletions(-) diff --git a/dsSchoolBuddy/ElasticSearch/T2_BgeM3.py b/dsSchoolBuddy/ElasticSearch/T2_BgeM3.py index 44799a92..1e9937c7 100644 --- a/dsSchoolBuddy/ElasticSearch/T2_BgeM3.py +++ b/dsSchoolBuddy/ElasticSearch/T2_BgeM3.py @@ -1,10 +1,15 @@ -# pip install pydantic +# pip install pydantic requests from langchain_core.documents import Document from langchain_core.vectorstores import InMemoryVectorStore from langchain_openai import OpenAIEmbeddings from langchain_text_splitters import RecursiveCharacterTextSplitter from pydantic import SecretStr # 导入 SecretStr -from Config.Config import EMBED_MODEL_NAME, EMBED_BASE_URL, EMBED_API_KEY +import requests +import json +from Config.Config import ( + EMBED_MODEL_NAME, EMBED_BASE_URL, EMBED_API_KEY, + RERANK_MODEL, RERANK_BASE_URL, RERANK_BINDING_API_KEY +) # 模拟长字符串文档内容 @@ -43,15 +48,56 @@ embeddings = OpenAIEmbeddings( vector_store = InMemoryVectorStore(embeddings) ids = vector_store.add_documents(documents=all_splits) -# 向量查询 -results = vector_store.similarity_search( - "混凝土", k=2 -) +# 向量查询 - 获取更多结果用于重排 +query = "混凝土" +results = vector_store.similarity_search(query, k=4) # 获取4个结果用于重排 -# 打印所有查询结果 -print("查询结果数量:", len(results)) -print("查询结果:") -for i, result in enumerate(results): - print(f"结果 {i+1}:") +print("向量搜索结果数量:", len(results)) + +# 存储重排后的文档和分数 +reranked_docs_with_scores = [] + +# 调用重排模型 +if len(results) > 1: + # 准备重排请求数据 + rerank_data = { + "model": RERANK_MODEL, + "query": query, + "documents": [doc.page_content for doc in results], + "top_n": len(results) + } + + # 调用SiliconFlow API进行重排 + headers = { + "Content-Type": "application/json", + "Authorization": f"Bearer {RERANK_BINDING_API_KEY}" + } + + try: + response = requests.post(RERANK_BASE_URL, headers=headers, data=json.dumps(rerank_data)) + response.raise_for_status() # 检查请求是否成功 + rerank_result = response.json() + + # 处理重排结果,保留分数 + for item in rerank_result.get("results", []): + doc_idx = item.get("index") + score = item.get("score", 0.0) # 获取可信度分数 + if 0 <= doc_idx < len(results): + reranked_docs_with_scores.append((results[doc_idx], score)) + + print("重排后结果数量:", len(reranked_docs_with_scores)) + except Exception as e: + print(f"重排模型调用失败: {e}") + print("将使用原始搜索结果") + # 使用原始结果,分数设为0.0 + reranked_docs_with_scores = [(doc, 0.0) for doc in results] +else: + # 只有一个结果,无需重排 + reranked_docs_with_scores = [(doc, 1.0) for doc in results] # 单个结果可信度设为1.0 + +# 打印所有查询结果及其可信度 +print("最终查询结果:") +for i, (result, score) in enumerate(reranked_docs_with_scores): + print(f"结果 {i+1} (可信度: {score:.4f}):") print(result.page_content) print("---")