Jamba V0.1 Chat Multilingual

🚀 Jamba-v0.1聊天多語言模型

本模型是一個基於transformers庫微調的聊天模型，通過對ai21labs/Jamba-v0.1進行微調，使其能夠作為聊天機器人使用。經過初步測試，該模型在英語和其他語言的對話中表現良好。

🚀 快速開始

本部分將介紹如何快速使用該模型進行文本生成。

安裝依賴

Jamba要求使用transformers版本4.39.0或更高：

pip install transformers>=4.39.0

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

pip install mamba-ssm causal-conv1d>=1.2.0

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

代碼示例

from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
import torch

double_quant_config = BitsAndBytesConfig(
   load_in_4bit=True,
   bnb_4bit_use_double_quant=True,
   bnb_4bit_compute_dtype=torch.float16
)

model = AutoModelForCausalLM.from_pretrained("lightblue/Jamba-v0.1-chat-multilingual", device_map="auto", quantization_config=double_quant_config, trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained("lightblue/Jamba-v0.1-chat-multilingual")

input_text = """<|im_start|>system 
You are a helpful AI assistant.
<|im_end|> 
<|im_start|>user
What is the most interesting fact about kangaroos that you know?
<|im_end|> 
<|im_start|>assistant
"""

input_ids = tokenizer(input_text, return_tensors='pt').to(model.device)["input_ids"]

outputs = model.generate(input_ids, max_new_tokens=256, temperature=0.0, repetition_penalty=1.1)

print(tokenizer.batch_decode(outputs)[0])
# <|startoftext|><|im_start|>system 
# You are a helpful AI assistant.
# <|im_end|> 
# <|im_start|>user
# What is the most interesting fact about kangaroos that you know?
# <|im_end|> 
# <|im_start|>assistant
# One of the most interesting facts about kangaroos is their unique reproductive system, which involves embryonic diapause and multiple births. Female kangaroos can have up to three joeys at different stages of development simultaneously: one fully developed in the pouch, another developing inside her body, and an embryo waiting for its turn.<|im_end|>
# <|endoftext|>

注意事項

• 此代碼會自動在任何輸入後附加 "<|startoftext|>" 特殊標記。在推理時，所有輸入都需要附加此標記，因為初步測試表明，不添加此標記會導致輸出錯誤。
• 你可以在不使用優化的Mamba內核的情況下運行模型，但不建議這樣做，因為這會導致顯著的延遲增加。若要這樣做，在加載模型時需要指定use_mamba_kernels=False。

✨ 主要特性

• 多語言支持：該模型經過訓練，可以在英語和其他語言的對話中表現良好。
• 微調實驗：通過對ai21labs/Jamba-v0.1進行微調，探索其作為聊天機器人的潛力。
• 可控輸出：可以通過系統消息來引導模型的輸出，例如控制輸出的語言難度、是否押韻等。

📦 安裝指南

依賴安裝

Jamba要求使用transformers版本4.39.0或更高：

pip install transformers>=4.39.0

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

pip install mamba-ssm causal-conv1d>=1.2.0

設備要求

你需要將模型部署在CUDA設備上。

💻 使用示例

基礎用法

from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
import torch

double_quant_config = BitsAndBytesConfig(
   load_in_4bit=True,
   bnb_4bit_use_double_quant=True,
   bnb_4bit_compute_dtype=torch.float16
)

model = AutoModelForCausalLM.from_pretrained("lightblue/Jamba-v0.1-chat-multilingual", device_map="auto", quantization_config=double_quant_config, trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained("lightblue/Jamba-v0.1-chat-multilingual")

input_text = """<|im_start|>system 
You are a helpful AI assistant.
<|im_end|> 
<|im_start|>user
What is the most interesting fact about kangaroos that you know?
<|im_end|> 
<|im_start|>assistant
"""

input_ids = tokenizer(input_text, return_tensors='pt').to(model.device)["input_ids"]

outputs = model.generate(input_ids, max_new_tokens=256, temperature=0.0, repetition_penalty=1.1)

print(tokenizer.batch_decode(outputs)[0])

高級用法

你可以通過系統消息來控制模型的輸出，例如讓模型用簡單英語、複雜英語或押韻的方式回答問題。

# 系統消息控制輸出為簡單英語
input_text = """<|startoftext|><|im_start|>system 
You are a helpful AI assistant. You write all answers in very simple English.
<|im_end|> 
<|im_start|>user
Write a 50 word analysis of why sausages are better than bacon.
<|im_end|> 
<|im_start|>assistant
"""

input_ids = tokenizer(input_text, return_tensors='pt').to(model.device)["input_ids"]

outputs = model.generate(input_ids, max_new_tokens=256, temperature=0.0, repetition_penalty=1.1)

print(tokenizer.batch_decode(outputs)[0])

# 系統消息控制輸出為複雜英語
input_text = """<|startoftext|><|im_start|>system 
You are a helpful AI assistant. You write all answers in very complex English.
<|im_end|> 
<|im_start|>user
Write a 50 word analysis of why sausages are better than bacon.
<|im_end|> 
<|im_start|>assistant
"""

input_ids = tokenizer(input_text, return_tensors='pt').to(model.device)["input_ids"]

outputs = model.generate(input_ids, max_new_tokens=256, temperature=0.0, repetition_penalty=1.1)

print(tokenizer.batch_decode(outputs)[0])

# 系統消息控制輸出為押韻
input_text = """<|startoftext|><|im_start|>system 
You are an AI assistant that answers all questions in rhyme.
<|im_end|> 
<|im_start|>user
Why is the sky blue?
<|im_end|> 
<|im_start|>assistant
"""

input_ids = tokenizer(input_text, return_tensors='pt').to(model.device)["input_ids"]

outputs = model.generate(input_ids, max_new_tokens=256, temperature=0.0, repetition_penalty=1.1)

print(tokenizer.batch_decode(outputs)[0])

📚 詳細文檔

模型詳情

初始測試結果

• 知識表現：模型在某些領域表現出了不錯的知識水平，但在其他領域可能會產生錯誤信息。
• 系統消息控制：通過系統消息可以比較容易地引導模型的輸出，例如控制輸出的語言難度、是否押韻等。
• 多輪對話準確性：初步測試顯示，模型在多輪對話中的準確性表現尚可。
• 多語言測試：模型能夠比較可靠地用輸入的語言進行回覆。

訓練詳情

訓練數據

• jondurbin/airoboros-3.2：一個約59K示例的英語LLM任務數據集，主要由GPT-4生成。該數據集包含了各種類型的任務，因此被選作主要的訓練數據。
• openchat/openchat_sharegpt4_dataset (僅GPT-4回覆)：一個約6K示例的多語言多輪對話數據集，包含了用戶與GPT-4的對話。由於jondurbin/airoboros-3.2中幾乎沒有多語言數據，而該項目是一家日本AI公司，需要模型能夠輸出日語，因此選擇了這個數據集來補充多語言數據。

數據準備代碼

import os
import pandas as pd
from datasets import load_dataset, Dataset, concatenate_datasets

os.environ['HF_HOME'] = "/workspace/hf_home"
os.environ['HF_HUB_ENABLE_HF_TRANSFER'] = "1"

boros_dataset = load_dataset("jondurbin/airoboros-3.2", split='train')

gpt4_df = pd.read_json("https://huggingface.co/datasets/openchat/openchat_sharegpt4_dataset/resolve/main/sharegpt_gpt4.json?download=true")
gpt4_df["conversations"] = gpt4_df["items"].apply(lambda x: [{'from': 'system', 'value': 'You are GPT-4, a helpful assistant.'}] + x)

gpt4_dataset = Dataset.from_pandas(gpt4_df[["conversations"]])

dataset = concatenate_datasets([gpt4_dataset, boros_dataset]).shuffle()

dataset.select_columns(["conversations"]).to_json("/workspace/airoboros-3.2_plus_openchat_sharegpt4.json")

訓練配置

base_model: ai21labs/Jamba-v0.1
trust_remote_code: true

load_in_8bit: false
load_in_4bit: true
strict: false

datasets:
  - path: /workspace/airoboros-3.2_plus_openchat_sharegpt4.json
    ds_type: json
    type: sharegpt
    conversation: chatml
dataset_prepared_path:
val_set_size: 0.01
output_dir: ./airoboros-3.2_plus_openchat_sharegpt4_one_epoch

sequence_len: 6000
sample_packing: true
pad_to_sequence_len: false
eval_sample_packing: true

use_wandb: true
wandb_project: axolotl
wandb_entity: peterd
wandb_name: airoboros-3.2_plus_openchat_sharegpt4

adapter: qlora
lora_r: 8
lora_alpha: 16
lora_dropout: 0.05
lora_target_linear: true

low_cpu_mem_usage: true
gradient_accumulation_steps: 4
micro_batch_size: 1
num_epochs: 1
optimizer: paged_adamw_8bit
lr_scheduler: cosine
learning_rate: 0.0002

train_on_inputs: false
group_by_length: false
bf16: auto
fp16:
tf32: false

gradient_checkpointing: true
gradient_checkpointing_kwargs:
  use_reentrant: false
early_stopping_patience:
resume_from_checkpoint:
local_rank:
logging_steps: 1
xformers_attention:
flash_attention: true

warmup_steps: 10
evals_per_epoch: 5
saves_per_epoch: 5
debug:
deepspeed: /workspace/axolotl/deepspeed_configs/zero2.json
weight_decay: 0.0
special_tokens:

訓練圖表

🔧 技術細節

本模型基於ai21labs/Jamba-v0.1進行微調，使用了QLoRA方法，並在A100 (80GB) x 4環境中進行了3小時的訓練。訓練過程中使用了兩個開源數據集，一個是英語數據集，另一個是多語言數據集，以提高模型的多語言能力。

📄 許可證

本模型使用Apache 2.0許可證。

👨‍💻 開發者

• 主要開發者 - Peter Devine ptrdvn
• 行政主管 - Shunichi Taniguchi shun1taniguchi