AI21 Jamba Mini 1.5

🚀 AI21 Jamba 1.5模型

AI21 Jamba 1.5是一系列先進的基礎模型，具備高效的長上下文處理能力和出色的性能。它們在多種語言和任務上表現優異，適用於商業場景，如函數調用、結構化輸出等。

🚀 快速開始

請注意，此版本將於2024年5月6日棄用。我們建議您過渡到新版本，可點擊此處查看。

✨ 主要特性

• 先進架構：AI21 Jamba 1.5系列模型是最先進的混合SSM - Transformer指令跟隨基礎模型。
• 高效推理：是市場上最強大、最高效的長上下文模型，推理速度比同類領先模型快達2.5倍。
• 多語言支持：支持英語、西班牙語、法語、葡萄牙語、意大利語、荷蘭語、德語、阿拉伯語和希伯來語。
• 商業優化：針對商業用例和功能進行了優化，如函數調用、結構化輸出（JSON）和基於文檔的生成。
• 靈活授權：根據Jamba開放模型許可證發佈，允許在許可條款下進行全面的研究和商業使用。

📦 安裝指南

運行優化的Mamba實現

要運行優化的Mamba實現，首先需要安裝mamba-ssm和causal-conv1d：

pip install mamba-ssm causal-conv1d>=1.2.0

同時，模型需要部署在CUDA設備上。

安裝vLLM

使用vLLM進行高效推理，需要安裝vLLM（要求版本0.5.4或更高）：

pip install vllm>=0.5.4

使用ExpertsInt8量化

使用ExpertsInt8量化技術，需要安裝vllm版本0.5.5或更高：

pip install vllm>=0.5.5

安裝transformers

使用transformers庫時，確保不使用4.44.0和4.44.1版本，因為這些版本存在限制Jamba架構運行的bug。

安裝trl進行微調

使用SFTTrainer進行微調時，需要安裝trl：

pip install trl

安裝bitsandbytes進行4位量化

使用QLoRA進行微調時，需要安裝bitsandbytes：

pip install bitsandbytes

💻 使用示例

使用vLLM運行模型

from vllm import LLM, SamplingParams
from transformers import AutoTokenizer

model = "ai21labs/AI21-Jamba-1.5-Mini"
number_gpus = 2

llm = LLM(model=model,
          max_model_len=200*1024,
          tensor_parallel_size=number_gpus)

tokenizer = AutoTokenizer.from_pretrained(model)

messages = [
   {"role": "system", "content": "You are an ancient oracle who speaks in cryptic but wise phrases, always hinting at deeper meanings."},
   {"role": "user", "content": "Hello!"},
]

prompts = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)

sampling_params = SamplingParams(temperature=0.4, top_p=0.95, max_tokens=100) 
outputs = llm.generate(prompts, sampling_params)

generated_text = outputs[0].outputs[0].text
print(generated_text)
#Output: Seek and you shall find. The path is winding, but the journey is enlightening. What wisdom do you seek from the ancient echoes?

使用ExpertsInt8量化運行模型

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.5-Mini",
          max_model_len=100*1024,
          quantization="experts_int8")

使用transformers運行模型

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("ai21labs/AI21-Jamba-1.5-Mini",
                                             torch_dtype=torch.bfloat16,
                                             attn_implementation="flash_attention_2",
                                             device_map="auto")

tokenizer = AutoTokenizer.from_pretrained("ai21labs/AI21-Jamba-1.5-Mini")

messages = [
   {"role": "system", "content": "You are an ancient oracle who speaks in cryptic but wise phrases, always hinting at deeper meanings."},
   {"role": "user", "content": "Hello!"},
]

input_ids = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors='pt').to(model.device)

outputs = model.generate(input_ids, max_new_tokens=216)

# Decode the output
conversation = tokenizer.decode(outputs[0], skip_special_tokens=True)

# Split the conversation to get only the assistant's response
assistant_response = conversation.split(messages[-1]['content'])[1].strip()
print(assistant_response)
# Output: Seek and you shall find. The path is winding, but the journey is enlightening. What wisdom do you seek from the ancient echoes?

以8位精度加載模型

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.5-Mini",
                                             torch_dtype=torch.bfloat16,
                                             attn_implementation="flash_attention_2",
                                             quantization_config=quantization_config)

在CPU上加載模型

from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained("ai21labs/AI21-Jamba-1.5-Mini",
                                             use_mamba_kernels=False)

工具使用示例

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("ai21labs/AI21-Jamba-1.5-Mini")

messages = [
    {
        "role": "user", 
        "content": "What's the weather like right now in Jerusalem and in London?"
    }
]

tools = [
    {
        'type': 'function', 
        'function': {
            'name': 'get_current_weather', 
            'description': 'Get the current weather', 
            'parameters': {
                'type': 'object', 
                'properties': {
                    'location': {'type': 'string', 'description': 'The city and state, e.g. San Francisco, CA'}, 
                    'format': {'type': 'string', 'enum': ['celsius', 'fahrenheit'], 'description': 'The temperature unit to use. Infer this from the users location.'}
                }, 
                'required': ['location', 'format']
            }
        }
    }
]

prompt = tokenizer.apply_chat_template(
    messages,
    tools=tools,
    tokenize=False,
)

將工具響應反饋給模型

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("ai21labs/AI21-Jamba-1.5-Mini")

# Note that you must send the tool responses in the same order as the model called the tools:
messages = [
    {
        "role": "user",
        "content": "What's the weather like right now in Jerusalem and in London?"
    },
    {
        "role": "assistant",
        "content": null,
        "tool_calls": [
            {
                "name": "get_current_weather",
                "arguments": "{\"location\": \"Jerusalem\", \"format\": \"celsius\"}"
            },
            {
                "name": "get_current_weather",
                "arguments": "{\"location\": \"London\", \"format\": \"celsius\"}"
            }
        ]
    },
    {
        "role": "tool",
        "content": "The weather in Jerusalem is 18 degrees celsius."
    },
    {
        "role": "tool",
        "content": "The weather in London is 8 degrees celsius."
    }
]

tool_use_prompt = tokenizer.apply_chat_template(
    messages,
    tools=tools,
    tokenize=False,
)

將文檔附加到Jamba 1.5提示

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("ai21labs/AI21-Jamba-1.5-Mini")

messages = [
        {
            "role": "user",
            "content": "Who wrote Harry Potter?"
        }
]

documents = [
        {
            "text": "Harry Potter is a series of seven fantasy novels written by British author J. K. Rowling.",
            "title": "Harry Potter"
        },
        {
            "text": "The Great Gatsby is a novel by American writer F. Scott Fitzgerald.",
            "title": "The Great Gatsby",
            "country": "United States",
            "genre": "Novel"

        }
]

prompt = tokenizer.apply_chat_template(
    messages,
    documents=documents,
    tokenize=False,
)

# Output: J. K. Rowling

使用JSON模式

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("ai21labs/AI21-Jamba-1.5-Mini")
messages = [
    {'role':'user', 
     'content':'Describe the first American president. Include year of birth (number) and name (string).'}
    ]
prompt = tokenizer.apply_chat_template(messages,
                                       tokenize=False,
                                       add_generation_prompt=False,
                                       knobs={"response_format": "json_object", "is_set": True})

#Output: "{ "year of birth": 1732, "name": "George Washington." }"

📚 詳細文檔

模型詳情

常見基準測試結果

RULER基準測試 - 有效上下文長度

多語言MMLU

🔧 技術細節

工具使用

Jamba 1.5支持根據Huggingface的工具使用API進行工具使用。用戶定義的工具會插入到聊天模板的專用部分，模型經過訓練可以識別該部分。

基於文檔的生成

Jamba在聊天模板中包含特定的“文檔”部分，用於標準化基於給定文檔集回答問題或遵循指令的過程。模型經過訓練會關注該部分，以這種方式格式化任務時，基於文檔的生成任務表現會更好。

JSON模式

Jamba 1.5經過訓練，具有特定的“旋鈕”，可幫助引導模型實現常見請求的行為。通過在系統消息中包含特定的預定義文本，可以啟用每種行為。為了方便使用，這些行為在Jamba 1.5的聊天模板中作為標誌提供，可以通過向聊天模板傳遞適當的參數來切換。

📄 許可證

該模型根據Jamba開放模型許可證發佈，這是一種寬鬆的許可證，允許在許可條款下進行全面的研究和商業使用。如果您需要為自己的需求許可該模型，請與我們聯繫。

關於AI21

AI21為企業構建可靠、實用且可擴展的AI解決方案。Jamba模型可在AI21 Studio和領先的雲合作伙伴處使用。如果您想了解Jamba 1.5 Mini和Jamba 1.5 Large如何為您的組織帶來實際價值，歡迎與我們交流。