How to tune Mistral AI 7B LLM with Hugging Face AutoTrain

With the progress of LLM research around the world, many models have become more accessible. One of the small but powerful open source models is ai/” rel=”noopener” target=”_blank”>Mistral ai 7B LLM. The model has adaptability in many use cases, showing better performance than LlaMA 2 13B in all benchmarks, using a sliding window attention (SWA) mechanism and be easy to implement.

The overall performance benchmark of the Mistral 7 B can be seen in the image below.

<img decoding="async" alt="How to tune Mistral ai 7B LLM with Hugging Face AutoTrain” width=”100%” src=”https://technicalterrence.com/wp-content/uploads/2023/11/1699797826_316_How-to-tune-Mistral-AI-7B-LLM-with-Hugging-Face.png”/><img decoding="async" src="https://technicalterrence.com/wp-content/uploads/2023/11/1699797826_316_How-to-tune-Mistral-AI-7B-LLM-with-Hugging-Face.png" alt="How to tune Mistral ai 7B LLM with Hugging Face AutoTrain” width=”100%”/>
Mistral 7B performance benchmark (Jiang et al. (2023))

The Mistral 7B model is available in the HugsFace also. With this, we can use Hugging Face AutoTrain to tune the model for our use cases. Embracing Face’s AutoTrain is a no-code platform with Python API that we can use to easily fine-tune any LLM model available on HugginFace.

This tutorial will teach us how to adjust Mistral ai 7B LLM with Hugging Face AutoTrain. How does it work? Let’s get into it.

To wrap the LLM with the Python API, we need to install the Python package, which you can run using the following code.

pip install -U autotrain-advanced

Additionally, we would use the Alpaca sample data set from HugsFacewhich required the Datasets package to acquire and the Transformers package to manipulate the Hugging Face model.

pip install datasets transformers

Next, we need to format our data to fit the Mistral 7B model. In general, there are two fundamental models that Mistral launched: ai/llm/mistral-v0.1″ rel=”noopener” target=”_blank”>Mistral 7B v0.1 and ai/llm/mistral-instruct-v0.1″ rel=”noopener” target=”_blank”>Mistral 7B Instruction v0.1. The Mistral 7B v0.1 is the basic model and the Mistral 7B Instruct v0.1 is a Mistral 7B v0.1 model that has been refined for conversation and answering questions.

We would need a CSV file containing a column of text to perform the adjustment with Hugging Face AutoTrain. However, we would use a different text format for the base and instructional models during adjustment.

First, let’s look at the data set we used for our sample.

from datasets import load_dataset
import pandas as pd

# Load the dataset
train= load_dataset("tatsu-lab/alpaca",split="train(:10%)")
train = pd.DataFrame(train)

The code above would sample ten percent of the actual data. We would only need that amount for this tutorial since it would take longer to train for larger data. Our data sample looks like the image below.

<img decoding="async" alt="How to tune Mistral ai 7B LLM with Hugging Face AutoTrain” width=”100%” src=”https://technicalterrence.com/wp-content/uploads/2023/11/1699797826_808_How-to-tune-Mistral-AI-7B-LLM-with-Hugging-Face.png”/><img decoding="async" src="https://technicalterrence.com/wp-content/uploads/2023/11/1699797826_808_How-to-tune-Mistral-AI-7B-LLM-with-Hugging-Face.png" alt="How to tune Mistral ai 7B LLM with Hugging Face AutoTrain” width=”100%”/>
Image by author

The data set already contains the formatted text columns we need to fit our LLM model. That’s why we don’t need to do anything. However, I would provide you with code if you have another data set that needs formatting.

def text_formatting(data):

    # If the input column is not empty
    if data('input'):

        text = f"""Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.\n\n### Instruction:\n{data("instruction")} \n\n### Input:\n{data("input")}\n\n### Response:\n{data("output")}"""

    else:

        text = f"""Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n### Instruction:\n{data("instruction")}\n\n### Response:\n{data("output")}""" 

    return text

train('text') = train.apply(text_formatting, axis =1)

For Hugging Face AutoTrain, we would need the data in CSV format so we can save it with the following code.

train.to_csv('train.csv', index = False)

Then move the CSV result to a folder called data. That’s all you need to prepare the data set for tuning Mistral 7B v0.1.

If you want to adjust Mistral 7B Instruct v0.1 to chat and answer questions, we need to follow the chat template format provided by Mistral, shown in the code block below.

<s>(INST) Instruction (/INST) Model answer</s>(INST) Follow-up instruction (/INST)

If we use our example data set above, we need to reformat the text column. We would use just the data without any input to the chat model.

train_chat = train(train('input') == '').reset_index(drop = True).copy()

Then, we could reformat the data with the following code.

def chat_formatting(data):

  text = f"<s>(INST) {data('instruction')} (/INST) {data('output')} </s>"

  return text

train_chat('text') = train_chat.apply(chat_formatting, axis =1)
train_chat.to_csv('train_chat.csv', index =False)

We will end up with an appropriate data set to fit the Mistral 7B Instruct v0.1 model.

<img decoding="async" alt="How to tune Mistral ai 7B LLM with Hugging Face AutoTrain” width=”100%” src=”https://technicalterrence.com/wp-content/uploads/2023/11/1699797826_391_How-to-tune-Mistral-AI-7B-LLM-with-Hugging-Face.png”/><img decoding="async" src="https://technicalterrence.com/wp-content/uploads/2023/11/1699797826_391_How-to-tune-Mistral-AI-7B-LLM-with-Hugging-Face.png" alt="How to tune Mistral ai 7B LLM with Hugging Face AutoTrain” width=”100%”/>
Image by author

With all the preparation set up, we can now start the AutoTrain to fine-tune our Mistral model.

Let’s configure the Hugging Face AutoTrain environment to tune the Mistral model. First, let’s run the AutoTrain setup using the following command.

Next, we would provide the information necessary for AutoTrain to run. For this tutorial, we will be using Mistral 7B Instruct v0.1.

project_name="my_autotrain_llm"
model_name="mistralai/Mistral-7B-Instruct-v0.1"

Then we would add the Hugging Face information if you want to submit your model to the repository.

push_to_hub = False
hf_token = "YOUR HF TOKEN"
repo_id = "username/repo_name"

Finally, we would start the information of the model parameters in the following variables. You can change them to see if the result is good.

learning_rate = 2e-4
num_epochs = 4
batch_size = 1
block_size = 1024
trainer = "sft"
warmup_ratio = 0.1
weight_decay = 0.01
gradient_accumulation = 4
use_fp16 = True
use_peft = True
use_int4 = True
lora_r = 16
lora_alpha = 32
lora_dropout = 0.045

We can modify many parameters, but we will not discuss them in this article. Some tips to improve LLM tuning include using a lower learning rate to maintain previously learned representations and vice versa, avoiding overfitting by adjusting the number of epochs, using a larger batch size for greater stability, or adjusting gradient accumulation. if you have a memory problem.

When all the information is ready, we will configure the environment to accept all the information that we have previously configured.

import os
os.environ("PROJECT_NAME") = project_name
os.environ("MODEL_NAME") = model_name
os.environ("PUSH_TO_HUB") = str(push_to_hub)
os.environ("HF_TOKEN") = hf_token
os.environ("REPO_ID") = repo_id
os.environ("LEARNING_RATE") = str(learning_rate)
os.environ("NUM_EPOCHS") = str(num_epochs)
os.environ("BATCH_SIZE") = str(batch_size)
os.environ("BLOCK_SIZE") = str(block_size)
os.environ("WARMUP_RATIO") = str(warmup_ratio)
os.environ("WEIGHT_DECAY") = str(weight_decay)
os.environ("GRADIENT_ACCUMULATION") = str(gradient_accumulation)
os.environ("USE_FP16") = str(use_fp16)
os.environ("USE_PEFT") = str(use_peft)
os.environ("USE_INT4") = str(use_int4)
os.environ("LORA_R") = str(lora_r)
os.environ("LORA_ALPHA") = str(lora_alpha)
os.environ("LORA_DROPOUT") = str(lora_dropout)

We would use the following command to run AutoTrain in our notebook.

!autotrain llm \
--train \
--model ${MODEL_NAME} \
--project-name ${PROJECT_NAME} \
--data-path data/ \
--text-column text \
--lr ${LEARNING_RATE} \
--batch-size ${BATCH_SIZE} \
--epochs ${NUM_EPOCHS} \
--block-size ${BLOCK_SIZE} \
--warmup-ratio ${WARMUP_RATIO} \
--lora-r ${LORA_R} \
--lora-alpha ${LORA_ALPHA} \
--lora-dropout ${LORA_DROPOUT} \
--weight-decay ${WEIGHT_DECAY} \
--gradient-accumulation ${GRADIENT_ACCUMULATION} \
$( (( "$USE_FP16" == "True" )) && echo "--fp16" ) \
$( (( "$USE_PEFT" == "True" )) && echo "--use-peft" ) \
$( (( "$USE_INT4" == "True" )) && echo "--use-int4" ) \
$( (( "$PUSH_TO_HUB" == "True" )) && echo "--push-to-hub --token ${HF_TOKEN} --repo-id ${REPO_ID}" )

If the adjustment process is successful, we will have a new directory of our adjusted model. We would use this directory to test our newly tuned model.

from transformers import AutoModelForCausalLM, AutoTokenizer

model_path = "my_autotrain_llm"
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForCausalLM.from_pretrained(model_path)

With the model and tokenizer ready to use, we would test the model with an input example.

input_text = "Give three tips for staying healthy."
input_ids = tokenizer.encode(input_text, return_tensors="pt")
output = model.generate(input_ids, max_new_tokens = 200)
predicted_text = tokenizer.decode(output(0), skip_special_tokens=True)
print(predicted_text)

Production:

Give three tips to stay healthy.

1. Eat a balanced diet: Be sure to include plenty of fruits, vegetables, lean proteins, and whole grains in your diet. This will help you get the nutrients you need to stay healthy and energetic.
2. Exercise regularly: Try to get at least 30 minutes of moderate exercise, such as brisk walking or cycling, every day. This will help you maintain a healthy weight, reduce your risk of chronic diseases, and improve your overall physical and mental health.
3. Get enough sleep: Aim for 7 to 9 hours of quality sleep each night. This will help you feel more rested and alert during the day, and will also help you maintain a healthy weight and reduce your risk of chronic diseases.

The model result has been close to the actual result of our training data, as shown in the image below.

1. Eat a balanced diet and be sure to include plenty of fruits and vegetables.
2. Exercise regularly to keep your body active and strong.
3. Get enough sleep and maintain a consistent sleep schedule.

The Mistral models are certainly powerful for their size, as a simple fine-tuning has already yielded promising results. Test your data set to see if it suits your work.

The Mistral ai 7B family model is a powerful LLM model that boasts superior performance than LLaMA and great adaptability. Since the model is available in Hugging Face, we can employ HuggingFace AutoTrain to tune the model. There are currently two models available for tuning at Hugging Face; Mistral 7B v0.1 for the base foundation model and Mistral 7B Instruct v0.1 for chatting and answering questions. Fine tuning yielded promising results even with a rapid training process.

Cornellius Yudha Wijaya He is an assistant data science manager and data writer. While working full-time at Allianz Indonesia, she loves sharing Python tips and data through social media and print media.