AI21 Labs Jamba-Instruct model is now available in Amazon Bedrock

We are excited to announce the availability of the Jamba-Instruct large language model (LLM) in Amazon Bedrock. Jamba-Instruct is built by AI21 Labs, and most notably supports a 256,000-token context window, making it especially useful for processing large documents and complex Retrieval Augmented Generation (RAG) applications.

What is Jamba-Instruct

Jamba-Instruct is an instruction-tuned version of the Jamba base model, previously open sourced by AI21 Labs, which combines a production grade-model, Structured State Space (SSM) technology, and Transformer architecture. With the SSM approach, Jamba-Instruct is able to achieve the largest context window length in its model size class while also delivering the performance traditional transformer-based models provide. These models yield a performance boost over AI21’s previous generation of models, the Jurassic-2 family of models. For more information about the hybrid SSM/Transformer architecture, refer to the Jamba: A Hybrid Transformer-Mamba Language Model whitepaper.

Get started with Jamba-Instruct

To get started with Jamba-Instruct models in Amazon Bedrock, first you need to get access to the model.

1. On the Amazon Bedrock console, choose Model access in the navigation pane.
2. Choose Modify model access.
3. Select the AI21 Labs models you want to use and choose Next.
4. Choose Submit to request model access.

For more information, refer to Model access.

Next, you can test the model either in the Amazon Bedrock Text or Chat playground.

Example use cases for Jamba-Instruct

Jamba-Instruct’s long context length is particularly well-suited for complex Retrieval Augmented Generation (RAG) workloads, or potentially complex document analysis. For example, it would be suitable for detecting contradictions between different documents or analyzing one document in the context of another. The following is an example prompt suitable for this use case:

You are an expert research assistant; 
you are to note any contradictions between the first document and second document provided: 

Document 1: 
{the document content} 

Document 2: 
{the document content} 

Contradictions:

You can also use Jamba for query augmentation, a technique where an original query is transformed into related queries, for purposes of optimizing RAG applications. For example:

You are a curious and novel researcher, 
who is highly interested in getting all the relevant information on a specific topic. 
Given an original query, you would like to generate up to 10 related queries. 
These queries should be grounded in the original query, but nevertheless new:

Original Query:
{Original Query}

New Queries:

You can also use Jamba for standard LLM operations, such as summarization and entity extraction.

Prompt guidance for Jamba-Instruct can be found in the AI21 model documentation. For more information about Jamba-Instruct, including relevant benchmarks, refer to Built for the Enterprise: Introducing AI21’s Jamba-Instruct Model.

Programmatic access

You can also access Jamba-Instruct through an API, using Amazon Bedrock and AWS SDK for Python (Boto3). For installation and setup instructions, refer to the quickstart. The following is an example code snippet:

import boto3
import json

bedrock = boto3.client(service_name="bedrock-runtime")

prompt = "[INST] INSERT YOUR PROMPT HERE [/INST]"

body = json.dumps({
    "prompt": prompt,
    "max_tokens": 256,
    "top_p": 0.8,
    "temperature": 0.7,
})

modelId = "ai21.jamba-instruct-v1:0"

accept = "application/json"
contentType = "application/json"

response = bedrock.invoke_model(
    body=body,
    modelId=modelId,
    accept=accept,
    contentType=contentType
)

print(json.loads(response.get('body').read()))