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

feat(model): Support deploy rerank model (#1522)

Browse Source
This commit is contained in:
Fangyin Cheng
2024-05-16 14:50:16 +08:00
committed by GitHub
parent 559affe87d
commit 593e974405
29 changed files with 814 additions and 75 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
dbgpt/app/component_configs.py
View File

@@ -1,6 +1,7 @@
from __future__ import annotations
import logging
from typing import Optional
from dbgpt._private.config import Config
from dbgpt.app.base import WebServerParameters
@@ -18,9 +19,14 @@ def initialize_components(
system_app: SystemApp,
embedding_model_name: str,
embedding_model_path: str,
rerank_model_name: Optional[str] = None,
rerank_model_path: Optional[str] = None,
):
# Lazy import to avoid high time cost
from dbgpt.app.initialization.embedding_component import _initialize_embedding_model
from dbgpt.app.initialization.embedding_component import (
_initialize_embedding_model,
_initialize_rerank_model,
)
from dbgpt.app.initialization.scheduler import DefaultScheduler
from dbgpt.app.initialization.serve_initialization import register_serve_apps
from dbgpt.datasource.manages.connector_manager import ConnectorManager
@@ -45,6 +51,7 @@ def initialize_components(
_initialize_embedding_model(
param, system_app, embedding_model_name, embedding_model_path
)
_initialize_rerank_model(param, system_app, rerank_model_name, rerank_model_path)
_initialize_model_cache(system_app)
_initialize_awel(system_app, param)
# Initialize resource manager of agent
@@ -89,6 +96,11 @@ def _initialize_agent(system_app: SystemApp):
def _initialize_resource_manager(system_app: SystemApp):
from dbgpt.agent.expand.resources.dbgpt_tool import list_dbgpt_support_models
from dbgpt.agent.expand.resources.host_tool import (
get_current_host_cpu_status,
get_current_host_memory_status,
get_current_host_system_load,
)
from dbgpt.agent.expand.resources.search_tool import baidu_search
from dbgpt.agent.resource.base import ResourceType
from dbgpt.agent.resource.manage import get_resource_manager, initialize_resource
@@ -104,6 +116,10 @@ def _initialize_resource_manager(system_app: SystemApp):
# Register a search tool
rm.register_resource(resource_instance=baidu_search)
rm.register_resource(resource_instance=list_dbgpt_support_models)
# Register host tools
rm.register_resource(resource_instance=get_current_host_cpu_status)
rm.register_resource(resource_instance=get_current_host_memory_status)
rm.register_resource(resource_instance=get_current_host_system_load)
def _initialize_openapi(system_app: SystemApp):