From 50ae55bb18def5abdaf7f7306a5ef38d9188fb48 Mon Sep 17 00:00:00 2001 From: changyun <45411281+cyun9601@users.noreply.github.com> Date: Wed, 8 May 2024 14:25:04 +0900 Subject: [PATCH 1/3] Update prompts.py Fixed an issue where the LLM model and custom prompts could not distinguish between main and other tasks. Signed-off-by: changyun <45411281+cyun9601@users.noreply.github.com> --- nemoguardrails/llm/prompts.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/nemoguardrails/llm/prompts.py b/nemoguardrails/llm/prompts.py index 73e5dba8f..e8c3b7aae 100644 --- a/nemoguardrails/llm/prompts.py +++ b/nemoguardrails/llm/prompts.py @@ -122,10 +122,14 @@ def get_prompt(config: RailsConfig, task: Union[str, Task]) -> TaskPrompt: task_model = "unknown" if config.models: - task_model = config.models[0].engine - if config.models[0].model: - task_model += "/" + config.models[0].model - + _models = [model for model in config.models if model.type == task.value] + if not _models: + _models = [model for model in config.models if model.type == "main"] + + task_model = _models[0].engine + if _models[0].model: + task_model += "/" + _models[0].model + task_prompting_mode = "standard" if config.prompting_mode: # if exists in config, overwrite, else, default to "standard" From c8c3e294a18ce847840451a95245be9541d322cb Mon Sep 17 00:00:00 2001 From: changyun <45411281+cyun9601@users.noreply.github.com> Date: Wed, 8 May 2024 15:08:22 +0900 Subject: [PATCH 2/3] Fix select model type and custom prompts task.py Fixed an issue where the LLM model and custom prompts could not distinguish between main and other tasks. Signed-off-by: changyun <45411281+cyun9601@users.noreply.github.com> --- nemoguardrails/llm/prompts.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/nemoguardrails/llm/prompts.py b/nemoguardrails/llm/prompts.py index e8c3b7aae..b0b88e09b 100644 --- a/nemoguardrails/llm/prompts.py +++ b/nemoguardrails/llm/prompts.py @@ -114,8 +114,6 @@ def _get_prompt( def get_prompt(config: RailsConfig, task: Union[str, Task]) -> TaskPrompt: """Return the prompt for the given task.""" - # Currently, we use the main model for all tasks - # TODO: add support to use different models for different tasks # Fetch current task parameters like name, models to use, and the prompting mode task_name = str(task.value) if isinstance(task, Task) else task From c519a1d43439ae60baa9bba433dcbf172a3c55f1 Mon Sep 17 00:00:00 2001 From: changyun <45411281+cyun9601@users.noreply.github.com> Date: Fri, 10 May 2024 17:21:06 +0900 Subject: [PATCH 3/3] Update prompts.py Signed-off-by: changyun <45411281+cyun9601@users.noreply.github.com> --- nemoguardrails/llm/prompts.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/nemoguardrails/llm/prompts.py b/nemoguardrails/llm/prompts.py index b0b88e09b..2c64f70da 100644 --- a/nemoguardrails/llm/prompts.py +++ b/nemoguardrails/llm/prompts.py @@ -120,7 +120,7 @@ def get_prompt(config: RailsConfig, task: Union[str, Task]) -> TaskPrompt: task_model = "unknown" if config.models: - _models = [model for model in config.models if model.type == task.value] + _models = [model for model in config.models if model.type == task_name] if not _models: _models = [model for model in config.models if model.type == "main"]