Model parameters option to pass in model tuning, arbitrary parameters

[3coins]

Fixes #361
Fixes #444

Description

This PR brings model_parameters as an additional option to pass in arbitrary values to the provider class during initialization. This option is available in both Chat UI and Magics.

This is useful for passing parameters such as model tuning that affect the response generation by the model.
This is also an appropriate place to pass in custom attributes required by certain providers/models.

The accepted value is a dictionary, with top level keys as the model id (provider:model_id), and value
should be any arbitrary dictionary which is unpacked and passed as-is to the provider class.

Configuring as a startup option

In this sample, the bedrock provider will be created with the value for model_kwargs when ai21.j2-mid-v1 model is selected.

jupyter lab --AiExtension.model_parameters {"bedrock:ai21.j2-mid-v1":{"model_kwargs":{"maxTokens":200}}}

The above will result in the following LLM class to be generated.

BedrockProvider(model_kwargs={"maxTokens":200}, ...)

Here is another example, where anthropic provider will be created with the values for max_tokens and temperature, when claude-2 model is selected.

jupyter lab --AiExtension.model_parameters {"anthropic:claude-2":{"max_tokens":1024,"temperature":0.9}}

The above will result in the following LLM class to be generated.

AnthropicProvider(max_tokens=1024, temperature=0.9, ...)

Configuring as a config file

This configuration can also be specified in a config file in json format. The config can be loaded by specifying the path to the config file.

jupyter lab --config <config.json>

Here is an example for configuring the bedrock provider for ai21.j2-mid-v1 model.

{
    "AiExtension": { 
        "model_parameters": {
            "bedrock:ai21.j2-mid-v1": {
                "model_kwargs": {
                    "maxTokens": 200
                }
            }
        }
    }
}

Magic command samples

%%ai sagemaker-endpoint:jumpstart-dft-meta-textgeneration-llama-2-7b -q {"inputs":"<prompt>","parameters":{"max_new_tokens":64,"top_p":0.9,"temperature":0.6,"return_full_text":false}} -n us-east-1 -p [0].generation -m {"endpoint_kwargs":{"CustomAttributes":"accept_eula=true"}} -f text
Translate English to French:
sea otter => loutre de mer
peppermint => menthe poivrée
plush girafe => girafe peluche
cheese =>

%%ai bedrock:ai21.j2-mid-v1 -m {"model_kwargs":{"maxTokens":256}} -f text
Write me a short story on Python language

[ellisonbg]

Not sure this config should be in the UI. Might be better matched for traitlet's based config.

[3coins]

@ellisonbg
Although, providing these options as traitlets will allow users to specify them at runtime, Jupyter AI currently doesn’t have a mechanism to tie these attributes with a particular provider/model; we need a way to specify complete settings config as a runtime input. Also, in certain enterprise environments, a user might not have an ability to control the runtime parameters or have the option to restart the server and supply these values. I realize that the UX as an input field for these options might not be ideal, but providing this option in the settings panel gives a path forward to the user in controlling the model response.

[andrii-i]

Was only able to test UI so far (don't have Bedrock access yet), it looks and works well, new fields are shown, format verification on them works.

[JasonWeill]

@meeseeksdev please backport to 1.x

[lumberbot-app]

Owee, I'm MrMeeseeks, Look at me.

There seem to be a conflict, please backport manually. Here are approximate instructions:

1. Checkout backport branch and update it.

git checkout 1.x
git pull

2. Cherry pick the first parent branch of the this PR on top of the older branch:

git cherry-pick -x -m1 1fecfea6556501212ac9b7309edb7f0c98c69618

3. You will likely have some merge/cherry-pick conflict here, fix them and commit:

git commit -am 'Backport PR #430: Model parameters option to pass in model tuning, arbitrary parameters'

4. Push to a named branch:

git push YOURFORK 1.x:auto-backport-of-pr-430-on-1.x

5. Create a PR against branch 1.x, I would have named this PR:

"Backport PR #430 on branch 1.x (Model parameters option to pass in model tuning, arbitrary parameters)"

And apply the correct labels and milestones.

Congratulations — you did some good work! Hopefully your backport PR will be tested by the continuous integration and merged soon!

Remember to remove the Still Needs Manual Backport label once the PR gets merged.

If these instructions are inaccurate, feel free to suggest an improvement.