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