Update README.md

README.md

@@ -5,16 +5,13 @@ license_link: https://www.ai21.com/jamba-open-model-license/
 ---
 # Model Information
 
-The AI21 Jamba 1.6 family of models is state-of-the-art, hybrid SSM-Transformer instruction following foundation models. The Jamba models are the most powerful & efficient long-context models on the market, which deliver up to 2.5X faster inference than leading models of comparable sizes.
+Built with hybrid SSM-Transformer architecture, the Jamba 1.6 family of models outperform other open, instruction-following foundation models on quality, speed, and long context performance, and rival leading closed models on quality. As open models, Jamba Mini 1.6 (12B active/52B total) and Jamba Large 1.6 (94B active/398B total) are available for private deployment, either in VPC or on-premise, and demonstrate superior performance on the kind of long context tasks that matter most to enterprises, such as RAG workflows and grounded question answering across lengthy documents.
 
-The models demonstrate superior long context handling, speed, and quality. They mark the first time a non-Transformer model has been successfully scaled to the quality and strength of the market’s leading models. 
+The models are released under the Jamba Open Model License, a permissive license allowing full research use and commercial use under the license terms.
 
-[Jamba 1.6 Mini](https://huggingface.co/ai21labs/AI21-Jamba-1.6-Mini) (12B active/52B total) and [Jamba 1.6 Large](https://huggingface.co/ai21labs/AI21-Jamba-1.5-Large) (94B active/398B total) are also optimized for business use cases and capabilities such as function calling, structured output (JSON), and grounded generation.
-
-The models are released under the [Jamba Open Model License](https://www.ai21.com/licenses/jamba-open-model-license), a permissive license allowing full research use and commercial use under the license terms. If you need to license the model for your needs, [talk to us](https://www.ai21.com/talk-to-us). 
-
-For more details of this model, see the white paper and the release [blog post](https://www.ai21.com/blog/announcing-jamba-model-family).
+If you need to license the model for your needs, talk to us.
 
+For more details of this model, see the release blog post.
 ## Model Details
 
 - **Developed by:** [AI21](https://www.ai21.com)
@@ -25,23 +22,6 @@ For more details of this model, see the white paper and the release [blog post](
 - **Supported languages:** English, Spanish, French, Portuguese, Italian, Dutch, German, Arabic and Hebrew
 
 
-## Results on common benchmarks
-
-### RULER Benchmark - Effective context length
-
-|Models|Claimed Length|Effective Length|4K|8K|16K|32K|64K|128K|256K|
-|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|
-Jamba 1.6 Large (94B/398B)|256K|256K|<ins>96.7</ins>|<ins>96.6</ins>|<ins>96.4</ins>|<ins>96.0</ins>|<ins>95.4</ins>|<ins>95.1</ins>|<ins>93.9</ins>|
-Jamba 1.6 Mini (12B/52B)|256K|256K|<ins>95.7</ins>|<ins>95.2</ins>|<ins>94.7</ins>|<ins>93.8</ins>|<ins>92.7</ins>|<ins>89.8</ins>|<ins>86.1</ins> |
-Gemini 1.5 Pro|1M|>128K|<ins>96.7</ins>|<ins>95.8</ins>|<ins>96.0</ins>|<ins>95.9</ins>|<ins>95.9</ins>|<ins>94.4</ins>| -- |
-GPT-4 1106-preview |128K|64K|<ins>96.6</ins>|<ins>96.3</ins>|<ins>95.2</ins>|<ins>93.2</ins>|<ins>87.0</ins>|81.2| -- |
-Llama 3.1 70B|128K|64K|<ins>96.5</ins>|<ins>95.8</ins>|<ins>95.4</ins>|<ins>94.8</ins>|<ins>88.4</ins>|66.6| -- |
-Command R-plus (104B)|128K|32K|<ins>95.6</ins>|<ins>95.2</ins>|<ins>94.2</ins>|<ins>92.0</ins>|84.3|63.1| -- |
-Llama 3.1 8B|128K|32K|<ins>95.5</ins>|<ins>93.8</ins>|<ins>91.6</ins>|<ins>87.4</ins>|84.7|77.0| -- |
-Mistral Large 2 (123B)|128K|32K|<ins>96.2</ins>|<ins>96.1</ins>|<ins>95.1</ins>|<ins>93.0</ins>|78.8|23.7| -- |
-Mixtral 8x22B (39B/141B)|64K|32K|<ins>95.6</ins>|<ins>94.9</ins>|<ins>93.4</ins>|<ins>90.9</ins>|84.7|31.7| -- |
-Mixtral 8x7B (12.9B/46.7B)|32K|32K|<ins>94.9</ins>|<ins>92.1</ins>|<ins>92.5</ins>|<ins>85.9</ins>|72.4|44.5| -- |
-
 # Usage
 ## Prerequisites
 
@@ -54,18 +34,18 @@ You also have to have the model on a CUDA device.
 
 ## Run the model with vLLM
 
-The recommended way to perform efficient inference with Jamba 1.6 Mini is using [vLLM](https://docs.vllm.ai/en/latest/). First, make sure to install vLLM (version 0.5.4 or higher is required)
+The recommended way to perform efficient inference with Jamba Mini 1.6 is using [vLLM](https://docs.vllm.ai/en/latest/). First, make sure to install vLLM (version 0.5.4 or higher is required)
 ```bash
 pip install vllm>=0.5.4
 ```
 
-In the example below, `number_gpus` should match the number of GPUs you want to deploy Jamba 1.6 Mini on. A minimum of 2 80GB GPUs is required.
+In the example below, `number_gpus` should match the number of GPUs you want to deploy Jamba Mini 1.6 on. A minimum of 2 80GB GPUs is required.
 
 ```python
 from vllm import LLM, SamplingParams
 from transformers import AutoTokenizer
 
-model = "ai21labs/AI21-Jamba-1.6-Mini"
+model = "ai21labs/AI21-Jamba-Mini-1.6"
 number_gpus = 2
 
 llm = LLM(model=model,
@@ -94,7 +74,7 @@ With the default BF16 precision on 2 80GB A100 GPUs and default vLLM configurati
 <u>Note:</u> vLLM's `main` branch has some memory utilization improvements specific to the Jamba architecture that allow using the full 256K context length on 2 80 GPUs. You can [build vLLM from source](https://docs.vllm.ai/en/latest/getting_started/installation.html#build-from-source) if you wish to make use of them.
 
 ### ExpertsInt8 quantization
-We've developed an innovative and efficient quantization technique, [ExpertsInt8](https://www.ai21.com/blog/announcing-jamba-model-family#:~:text=Like%20all%20models%20in%20its%20size%20class%2C%20Jamba%201.5%20Large%20can%E2%80%99t%20be%20loaded%20in%20full%20(FP32)%20or%20half%20(FP16/BF16)%20precision%20on%20a%20single%20node%20of%208%20GPUs.%20Dissatisfied%20with%20currently%20available%20quantization%20techniques%2C%20we%20developed%20ExpertsInt8%2C%20a%20novel%20quantization%20technique%20tailored%20for%20MoE%20models.), designed for MoE models deployed in vLLM, including Jamba models. Using it, you'll be able to deploy Jamba 1.5 Mini on a single 80GB GPU.
+We've developed an innovative and efficient quantization technique, [ExpertsInt8](https://www.ai21.com/blog/announcing-jamba-model-family#:~:text=Like%20all%20models%20in%20its%20size%20class%2C%20Jamba%201.6%20Large%20can%E2%80%99t%20be%20loaded%20in%20full%20(FP32)%20or%20half%20(FP16/BF16)%20precision%20on%20a%20single%20node%20of%208%20GPUs.%20Dissatisfied%20with%20currently%20available%20quantization%20techniques%2C%20we%20developed%20ExpertsInt8%2C%20a%20novel%20quantization%20technique%20tailored%20for%20MoE%20models.), designed for MoE models deployed in vLLM, including Jamba models. Using it, you'll be able to deploy Jamba Mini 1.6 on a single 80GB GPU.
 
 In order to use ExpertsInt8, you need to use vllm version 0.5.5 or higher: `pip install vllm>=0.5.5`
 
@@ -104,7 +84,7 @@ import os
 os.environ['VLLM_FUSED_MOE_CHUNK_SIZE']='32768'    # This is a workaround a bug in vLLM's fused_moe kernel
 
 from vllm import LLM
-llm = LLM(model="ai21labs/AI21-Jamba-1.6-Mini",
+llm = LLM(model="ai21labs/AI21-Jamba-Mini-1.6",
           max_model_len=100*1024,
           quantization="experts_int8")
 ```
@@ -112,18 +92,18 @@ llm = LLM(model="ai21labs/AI21-Jamba-1.6-Mini",
 
 ## Run the model with `transformers`
 
-The following example loads Jamba 1.6 Mini to the GPU in BF16 precision, uses optimized [FlashAttention2](https://github.com/Dao-AILab/flash-attention) and Mamba kernels, and parallelizes the model across multiple GPUs using [accelerate](https://huggingface.co/docs/accelerate/index). Note that in half precision (FP16/BF16), Jamba 1.5 Mini is too large to fit on a single 80GB GPU, so you'll need at least 2 such GPUs.
+The following example loads Jamba Mini 1.6 to the GPU in BF16 precision, uses optimized [FlashAttention2](https://github.com/Dao-AILab/flash-attention) and Mamba kernels, and parallelizes the model across multiple GPUs using [accelerate](https://huggingface.co/docs/accelerate/index). Note that in half precision (FP16/BF16), Jamba Mini 1.6 is too large to fit on a single 80GB GPU, so you'll need at least 2 such GPUs.
 
 ```python
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-model = AutoModelForCausalLM.from_pretrained("ai21labs/AI21-Jamba-1.6-Mini",
+model = AutoModelForCausalLM.from_pretrained("ai21labs/AI21-Jamba-Mini-1.6",
                                              torch_dtype=torch.bfloat16,
                                              attn_implementation="flash_attention_2",
                                              device_map="auto")
 
-tokenizer = AutoTokenizer.from_pretrained("ai21labs/AI21-Jamba-1.6-Mini")
+tokenizer = AutoTokenizer.from_pretrained("ai21labs/AI21-Jamba-Mini-1.6")
 
 messages = [
    {"role": "system", "content": "You are an ancient oracle who speaks in cryptic but wise phrases, always hinting at deeper meanings."},
@@ -145,7 +125,7 @@ print(assistant_response)
 
 <u>Note:</u> Versions 4.44.0 and 4.44.1 of `transformers` have a bug that restricts the ability to run the Jamba architecture. Make sure you're not using these versions.
 
-<u>Note:</u> If you're having trouble installing `mamba-ssm` and `causal-conv1d` for the optimized Mamba kernels, you can run Jamba 1.5 Mini without them, at the cost of extra latency. In order to do that, add the kwarg `use_mamba_kernels=False` when loading the model via `AutoModelForCausalLM.from_pretained()`.
+<u>Note:</u> If you're having trouble installing `mamba-ssm` and `causal-conv1d` for the optimized Mamba kernels, you can run Jamba Mini 1.6 without them, at the cost of extra latency. In order to do that, add the kwarg `use_mamba_kernels=False` when loading the model via `AutoModelForCausalLM.from_pretained()`.
 
 <details><summary><strong>Load the model in 8-bit</strong></summary>
 
@@ -155,7 +135,7 @@ print(assistant_response)
 from transformers import AutoModelForCausalLM, BitsAndBytesConfig
 quantization_config = BitsAndBytesConfig(load_in_8bit=True,
                                          llm_int8_skip_modules=["mamba"])
-model = AutoModelForCausalLM.from_pretrained("ai21labs/AI21-Jamba-1.6-Mini",
+model = AutoModelForCausalLM.from_pretrained("ai21labs/AI21-Jamba-Mini-1.6",
                                              torch_dtype=torch.bfloat16,
                                              attn_implementation="flash_attention_2",
                                              quantization_config=quantization_config)
@@ -165,11 +145,11 @@ model = AutoModelForCausalLM.from_pretrained("ai21labs/AI21-Jamba-1.6-Mini",
 
 <details><summary><strong>Load the model on CPU</strong></summary>
 
-If you don't have access to a GPU, you can also load and run Jamba 1.6 Mini on a CPU. Note this will result in poor inference performance.
+If you don't have access to a GPU, you can also load and run Jamba Mini 1.6 on a CPU. Note this will result in poor inference performance.
 
 ```python
 from transformers import AutoModelForCausalLM
-model = AutoModelForCausalLM.from_pretrained("ai21labs/AI21-Jamba-1.6-Mini",
+model = AutoModelForCausalLM.from_pretrained("ai21labs/AI21-Jamba-Mini-1.6",
                                              use_mamba_kernels=False)
 ```
 </details>
@@ -179,7 +159,7 @@ model = AutoModelForCausalLM.from_pretrained("ai21labs/AI21-Jamba-1.6-Mini",
 # Model features
 
 ## Tool use with Jamba
-Jamba 1.6 supports tool use capabilities in accordance with Huggingface's tool use API. The tools defined by the user are inserted into a dedicated section in the chat template which the model was trained to recognize. 
+Jamba Mini 1.6 supports tool use capabilities in accordance with Huggingface's tool use API. The tools defined by the user are inserted into a dedicated section in the chat template which the model was trained to recognize. 
 
 Given a conversation that contains tools, the model can output content, tool invocations or both.
 
@@ -189,7 +169,7 @@ Given a conversation that contains tools, the model can output content, tool inv
 ```python
 from transformers import AutoTokenizer
 
-tokenizer = AutoTokenizer.from_pretrained("ai21labs/AI21-Jamba-1.6-Mini")
+tokenizer = AutoTokenizer.from_pretrained("ai21labs/AI21-Jamba-Mini-1.6")
 
 messages = [
     {
@@ -242,7 +222,7 @@ The `arguments` field for each tool call can be either a dict or a JSON string.
 ```python
 from transformers import AutoTokenizer
 
-tokenizer = AutoTokenizer.from_pretrained("ai21labs/AI21-Jamba-1.6-Mini")
+tokenizer = AutoTokenizer.from_pretrained("ai21labs/AI21-Jamba-Mini-1.6")
 
 # Note that you must send the tool responses in the same order as the model called the tools:
 messages = [
@@ -293,12 +273,12 @@ A common use-case for LLMs is grounded generation and RAG, where the model is re
 
 Similar to tools, which are given as an external argument to the model in addition to the conversation, documents are provided in a similar way. To support document-level metadata, a document is defined as a dictionary with key-values of your choosing. These are formatted within the chat template. Two keys that get special treatment are "title", which is formatted at the top of the document if present, and "text" which is a required field and defines the actual text of the document.
 
-<details><summary><strong>Ataching documents to Jamba 1.6 prompt</strong></summary>
+<details><summary><strong>Ataching documents to Jamba Mini 1.6 prompt</strong></summary>
 
 ```python
 from transformers import AutoTokenizer
 
-tokenizer = AutoTokenizer.from_pretrained("ai21labs/AI21-Jamba-1.6-Mini")
+tokenizer = AutoTokenizer.from_pretrained("ai21labs/AI21-Jamba-Mini-1.6")
 
 messages = [
         {
@@ -334,16 +314,16 @@ prompt = tokenizer.apply_chat_template(
 </details>
 
 ## JSON mode
-Jamba 1.6 was trained with specific “knobs”, which help steer the model towards commonly requested behaviors. Each behavior is enabled by including specific pre-defined text in the system message. For ease of use, we've included them as flags in Jamba 1.5's chat template, so they can be toggled by passing appropriate arguments to the chat template.
+Jamba 1.6 was trained with specific “knobs”, which help steer the model towards commonly requested behaviors. Each behavior is enabled by including specific pre-defined text in the system message. For ease of use, we've included them as flags in Jamba 1.6's chat template, so they can be toggled by passing appropriate arguments to the chat template.
 
-Jamba 1.6 was trained to produce valid JSONs when requested to. It does so naturally, but when the JSON mode knob is activated the likelihood of a valid json increases considerably. In JSON mode, Jamba 1.5 will attempt to output a valid JSON regardless of the user request. However, it is highly recommended to specify information about the expected json schema in the user request or system message to get the best results, as shown in the example below.
+Jamba 1.6 was trained to produce valid JSONs when requested to. It does so naturally, but when the JSON mode knob is activated the likelihood of a valid json increases considerably. In JSON mode, Jamba 1.6 will attempt to output a valid JSON regardless of the user request. However, it is highly recommended to specify information about the expected json schema in the user request or system message to get the best results, as shown in the example below.
 
 <details><summary><strong>Usage of JSON knob in Jamba 1.6</strong></summary>
 
 ```python
 from transformers import AutoTokenizer
 
-tokenizer = AutoTokenizer.from_pretrained("ai21labs/AI21-Jamba-1.6-Mini")
+tokenizer = AutoTokenizer.from_pretrained("ai21labs/AI21-Jamba-Mini-1.6")
 messages = [
     {'role':'user', 
      'content':'Describe the first American president. Include year of birth (number) and name (string).'}
@@ -379,9 +359,9 @@ from datasets import load_dataset
 from trl import SFTTrainer, SFTConfig
 from peft import LoraConfig
 
-tokenizer = AutoTokenizer.from_pretrained("ai21labs/AI21-Jamba-1.6-Mini")
+tokenizer = AutoTokenizer.from_pretrained("ai21labs/AI21-Jamba-Mini-1.6")
 model = AutoModelForCausalLM.from_pretrained(
-    "ai21labs/AI21-Jamba-1.6-Mini",
+    "ai21labs/AI21-Jamba-Mini-1.6",
     device_map="auto",
     torch_dtype=torch.bfloat16,
     attn_implementation="flash_attention_2",
@@ -434,14 +414,14 @@ from datasets import load_dataset
 from trl import SFTTrainer, SFTConfig
 from peft import LoraConfig
 
-tokenizer = AutoTokenizer.from_pretrained("ai21labs/AI21-Jamba-1.6-Mini")
+tokenizer = AutoTokenizer.from_pretrained("ai21labs/AI21-Jamba-Mini-1.6")
 quantization_config = BitsAndBytesConfig(
     load_in_4bit=True,
     bnb_4bit_quant_type="nf4",
     bnb_4bit_compute_dtype=torch.bfloat16,
 )
 model = AutoModelForCausalLM.from_pretrained(
-    "ai21labs/AI21-Jamba-1.6-Mini",
+    "ai21labs/AI21-Jamba-Mini-1.6",
     device_map="auto",
     quantization_config=quantization_config,
     torch_dtype=torch.bfloat16,
@@ -489,4 +469,4 @@ pip install bitsandbytes
 # About AI21
 
 AI21 builds reliable, practical, and scalable AI solutions for the enterprise. The Jamba models are available in the [AI21 Studio](https://www.ai21.com/studio) and in leading cloud partners.
-To learn more about how Jamba 1.6 Mini and Jamba 1.6 Large can bring real world value to your organization, let’s talk.
+To learn more about how Jamba Mini 1.6 and Jamba Large 1.6 can bring real world value to your organization, let’s talk.