Support SageMaker Endpoints in chat

[dlqqq]

Description

• See Specify request/response data types for SageMaker endpoints in chat UI #193

Allows for usage of language models hosted on SageMaker Endpoints. The only constraint is that both the request and response of the model are in JSON.

Screen.Recording.2023-05-30.at.10.44.09.AM.mov

Providers may now declare fields, which are keyword arguments expected by the constructor. Each field is declared via a Field object, which is defined as follows:

class TextField(BaseModel):
    type: Literal["text"] = "text"
    key: str
    label: str

class MultilineTextField(BaseModel):
    type: Literal["text-multiline"] = "text-multiline"
    key: str
    label: str

Field = Union[TextField, MultilineTextField]

The backend stores the values of each fields under config.fields.[<model-id>], where <model-id> is the global model ID of a model. The backend will also automatically read this config object and pass fields for a model to a model provider's constructor as keyword arguments.

In the case of SageMaker Endpoints (SMEP), there are 3 additional fields declared:

class SmEndpointProvider(BaseProvider, SagemakerEndpoint):
    id = "sagemaker-endpoint"
    ...
    fields = [
        TextField(
            key="region_name",
            label="Region name",
        ),
        MultilineTextField(
            key="request_schema",
            label="Request schema",
        ),
        TextField(
            key="response_path",
            label="Response path",
        )
    ]

The constructor of the SMEP provider has been modified to accept two new keyword arguments: request_schema and response_path.

• request_schema is a JSON string. Any values that match the exact literal "<prompt>" are substituted with the value of the prompt. For example, when using flan-t5-xl on SMEP via SM Jumpstart, the request schema should be {"text_inputs":"<prompt>"}
• response_path is a JSON path as defined in the specification. So for example, when using flan-t5-xl on SMEP, this should be generated_texts.[0].

The change to the SMEP provider constructor should enable usage in magics fairly easily, but unfortunately I cannot finish this before my vacation tomorrow.

Follow-up items

• Replicate this on SageMaker and verify that you're able to reproduce my results.
• Add documentation for this
• Do additional UI testing and iron out bugs that you discover.
  • The UI state in the Chat Settings component is fairly complex; this should be refactored and decomposed into a larger set of smaller components. However, due to time constraints, this may not happen for the foreseeable future.
• Replicate this support in magics (see cannot call SageMaker Endpoint in magics #36).
  • There needs to be a standard syntax for passing these additional kwargs/fields to a model provider. Ideally, each field can be specified as a separate argument to the magic, rather than necessitating each user pass an ugly and difficult-to-read JSON blob as a single argument.

[JasonWeill]

We should have at least basic info in the user docs to indicate how people can start using models in Jupyter AI via SageMaker.

[dlqqq]

I've added SageMaker Endpoints in the user documentation.

[3coins]

@JasonWeill
Tested successfully that this works now, one thing that tripped me was that while adding the request schema, I copied it from the Studio notebook, this failed json decode while running the prediction, as the version from studio was a python object vs JSON. Created #202 for adding validation.

Here are working request/response schemas for flan-t5-xl model.

Request schema

{"text_inputs":"<prompt>", "max_length":50, "num_return_sequences":3, "top_k":50, "top_p":0.95, "do_sample":true}

Response schema

generated_texts.[0]

[JasonWeill]

UI quibble: I don't love that the headings ("Language model", "Embedding model", etc) are now encased in rectangles with rounded corners. They look like buttons, but they're not operable as buttons.

[3coins]

@JasonWeill
Updated the headers, here is a screenshot.

[3coins]

The CI failure is not related to the changes here.