gavindiaz/NexQuant
1
0
Fork 0
mirror of https://github.com/NicolasBohn/NexQuant.git synced 2026-08-01 17:37:43 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: using different chat model in different part (#822)

Browse Source 
* using model in the chat_model_map in one tag

* add replace timer to DS loop

* fix CI

* fix CI

* add more custom config in chat_model_map

* fix CI

* fix CI

* fix CI

---------

Co-authored-by: Xu Yang <xuyang1@microsoft.com>
This commit is contained in:
XianBW
2025-04-24 18:22:51 +08:00
committed by GitHub
parent 44ccee864d
commit f909b1b6bc
5 changed files with 51 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rdagent/app/data_science/loop.py
+8 -3
View File
@@ -221,9 +221,13 @@ class DataScienceRDLoop(RDLoop):
@classmethod
def load(
cls, path: Union[str, Path], output_path: Optional[Union[str, Path]] = None, do_truncate: bool = False
cls,
path: Union[str, Path],
output_path: Optional[Union[str, Path]] = None,
do_truncate: bool = False,
replace_timer: bool = True,
) -> "LoopBase":
session = super().load(path, output_path, do_truncate)
session = super().load(path, output_path, do_truncate, replace_timer)
if (
DS_RD_SETTING.enable_knowledge_base
and DS_RD_SETTING.knowledge_base_version == "v1"
@@ -255,6 +259,7 @@ def main(
competition="bms-molecular-translation",
do_truncate=True,
timeout=None,
replace_timer=True,
):
"""
@@ -297,7 +302,7 @@ def main(
if path is None:
kaggle_loop = DataScienceRDLoop(DS_RD_SETTING)
else:
kaggle_loop = DataScienceRDLoop.load(path, output_path, do_truncate)
kaggle_loop = DataScienceRDLoop.load(path, output_path, do_truncate, replace_timer)
kaggle_loop.run(step_n=step_n, loop_n=loop_n, all_duration=timeout)
rdagent/oai/backend/deprec.py
+13 -13
View File
@@ -152,7 +152,7 @@ class DeprecBackend(APIBackend):
self.gcr_endpoint_max_token = LLM_SETTINGS.gcr_endpoint_max_token
if not os.environ.get("PYTHONHTTPSVERIFY", "") and hasattr(ssl, "_create_unverified_context"):
ssl._create_default_https_context = ssl._create_unverified_context # noqa: SLF001
self.chat_model_map = json.loads(LLM_SETTINGS.chat_model_map)
self.chat_model_map = LLM_SETTINGS.chat_model_map
self.chat_model = LLM_SETTINGS.chat_model
self.encoder = None
elif LLM_SETTINGS.chat_use_azure_deepseek:
@@ -160,7 +160,7 @@ class DeprecBackend(APIBackend):
endpoint=LLM_SETTINGS.chat_azure_deepseek_endpoint,
credential=AzureKeyCredential(LLM_SETTINGS.chat_azure_deepseek_key),
)
self.chat_model_map = json.loads(LLM_SETTINGS.chat_model_map)
self.chat_model_map = LLM_SETTINGS.chat_model_map
self.encoder = None
self.chat_model = "deepseek-R1"
self.chat_stream = LLM_SETTINGS.chat_stream
@@ -181,7 +181,7 @@ class DeprecBackend(APIBackend):
)
self.chat_model = LLM_SETTINGS.chat_model
self.chat_model_map = json.loads(LLM_SETTINGS.chat_model_map)
self.chat_model_map = LLM_SETTINGS.chat_model_map
self.encoder = self._get_encoder()
self.chat_openai_base_url = LLM_SETTINGS.chat_openai_base_url
self.embedding_openai_base_url = LLM_SETTINGS.embedding_openai_base_url
@@ -303,19 +303,20 @@ class DeprecBackend(APIBackend):
logger.info(self._build_log_messages(messages), tag="llm_messages")
# TODO: fail to use loguru adaptor due to stream response
model = LLM_SETTINGS.chat_model
temperature = LLM_SETTINGS.chat_temperature
max_tokens = LLM_SETTINGS.chat_max_tokens
frequency_penalty = LLM_SETTINGS.chat_frequency_penalty
presence_penalty = LLM_SETTINGS.chat_presence_penalty
# Use index 4 to skip the current function and intermediate calls,
# and get the locals of the caller's frame.
caller_locals = inspect.stack()[4].frame.f_locals
if "self" in caller_locals:
tag = caller_locals["self"].__class__.__name__
else:
tag = inspect.stack()[4].function
model = self.chat_model_map.get(tag, self.chat_model)
if self.chat_model_map:
for t, mc in self.chat_model_map.items():
if t in logger._tag:
model = mc.get("model", model)
temperature = float(mc.get("temperature", temperature))
if "max_tokens" in mc:
max_tokens = int(mc["max_tokens"])
break
finish_reason = None
if self.use_llama2:
@@ -394,7 +395,7 @@ class DeprecBackend(APIBackend):
logger.info(f"{LogColors.CYAN}Think:{think_part}{LogColors.END}", tag="llm_messages")
logger.info(f"{LogColors.CYAN}Response:{resp}{LogColors.END}", tag="llm_messages")
else:
call_kwargs = dict(
call_kwargs: dict[str, Any] = dict(
model=model,
messages=messages,
max_tokens=max_tokens,
@@ -443,7 +444,6 @@ class DeprecBackend(APIBackend):
logger.info(
json.dumps(
{
"tag": tag,
"total_tokens": response.usage.total_tokens,
"prompt_tokens": response.usage.prompt_tokens,
"completion_tokens": response.usage.completion_tokens,
rdagent/oai/backend/litellm.py
+23 -9
View File
@@ -1,4 +1,4 @@
from typing import Any
from typing import Any, Literal, cast
from litellm import (
completion,
@@ -78,18 +78,32 @@ class LiteLLMAPIBackend(APIBackend):
logger.info(self._build_log_messages(messages), tag="llm_messages")
# Call LiteLLM completion
model = LITELLM_SETTINGS.chat_model
temperature = LITELLM_SETTINGS.chat_temperature
max_tokens = LITELLM_SETTINGS.chat_max_tokens
reasoning_effort = LITELLM_SETTINGS.reasoning_effort
if LITELLM_SETTINGS.chat_model_map:
for t, mc in LITELLM_SETTINGS.chat_model_map.items():
if t in logger._tag:
model = mc["model"]
if "temperature" in mc:
temperature = float(mc["temperature"])
if "max_tokens" in mc:
max_tokens = int(mc["max_tokens"])
if "reasoning_effort" in mc and mc["reasoning_effort"] in ["low", "medium", "high"]:
reasoning_effort = cast(Literal["low", "medium", "high"], mc["reasoning_effort"])
break
response = completion(
model=LITELLM_SETTINGS.chat_model,
model=model,
messages=messages,
stream=LITELLM_SETTINGS.chat_stream,
temperature=LITELLM_SETTINGS.chat_temperature,
max_tokens=LITELLM_SETTINGS.chat_max_tokens,
reasoning_effort=LITELLM_SETTINGS.reasoning_effort,
temperature=temperature,
max_tokens=max_tokens,
reasoning_effort=reasoning_effort,
**kwargs,
)
logger.info(
f"{LogColors.GREEN}Using chat model{LogColors.END} {LITELLM_SETTINGS.chat_model}", tag="llm_messages"
)
logger.info(f"{LogColors.GREEN}Using chat model{LogColors.END} {model}", tag="llm_messages")
if LITELLM_SETTINGS.chat_stream:
logger.info(f"{LogColors.BLUE}assistant:{LogColors.END}", tag="llm_messages")
@@ -117,7 +131,7 @@ class LiteLLMAPIBackend(APIBackend):
logger.info(f"{LogColors.BLUE}assistant:{LogColors.END} {finish_reason_str}\n{content}", tag="llm_messages")
global ACC_COST
cost = completion_cost(model=LITELLM_SETTINGS.chat_model, messages=messages, completion=content)
cost = completion_cost(model=model, messages=messages, completion=content)
ACC_COST += cost
logger.info(
f"Current Cost: ${float(cost):.10f}; Accumulated Cost: ${float(ACC_COST):.10f}; {finish_reason=}",
rdagent/oai/llm_conf.py
+1 -1
View File
@@ -114,7 +114,7 @@ class LLMSettings(ExtendedBaseSettings):
chat_azure_deepseek_endpoint: str = ""
chat_azure_deepseek_key: str = ""
chat_model_map: str = "{}"
chat_model_map: dict[str, dict[str, str]] = {}
LLM_SETTINGS = LLMSettings()
rdagent/utils/workflow.py
+6 -2
View File
@@ -203,7 +203,11 @@ class LoopBase:
@classmethod
def load(
cls, path: Union[str, Path], output_path: Optional[Union[str, Path]] = None, do_truncate: bool = False
cls,
path: Union[str, Path],
output_path: Optional[Union[str, Path]] = None,
do_truncate: bool = False,
replace_timer: bool = True,
) -> "LoopBase":
path = Path(path)
with path.open("rb") as f:
@@ -223,7 +227,7 @@ class LoopBase:
max_loop = max(session.loop_trace.keys())
logger.storage.truncate(time=session.loop_trace[max_loop][-1].end)
if session.timer.started:
if session.timer.started and replace_timer:
RD_Agent_TIMER_wrapper.replace_timer(session.timer)
RD_Agent_TIMER_wrapper.timer.restart_by_remain_time()
return session