Introduction to JumpStart - Text Embedding

This notebook's CI test result for us-west-2 is as follows. CI test results in other regions can be found at the end of the notebook.

Welcome to Amazon SageMaker JumpStart! You can use JumpStart to solve many Machine Learning tasks through one-click in SageMaker Studio, or through SageMaker JumpStart API.

In this demo notebook, we demonstrate how to use the JumpStart API for Text Embedding. Text Embedding is the task of converting text into meaningful vector representations. We show how to use pre-trained model to find text embeddings. Additionally, we also show how to leverage these embeddings for Semantic Textual Similarity.

1. Set Up
2. Select a model
3. Retrieve JumpStart Artifacts & Deploy an Endpoint
4. Query endpoint and parse response
5. Semantic Textual Similarity
6. Clean up the endpoint

Note: This notebook was tested on ml.t3.medium instance in Amazon SageMaker Studio with Python 3 (Data Science) kernel and in Amazon SageMaker Notebook instance with conda_python3 kernel.

1. Set Up

Before executing the notebook, there are some initial steps required for set up. This notebook requires latest version of sagemaker and ipywidgets

Permissions and environment variables

To host on Amazon SageMaker, we need to set up and authenticate the use of AWS services. Here, we use the execution role associated with the current notebook as the AWS account role with SageMaker access.

2. Select a model

Here, we download jumpstart model_manifest file from the jumpstart s3 bucket, filter-out all the Text Embedding models and select a model for inference.

Chose a model for Inference

3. Retrieve JumpStart Artifacts & Deploy an Endpoint

We start by retrieving the deploy_image_uri, deploy_source_uri, and model_uri for the pre-trained model. To host the pre-trained model, we create an instance of sagemaker.model.Model and deploy it. This may take a few minutes.

4. Query endpoint and parse response

Input to the endpoint is any string of text dumped in json and encoded in utf-8 format. Output of the endpoint is a json with the text embedding.

Below, we put in some example input text. You can put in any text and the model will summarize the text.

5. Semantic Textual Similarity

A use case of sentence embedding is to cluster together sentences with similar semantic meaning. In the example below we compute the embeddings of sentences in three categories: pets, cities in the U.S., and color. We see that sentences originating from the same category have much closer embedding vectors than those from different categories.

Specifically, the code will do the following:

• The endpoint that you have created above will output an embedding vector for each sentence;
• The distance between any pair of sentences is computed by the cosine similarity of corresponded embedding vectors;
• A heatmap is created to visualize the distance between any pair of sentences in the embedding space. Darker the color, larger the cosine similarity (smaller the distance).

Note. Cosine similarity of two vectors is the inner product of the normalized vectors (scale down to have length 1).

6. Clean up the endpoint

Notebook CI Test Results

This notebook was tested in multiple regions. The test results are as follows, except for us-west-2 which is shown at the top of the notebook.