107 lines
4.4 KiB
Python
107 lines
4.4 KiB
Python
# 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
|
||
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
|
||
)
|
||
|
||
|
||
# 模拟长字符串文档内容
|
||
long_text = """混凝土是一种广泛使用的建筑材料，由水泥、砂、石子和水混合而成。它具有高强度、耐久性和良好的可塑性，被广泛应用于建筑、桥梁、道路等土木工程领域。
|
||
|
||
混凝土的历史可以追溯到古罗马时期，当时人们使用火山灰、石灰和碎石混合制成类似混凝土的材料。现代混凝土技术始于19世纪，随着波特兰水泥的发明而得到快速发展。
|
||
|
||
混凝土的性能取决于其配合比，包括水灰比、砂率等参数。水灰比是影响混凝土强度的关键因素，较小的水灰比通常会产生更高强度的混凝土。
|
||
|
||
为了改善混凝土的性能，常常会添加各种外加剂，如减水剂、早强剂、缓凝剂等。此外，还可以使用纤维增强、聚合物改性等技术来提高混凝土的韧性和耐久性。
|
||
|
||
在施工过程中，混凝土需要适当的养护，以确保其强度正常发展。养护措施包括浇水、覆盖保湿、蒸汽养护等。
|
||
|
||
随着建筑技术的发展，高性能混凝土、自密实混凝土、再生骨料混凝土等新型混凝土不断涌现，为土木工程领域提供了更多的选择。"""
|
||
|
||
# 创建文档对象
|
||
docs = [Document(page_content=long_text, metadata={"source": "simulated_document"})]
|
||
print(f"文档数量：{len(docs)} 个")
|
||
|
||
# 切割文档
|
||
text_splitter = RecursiveCharacterTextSplitter(
|
||
chunk_size=200, chunk_overlap=100, add_start_index=True
|
||
)
|
||
all_splits = text_splitter.split_documents(docs)
|
||
|
||
print(f"切割后的文档块数量：{len(all_splits)}")
|
||
|
||
# 嵌入模型
|
||
embeddings = OpenAIEmbeddings(
|
||
model=EMBED_MODEL_NAME,
|
||
base_url=EMBED_BASE_URL,
|
||
api_key=SecretStr(EMBED_API_KEY) # 包装成 SecretStr 类型
|
||
)
|
||
|
||
# 向量存储
|
||
vector_store = InMemoryVectorStore(embeddings)
|
||
ids = vector_store.add_documents(documents=all_splits)
|
||
|
||
# 向量查询 - 获取更多结果用于重排
|
||
query = "混凝土"
|
||
results = vector_store.similarity_search(query, k=4) # 获取4个结果用于重排
|
||
|
||
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()
|
||
|
||
# 处理重排结果，提取relevance_score
|
||
if "results" in rerank_result:
|
||
for item in rerank_result["results"]:
|
||
doc_idx = item.get("index")
|
||
score = item.get("relevance_score", 0.0)
|
||
if 0 <= doc_idx < len(results):
|
||
reranked_docs_with_scores.append((results[doc_idx], score))
|
||
else:
|
||
print("警告: 无法识别重排API响应格式")
|
||
reranked_docs_with_scores = [(doc, 0.0) for doc in results]
|
||
|
||
print("重排后结果数量：", len(reranked_docs_with_scores))
|
||
except Exception as e:
|
||
print(f"重排模型调用失败: {e}")
|
||
print("将使用原始搜索结果")
|
||
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("---")
|