--- library_name: transformers license: apache-2.0 base_model: ai21labs/Jamba-v0.1 datasets: - jondurbin/airoboros-3.2 - openchat/openchat_sharegpt4_dataset pipeline_tag: text-generation --- # Model Overview This model was trained as a small-scale experiment to determine how easy it is to fine-tune [ai21labs/Jamba-v0.1](https://huggingface.co/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: ```bash pip install transformers>=4.39.0 ``` In order to run optimized Mamba implementations, you first need to install `mamba-ssm` and `causal-conv1d`: ```bash 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. ```python 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|> ## : 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](https://huggingface.co/datasets/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](https://huggingface.co/jondurbin/bagel-7b-v0.4), [NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO](https://huggingface.co/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](https://huggingface.co/datasets/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 ```python 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 ```yaml 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: ```
[Built with Axolotl](https://github.com/OpenAccess-AI-Collective/axolotl)
Training graphs ![image/png](https://cdn-uploads.huggingface.co/production/uploads/64b63f8ad57e02621dc93c8b/umxTIsNRHUtKS_kL81Uyf.png) ![image/png](https://cdn-uploads.huggingface.co/production/uploads/64b63f8ad57e02621dc93c8b/mpuCoL99rxX8RCgXH1CJo.png) ![image/png](https://cdn-uploads.huggingface.co/production/uploads/64b63f8ad57e02621dc93c8b/5FvwYNdte-bgzEvcvFO8I.png)

# Developers Lead developer - Peter Devine [ptrdvn](https://huggingface.co/ptrdvn) Administrative supervisor - Shunichi Taniguchi [shun1taniguchi](https://huggingface.co/shun1taniguchi)