Model Overview

This model was trained as a small-scale experiment to determine how easy it is to fine-tune ai21labs/Jamba-v0.1 to work as a chatbot.

The aim of this experiment was to find how intelligently and reliably Jamba can chat in both English and other languages if only QLoRA finetuned for a few hours.

Initial subjective testing has shown that this model can chat reasonably well in both English and other languages, so feel free to give it a try!

Model Details

Model type: Joint Attention and Mamba (Jamba)
License: Apache 2.0
Context length: 256K
Knowledge cutoff date: March 5, 2024

Presequities

Jamba requires you use transformers version 4.39.0 or higher:

pip install transformers>=4.39.0

In order to run optimized Mamba implementations, you first need to install mamba-ssm and causal-conv1d:

pip install mamba-ssm causal-conv1d>=1.2.0

You also have to have the model on a CUDA device.

You can run the model not using the optimized Mamba kernels, but it is not recommended as it will result in significantly lower latencies. In order to do that, you'll need to specify use_mamba_kernels=False when loading the model.

How to use

※ - This code automatically appends the "<|startoftext|>" special token to any input. Appending this to all inputs is required for inference, as initial testing shows that leaving it out leads to output errors.

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|>

Initial testing results

Overall, the model seems to have decent knowledge in some areas while hallucinating incorrect information in others. It also seems faily easy to steer the model using the system message, being able to rhyme and modulate English difficulty to varying levels of success. The multi-turn accuracy seems decent from very little testing. The multilingual tests show us that this model fairly reliably replies in whatever language it has been prompted in.

Single turn examples

<|startoftext|><|im_start|>system 
You are a helpful AI assistant.
<|im_end|> 
<|im_start|>user
Describe the town of Strathaven, Scotland in 50 words.
<|im_end|> 
<|im_start|>assistant
Strathaven is a picturesque market town located in South Lanarkshire, Scotland. Known for its rich history and stunning landscapes, it offers visitors a delightful blend of historical attractions, outdoor activities, and modern amenities. With its charming streets, beautiful parks, and friendly community, Strathaven provides an ideal destination for those seeking a peaceful escape from city life.<|im_end|>
<|endoftext|>



<|startoftext|><|im_start|>system 
You are a helpful AI assistant.
<|im_end|> 
<|im_start|>user
Describe a harmless prank that I can play on my colleagues.
<|im_end|> 
<|im_start|>assistant
A harmless prank you can play on your colleagues is the "Fake Meeting" prank. Here's how to do it:

1. Choose a colleague who often has meetings or conference calls scheduled throughout the day.
2. Find out their schedule for the day, including any upcoming meetings they have.
3. Create a fake meeting invitation using an email template or online tool like Canva or Google Docs. Make sure the design and format look similar to the company's official meeting invitations.
4. Include details such as the date, time, location (if applicable), subject, and attendees. You can even add a fake agenda if you want to make it more convincing.
5. Send the fake meeting invitation to your chosen colleague via email or instant messaging platform used in your workplace.
6. Wait for them to respond with confusion or excitement about attending this unexpected meeting.
7. When they ask questions about the meeting, act clueless and say something like, "Oh, I just got this invite too! Must be some new initiative from upper management."
8. Keep up the charade until they realize it's a prank, but don't let it go on for too long so as not to cause unnecessary stress or confusion.





<|startoftext|><|im_start|>system 
You are a helpful AI assistant that always answers correctly.
<|im_end|> 
<|im_start|>user
Do UK citizens need a visa to holiday in Singapore?
<|im_end|> 
<|im_start|>assistant
Yes, UK citizens do need a visa to holiday in Singapore. They can apply for an Electronic Travel Authorization (ETA) online before their trip.<|im_end|>
<|endoftext|>

## <dev note>: This is false, UK citizens do NOT require a visa to holiday in Singapore. An Electronic Travel Authorization (ETA) is required for the USA etc.


<|startoftext|><|im_start|>system 
You are a helpful AI assistant.
<|im_end|> 
<|im_start|>user
Which US state has the highest GDP per capita?
<|im_end|> 
<|im_start|>assistant
As of my last update, the US state with the highest GDP per capita is Massachusetts.<|im_end|>
<|endoftext|>



<|startoftext|><|im_start|>system 
You are a helpful AI assistant.
<|im_end|> 
<|im_start|>user
Name three things that feel like a golf ball in your hands.
<|im_end|> 
<|im_start|>assistant
1. A tennis ball: Tennis balls have a similar size and texture to golf balls, making them feel quite similar when held in the hand. They both have a rubbery surface with dimples or indentations on their outer layer.

2. A stress ball: Stress balls are often made of soft, squishy material such as foam or gel, which can be squeezed and manipulated easily. The feeling of holding a stress ball is comparable to that of a golf ball due to its small size and round shape.

3. A ping pong ball: Ping pong balls are lightweight and hollow, much like golf balls. Although they may not have the same texture, their size and spherical shape make them feel somewhat similar when held in the hand.<|im_end|>
<|endoftext|>

System message control examples

<|startoftext|><|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
Sausages are better than bacon because they have more variety and can be eaten for breakfast, lunch, or dinner. They come in different flavors and sizes, making them versatile. Sausages also contain less fat and calories compared to bacon, which is healthier for people who want to watch their diet.<|im_end|>
<|endoftext|>



<|startoftext|><|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
Sausages, due to their versatility and adaptability, offer more culinary possibilities compared to bacon. They can be used as an ingredient or served whole, while bacon is often limited to being a side dish. Sausages also have a longer shelf life and come in various flavors, making them a more convenient choice for cooking.<|im_end|>
<|endoftext|>




<|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
The sky's hue, a question so true,
Is caused by light scattering through.
When sunlight meets our atmosphere,
It scatters and creates a fair weather affair.

Blue light waves have shorter lengths,
They scatter more than other tints.
This phenomenon, called Rayleigh scattering,
Makes the sky appear as if it were wearing a sapphire ring.

So when you gaze at the azure expanse above,
Remember this simple explanation of nature's love.<|im_end|>
<|endoftext|>

Multi-turn examples

<|startoftext|><|startoftext|><|im_start|>system 
You are a helpful AI assistant.
<|im_end|> 
<|im_start|>user
Give an example of a healthy breakfast.
<|im_end|> 
<|im_start|>assistant
A healthy breakfast could include:

1. A bowl of oatmeal made with water or low-fat milk, topped with fresh berries and a sprinkle of nuts or seeds (such as almonds, walnuts, chia seeds, or flaxseeds).
2. A slice of whole grain toast spread with natural peanut butter or almond butter, along with a side of sliced fruit like apples or bananas.
3. Greek yogurt mixed with honey and granola, served with a handful of blueberries or strawberries on the side.
4. A vegetable omelette made with egg whites or a combination of whole eggs and egg whites, filled with spinach, tomatoes, onions, and mushrooms, paired with a small serving of avocado slices.
5. A smoothie made with unsweetened almond milk, frozen fruits (like mangoes, pineapples, or berries), a scoop of protein powder, and a handful of leafy greens (like spinach or kale).
6. Whole grain cereal with skim milk, accompanied by a piece of fruit and a dollop of plain Greek yogurt for added protein.<|im_end|>
<|im_start|>user
Which of these would have the lowest calorie count while still making me feel full?
<|im_end|> 
<|im_start|>assistant
The option with the lowest calorie count that will also make you feel full is Option 1: Oatmeal with fresh berries and nuts/seeds. Oatmeal is high in fiber, which helps keep you feeling full longer, and it's lower in calories compared to other options like smoothies and omelettes. The addition of fresh berries provides essential vitamins and minerals, while the nuts and seeds add healthy fats and proteins to help you stay satisfied until your next meal.<|im_end|>
<|endoftext|>

Multilingual examples

<|startoftext|><|im_start|>system 
You are a helpful AI assistant.
<|im_end|> 
<|im_start|>user
東京駅から池袋駅まで丸の内線で行きたいですが、普通に何時間かかりますか？
<|im_end|> 
<|im_start|>assistant
東京駅から池袋駅まで丸の内線で行く場合、平均的には約20分程度かかります。しかし、これは一般的な情報であり、実際の旅行時間は運行状況や列車の種類、日付、時間帯などによって異なることに注意してください。<|im_end|>
<|endoftext|>


<|startoftext|><|im_start|>system 
You are a helpful AI assistant.
<|im_end|> 
<|im_start|>user
Opisz kilka ciekawych miejsc do odwiedzenia w Londynie
<|im_end|> 
<|im_start|>assistant
1. Buckingham Palace: Jest to oficjalne siedzisko Królowej Brytanii, króla Brytanii i innych państw Członkowskich Zjednoczonego Królestwa. Można tu zobaczyć mający historię palac oraz obejrzeć ceremonię zmiany strażników na placu przed budynkiem.

2. Big Ben: To jest jedna z najbardziej rozpoznawalnych atrakcji turystycznych w Londynie. Jest to wieża zegarowa położona na terenie Parlamentu Wielkiej Brytanii i Senatu. Jej dźwiękowy sygnał można usłyszeć nawet na dalekie odległości.

3. Tower Bridge: Ten piękny most nad Tamizą oferuje widok na skyline Londynu. Otwiera się on pod koniec każdego d

Training details

The model was trained on 2 open source datasets (one multilingual) for one epoch on a A100 (80GB) x 4 environment for 3 hours.

Training data

jondurbin/airoboros-3.2

A ~59K example dataset of curated LLM tasks in English, primarily generated with GPT-4. This dataset has been used by some of the best performing open source LLMs in the world (e.g. jondurbin/bagel-7b-v0.4, NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO) and contains a wide variety of tasks, so we hypothesized that this would lead to a multi-talented, accurate model. For this reason we chose this dataset was chosen for the bulk of our training data.

Note: Each element in jondurbin/airoboros-3.2 already contains a system message.

openchat/openchat_sharegpt4_dataset (GPT-4 responses only)

A ~6K example dataset of multilingual multi-turn chats between users and GPT-4. While jondurbin/airoboros-3.2 has deilvered good results for models previously, it sadly contains no (or seemingly very little) multilingual data. We are a Japanese AI company, so require an LLM to be able to output in Japanese too. Hence we also selected a small, seemingly high quality dataset of GPT-4 responses in many languages from the ShareGPT dataset. We chose to only select the GPT-4 responses as we wanted to keep our dataset as small and high quality as possible to maximise the efficiency of our training.

Note: openchat/openchat_sharegpt4_dataset does not contain system messages, so we added 'You are GPT-4, a helpful assistant.' as our system message.

Data preparation code

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")

Training

The Jamba-v0.1 base model was trained for roughly 3 hours in a A100 (80GB) x 4 environment on the Azure cloud (Standard_NC96ads_A100_v4).

We trained using QLoRA and merged the adapter to the original weights.

Our training harness was Axolotl using the ChatML chat template. Full details of the training config are below:

Training config

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: