feat: 调整结构
This commit is contained in:
parent
148474de9b
commit
5b101f26c6
|
@ -173,13 +173,14 @@ class ModelManager:
|
||||||
|
|
||||||
|
|
||||||
class BaseRAG(ABC):
|
class BaseRAG(ABC):
|
||||||
|
|
||||||
def __init__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
vector_store_name: str = "default",
|
vector_store_name: str = "default",
|
||||||
embedding_config: Optional[Dict] = None,
|
|
||||||
retriever_top_k: int = 3,
|
retriever_top_k: int = 3,
|
||||||
llm: Optional[BaseLLM] = None,
|
|
||||||
persist_directory: str = "./chroma_db",
|
persist_directory: str = "./chroma_db",
|
||||||
|
llm: Optional[BaseLLM] = None,
|
||||||
|
embedding_config: Optional[Dict] = None,
|
||||||
rerank_config: Optional[Dict] = None,
|
rerank_config: Optional[Dict] = None,
|
||||||
):
|
):
|
||||||
"""
|
"""
|
||||||
|
@ -271,7 +272,7 @@ class BaseRAG(ABC):
|
||||||
# 根据分数排序
|
# 根据分数排序
|
||||||
doc_scores = list(zip(documents, scores))
|
doc_scores = list(zip(documents, scores))
|
||||||
doc_scores.sort(key=lambda x: x[1], reverse=True)
|
doc_scores.sort(key=lambda x: x[1], reverse=True)
|
||||||
|
|
||||||
# 返回top_k个文档
|
# 返回top_k个文档
|
||||||
result = [doc for doc, score in doc_scores[:top_k]]
|
result = [doc for doc, score in doc_scores[:top_k]]
|
||||||
print(f"重排输出: {len(result)} 个文档")
|
print(f"重排输出: {len(result)} 个文档")
|
||||||
|
|
Loading…
Reference in New Issue