feat: add rerank
This commit is contained in:
parent
a00847ff82
commit
1a83148928
|
|
@ -1,4 +1,4 @@
|
||||||
#!/usr/bin/env bash
|
#!/usr/bin/env bash
|
||||||
export $(cat .env | xargs)
|
export $(cat .env | xargs)
|
||||||
pytest -s -W ignore::DeprecationWarning -k test_search src/tests/test_nodes.py
|
pytest -s -W ignore::DeprecationWarning src/tests/test_nodes.py::test_rerank
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -14,6 +14,9 @@ class Config(TypedDict):
|
||||||
llm_api_key: str
|
llm_api_key: str
|
||||||
llm_api_host: str
|
llm_api_host: str
|
||||||
llm_model: str
|
llm_model: str
|
||||||
|
rerank_api_key: str
|
||||||
|
rerank_api_host: str
|
||||||
|
rerank_model: str
|
||||||
embedding_api_key: str
|
embedding_api_key: str
|
||||||
embedding_api_host: str
|
embedding_api_host: str
|
||||||
embedding_model: str
|
embedding_model: str
|
||||||
|
|
@ -33,6 +36,9 @@ def _read_config() -> Config:
|
||||||
"llm_api_host": os.getenv("LLM_API_HOST"),
|
"llm_api_host": os.getenv("LLM_API_HOST"),
|
||||||
"llm_api_key": os.getenv("LLM_API_KEY"),
|
"llm_api_key": os.getenv("LLM_API_KEY"),
|
||||||
"llm_model": os.getenv("LLM_MODEL"),
|
"llm_model": os.getenv("LLM_MODEL"),
|
||||||
|
"rerank_api_host": os.getenv("RERANK_API_HOST"),
|
||||||
|
"rerank_api_key": os.getenv("RERANK_API_KEY"),
|
||||||
|
"rerank_model": os.getenv("RERANK_MODEL"),
|
||||||
"embedding_api_host": os.getenv("EMBEDDING_API_HOST"),
|
"embedding_api_host": os.getenv("EMBEDDING_API_HOST"),
|
||||||
"embedding_api_key": os.getenv("EMBEDDING_API_KEY"),
|
"embedding_api_key": os.getenv("EMBEDDING_API_KEY"),
|
||||||
"embedding_model": os.getenv("EMBEDDING_MODEL"),
|
"embedding_model": os.getenv("EMBEDDING_MODEL"),
|
||||||
|
|
|
||||||
|
|
@ -1,4 +1,3 @@
|
||||||
from ast import List
|
|
||||||
import httpx
|
import httpx
|
||||||
import json
|
import json
|
||||||
from src.pipeline.config import config
|
from src.pipeline.config import config
|
||||||
|
|
@ -46,6 +45,23 @@ class AsyncLLm:
|
||||||
"Authorization": f"Bearer {self.embedding_api_key}",
|
"Authorization": f"Bearer {self.embedding_api_key}",
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
|
# rerank
|
||||||
|
self.rerank_api = config["rerank_api_host"].rstrip("/") + "/score"
|
||||||
|
self.rerank_model = config["rerank_model"]
|
||||||
|
self.rerank_api_key = config["rerank_api_key"]
|
||||||
|
self.rerank_client = httpx.AsyncClient(
|
||||||
|
http2=False,
|
||||||
|
trust_env=False,
|
||||||
|
timeout=httpx.Timeout(timeout),
|
||||||
|
limits=httpx.Limits(
|
||||||
|
max_connections=max_connections,
|
||||||
|
max_keepalive_connections=max_keepalive,
|
||||||
|
),
|
||||||
|
headers={
|
||||||
|
"Content-Type": "application/json",
|
||||||
|
"Authorization": f"Bearer {self.rerank_api_key}",
|
||||||
|
},
|
||||||
|
)
|
||||||
|
|
||||||
async def embedding(self, text: str) -> list[float]:
|
async def embedding(self, text: str) -> list[float]:
|
||||||
try:
|
try:
|
||||||
|
|
@ -66,7 +82,48 @@ class AsyncLLm:
|
||||||
|
|
||||||
return []
|
return []
|
||||||
|
|
||||||
async def rerank(self, query: str, documents: List[dict]) -> List[dict]:
|
async def rerank(self, query: str, documents: list[dict]) -> list[dict]:
|
||||||
|
|
||||||
|
logger.debug("\n\n default:\n\n" + "\n".join([str({"content": x["content"][:10], "score": x["score"]}) for x in documents]))
|
||||||
|
|
||||||
|
if not documents:
|
||||||
|
return []
|
||||||
|
|
||||||
|
text_1 = []
|
||||||
|
text_2 = []
|
||||||
|
|
||||||
|
for doc in documents:
|
||||||
|
text_1.append(query)
|
||||||
|
text_2.append(doc.get("content", "")[:1024])
|
||||||
|
|
||||||
|
# Qwen3 专用
|
||||||
|
payload = {
|
||||||
|
"model": self.rerank_model,
|
||||||
|
"text_1": text_1,
|
||||||
|
"text_2": text_2,
|
||||||
|
}
|
||||||
|
|
||||||
|
try:
|
||||||
|
resp = await self.rerank_client.post(self.rerank_api, json=payload)
|
||||||
|
resp.raise_for_status()
|
||||||
|
result = resp.json()
|
||||||
|
scores = [item["score"] for item in result["data"]]
|
||||||
|
|
||||||
|
reranked = []
|
||||||
|
for doc, score in zip(documents, scores):
|
||||||
|
d = dict(doc)
|
||||||
|
d["rerank_score"] = float(score)
|
||||||
|
reranked.append(d)
|
||||||
|
reranked.sort(key=lambda x: x["rerank_score"], reverse=True)
|
||||||
|
|
||||||
|
logger.debug("\n\n rerank: \n\n" + "\n".join([str({"content": x["content"][:10], "score": x["score"], "rerank_score": x["rerank_score"]}) for x in reranked]))
|
||||||
|
|
||||||
|
return reranked
|
||||||
|
except httpx.HTTPStatusError as e:
|
||||||
|
logger.error(e)
|
||||||
|
logger.error(f"Rerank HTTP error: {e.response.text}")
|
||||||
|
except Exception as e:
|
||||||
|
logger.exception("Rerank request failed")
|
||||||
return []
|
return []
|
||||||
|
|
||||||
async def chat(
|
async def chat(
|
||||||
|
|
|
||||||
|
|
@ -143,6 +143,7 @@ class SearchFromESNode(AsyncNode):
|
||||||
return await es.client.hybrid_search_es(**prep_res)
|
return await es.client.hybrid_search_es(**prep_res)
|
||||||
|
|
||||||
async def post_async(self, shared, prep_res, exec_res):
|
async def post_async(self, shared, prep_res, exec_res):
|
||||||
|
logger.debug(f"SearchFromESNode: length:{len(exec_res)}")
|
||||||
shared["results"] = exec_res
|
shared["results"] = exec_res
|
||||||
return "default"
|
return "default"
|
||||||
|
|
||||||
|
|
@ -154,12 +155,16 @@ class RerankNode(AsyncNode):
|
||||||
|
|
||||||
async def prep_async(self, shared):
|
async def prep_async(self, shared):
|
||||||
# 准备要重排的数据
|
# 准备要重排的数据
|
||||||
return {"query": shared["text"], "results": shared.get("results", [])}
|
return {
|
||||||
|
"query": shared["text"],
|
||||||
|
"results": shared.get("results", [])[: shared["rerank_top_k"]],
|
||||||
|
}
|
||||||
|
|
||||||
async def exec_async(self, prep_res):
|
async def exec_async(self, prep_res):
|
||||||
query = prep_res["query"]
|
query = prep_res["query"]
|
||||||
results = prep_res["results"]
|
results = prep_res["results"]
|
||||||
|
logger.debug(f"query: {query}")
|
||||||
|
logger.debug(f"results: {results}")
|
||||||
if not results:
|
if not results:
|
||||||
return []
|
return []
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -1,3 +1,4 @@
|
||||||
|
import asyncio
|
||||||
import pytest
|
import pytest
|
||||||
import json
|
import json
|
||||||
from src.pipeline.core.pocket_flow import AsyncFlow
|
from src.pipeline.core.pocket_flow import AsyncFlow
|
||||||
|
|
@ -78,13 +79,17 @@ async def test_rerank():
|
||||||
shared = {
|
shared = {
|
||||||
"text": "哪里盛产矿石",
|
"text": "哪里盛产矿石",
|
||||||
"index": "test_kb",
|
"index": "test_kb",
|
||||||
"top_k": 10,
|
"top_k": 50,
|
||||||
|
"rerank_top_k": 5,
|
||||||
|
"top_n": 3,
|
||||||
"results": [], # [{es_id, doc_id, title, type, created_at, score, content}]
|
"results": [], # [{es_id, doc_id, title, type, created_at, score, content}]
|
||||||
}
|
}
|
||||||
|
|
||||||
embeddingNode = nodes.EmbeddingNode()
|
embeddingNode = nodes.EmbeddingNode()
|
||||||
searchNode = nodes.SearchFromESNode()
|
searchNode = nodes.SearchFromESNode()
|
||||||
embeddingNode >> searchNode
|
rerankNode = nodes.RerankNode()
|
||||||
|
embeddingNode >> searchNode >> rerankNode
|
||||||
|
# embeddingNode >> searchNode
|
||||||
flow = AsyncFlow(embeddingNode)
|
flow = AsyncFlow(embeddingNode)
|
||||||
|
|
||||||
await flow.run_async(shared)
|
await flow.run_async(shared)
|
||||||
|
|
@ -94,7 +99,7 @@ async def test_rerank():
|
||||||
res = await llm.client.chat(
|
res = await llm.client.chat(
|
||||||
messages=[
|
messages=[
|
||||||
{"role": "system", "content": utils.rag_system_prompt()},
|
{"role": "system", "content": utils.rag_system_prompt()},
|
||||||
{"role": "user", "content": utils.rag_user_prompt(shared["text"], shared["results"])},
|
{"role": "user", "content": utils.rag_user_prompt(shared["text"], shared["results"][:shared["top_n"]])},
|
||||||
]
|
]
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue