Task specific models

[manikyabard]

Are there any immediate plans to complete the TODO "add support to use different models for different tasks" in llm/prompts.py? If the solution is thought out, I can potentially help with this.

[drazvan]

Hi @manikyabard ! It would be great if you could help. The solution is not fully fleshed out. Let me follow up on this tomorrow.

[drazvan]

@manikyabard: Let's split this in two parts.

First, we want a simple support for specifying the task_model as a parameter. And if it isn't specified, it should use the existing logic to use the main one. You can make a first Pull Request for this change only.

Second, the LLMTaskManager should actually provide a more complete interface for "running LLM tasks", i.e., given a task name and parameters, and potentially a model name, to take care of rendering the prompt, making the LLM call, applying the output parser and ultimately returning the final result. This pattern is used in many places, e.g., https://github.com/NVIDIA/NeMo-Guardrails/blob/develop/nemoguardrails/actions/llm/generation.py#L297-L316.

So, adding LLMTaskManager.run_task_async(task_name: str, model_name: str, task_args: dict) will provide a cleaner interface. For this to work, the LLMTaskManager will need to have access to the initialized LLMs somehow. One solution is for the _init_llms (https://github.com/NVIDIA/NeMo-Guardrails/blob/develop/nemoguardrails/rails/llm/llmrails.py#L253) function to store a dict, pass it to the Runtime when it is initialized, which in turn will pass it to the LLMTaskManager instance when initialized.

If you want to take on the second part as well, I'm happy to provide more guidance if needed.

Thanks!

[manikyabard]

Just to make sure I'm understanding this correctly, it looks like the model type would have to be the action/task name if we want the correct llm to be passed in during the action execution (according to _process_start_action in flows/runtime.py and _init_llms in llmrails.py).

Although determining the task_model should be fairly straightforward based on the type value, would a (list) parameter like tasks in the models config with values being a list of the task names make more sense?

Basically something like

models:
  - type: model_type
    engine: azure
    model: gpt-4
    parameters:
      tasks: ['generate_bot_message', 'generate_value']
      openai_api_key: "..."
      openai_api_version: "..."
      azure_endpoint: "..."
      deployment_name: "..."

instead of doing

models:
  - type: generate_bot_message
    engine: azure
    model: gpt-4
    parameters:
      openai_api_key: "..."
      openai_api_version: "..."
      azure_endpoint: "..."
      deployment_name: "..."
  - type: generate_value
    engine: azure
    model: gpt-4
    parameters:
      openai_api_key: "..."
      openai_api_version: "..."
      azure_endpoint: "..."
      deployment_name: "..."

just to make the config more concise.

Also, is the second part more so to remove the llm parameter from the action execution calls and somehow directly make use of the llm task manager to determine these things?

[manikyabard]

@drazvan I think I have a better understanding of how the current nemoguardrails version supports leveraging different llms for different actions by looking at the custom rag example which uses the tabular_llm.

Right now, it seems like the missing functionality is being able to specify a particular model config to be used with a particular action without modifying the action parameters (this would be good to support for changing some models for internal actions like generate_bot_message which seems to work well with GPT-4).

My understanding is that this can be done using the task_model parameter. It looks like there are two approaches for this. We can either infer this value based on other parameters (above comment), or specify this directly somehow.

Which approach do you think is better design? If we want to specify task_model directly, where's the best place for doing so?

[manikyabard]

@drazvan Any thoughts on this? Is this something which is being worked on already now?

[drazvan]

@manikyabard : sorry for the delay in responding. No, this is not being worked on. Adding a model parameter in the YAML configuration is the way to go. It should support the format "ENGINE / MODEL_NAME". When this is set for a task, regardless of the main LLM, the model specified for the task will be used.

...
    - task: fact_checking
+     model: "openai/gpt-3.5-turbo-instruct"
      content: |-
      ...