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如何为您的组织带来实际价值，欢迎与我们交流。