github/DB-GPT
1
0
Fork 0
mirror of https://github.com/csunny/DB-GPT.git synced 2025-09-13 21:21:08 +00:00
Code Issues Packages Projects Releases Wiki Activity

chore(model): Modify siliconflow name (#2194)

Browse Source
This commit is contained in:
Fangyin Cheng
2024-12-12 16:47:14 +08:00
committed by GitHub
parent 4da1809b31
commit dab55493f5
10 changed files with 48 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
dbgpt/rag/assembler/embedding.py
View File

@@ -1,4 +1,5 @@
"""Embedding Assembler."""
from concurrent.futures import ThreadPoolExecutor
from typing import Any, List, Optional
@@ -130,7 +131,11 @@ class EmbeddingAssembler(BaseAssembler):
Returns:
List[str]: List of chunk ids.
"""
return self._index_store.load_document(self._chunks)
max_chunks_once_load = kwargs.get("max_chunks_once_load", 10)
max_threads = kwargs.get("max_threads", 1)
return self._index_store.load_document_with_limit(
self._chunks, max_chunks_once_load, max_threads
)
async def apersist(self, **kwargs) -> List[str]:
"""Persist chunks into store.
@@ -139,7 +144,11 @@ class EmbeddingAssembler(BaseAssembler):
List[str]: List of chunk ids.
"""
# persist chunks into vector store
return await self._index_store.aload_document(self._chunks)
max_chunks_once_load = kwargs.get("max_chunks_once_load", 10)
max_threads = kwargs.get("max_threads", 1)
return await self._index_store.aload_document_with_limit(
self._chunks, max_chunks_once_load, max_threads
)
def _extract_info(self, chunks) -> List[Chunk]:
"""Extract info from chunks."""

4
dbgpt/rag/embedding/rerank.py
View File

@@ -175,10 +175,10 @@ class SiliconFlowRerankEmbeddings(OpenAPIRerankEmbeddings):
"""Initialize the SiliconFlowRerankEmbeddings."""
# If the API key is not provided, try to get it from the environment
if "api_key" not in kwargs:
kwargs["api_key"] = os.getenv("SILICON_FLOW_API_KEY")
kwargs["api_key"] = os.getenv("SILICONFLOW_API_KEY")
if "api_url" not in kwargs:
env_api_url = os.getenv("SILICON_FLOW_API_BASE")
env_api_url = os.getenv("SILICONFLOW_API_BASE")
if env_api_url:
env_api_url = env_api_url.rstrip("/")
kwargs["api_url"] = env_api_url + "/rerank"