xbe.at

LLMs and generative AI / Fine tuning and alignment

Fine-Tuning MISTRAL AI Model with Python

5 min read, 15 sections

Contents
  1. Fine-Tuning MISTRAL AI Model
  2. Understanding MISTRAL AI
  3. Preparing the Dataset
  4. Tokenization and Encoding
  5. Defining the Training Arguments
  6. Creating a Custom Dataset
  7. Defining the Model for Fine-Tuning
  8. Setting Up the Trainer
  9. Fine-Tuning the Model
  10. Evaluating the Fine-Tuned Model
  11. Saving the Fine-Tuned Model
  12. Using the Fine-Tuned Model
  13. Real-Life Example: Customer Review Analysis
  14. Real-Life Example: Social Media Sentiment Analysis
  15. Additional Resources

Fine-Tuning MISTRAL AI Model

Fine-tuning the MISTRAL AI model involves adapting a pre-trained model to specific tasks or domains. This process enhances the model's performance on targeted applications while leveraging the knowledge acquired during pre-training. We'll explore the steps, techniques, and best practices for fine-tuning MISTRAL using Python.

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model_name = "mistralai/Mistral-7B-v0.1"
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

print(f"Loaded {model_name} model and tokenizer")

Understanding MISTRAL AI

MISTRAL AI is a powerful language model developed by Mistral AI. It's designed to generate human-like text and can be fine-tuned for various natural language processing tasks. The model uses a transformer architecture and has been pre-trained on a diverse corpus of text data.

# Exploring model architecture
print(f"Model architecture: {model.__class__.__name__}")
print(f"Number of parameters: {model.num_parameters():,}")
print(f"Vocabulary size: {len(tokenizer)}")

Preparing the Dataset

Before fine-tuning, we need to prepare our dataset. This involves collecting and preprocessing the data specific to our task. For this example, we'll use a sentiment analysis dataset.

from datasets import load_dataset

dataset = load_dataset("imdb")
train_data = dataset["train"]
test_data = dataset["test"]

print(f"Training samples: {len(train_data)}")
print(f"Testing samples: {len(test_data)}")
print(f"Sample text: {train_data[0]['text'][:100]}...")
print(f"Sample label: {train_data[0]['label']}")

Tokenization and Encoding

To prepare our data for the model, we need to tokenize and encode it. This process converts text into numerical representations that the model can understand.

def tokenize_function(examples):
    return tokenizer(examples["text"], padding="max_length", truncation=True, max_length=512)

tokenized_train = train_data.map(tokenize_function, batched=True)
tokenized_test = test_data.map(tokenize_function, batched=True)

print(f"Encoded input_ids shape: {tokenized_train[0]['input_ids'][:10]}")
print(f"Encoded attention_mask shape: {tokenized_train[0]['attention_mask'][:10]}")

Defining the Training Arguments

We need to set up the training arguments that control various aspects of the fine-tuning process, such as learning rate, batch size, and number of epochs.

from transformers import TrainingArguments

training_args = TrainingArguments(
    output_dir="./results",
    num_train_epochs=3,
    per_device_train_batch_size=8,
    per_device_eval_batch_size=8,
    warmup_steps=500,
    weight_decay=0.01,
    logging_dir="./logs",
    logging_steps=10,
)

print(f"Learning rate: {training_args.learning_rate}")
print(f"Batch size: {training_args.per_device_train_batch_size}")
print(f"Number of epochs: {training_args.num_train_epochs}")

Creating a Custom Dataset

To fine-tune MISTRAL for our specific task, we need to create a custom dataset that combines the input text with the target labels.

import torch
from torch.utils.data import Dataset

class SentimentDataset(Dataset):
    def __init__(self, encodings, labels):
        self.encodings = encodings
        self.labels = labels

    def __getitem__(self, idx):
        item = {key: torch.tensor(val[idx]) for key, val in self.encodings.items()}
        item['labels'] = torch.tensor(self.labels[idx])
        return item

    def __len__(self):
        return len(self.labels)

train_dataset = SentimentDataset(tokenized_train, train_data["label"])
test_dataset = SentimentDataset(tokenized_test, test_data["label"])

print(f"Training dataset size: {len(train_dataset)}")
print(f"Testing dataset size: {len(test_dataset)}")

Defining the Model for Fine-Tuning

We need to adapt the MISTRAL model for our specific task. In this case, we'll add a classification head for sentiment analysis.

from transformers import AutoModelForSequenceClassification

num_labels = 2  # Binary classification for sentiment analysis
model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=num_labels)

print(f"Model architecture: {model.__class__.__name__}")
print(f"Number of parameters: {model.num_parameters():,}")

Setting Up the Trainer

The Trainer class from the transformers library simplifies the fine-tuning process by handling the training loop, evaluation, and logging.

from transformers import Trainer

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=test_dataset,
)

print("Trainer initialized with the following components:")
print(f"Model: {trainer.model.__class__.__name__}")
print(f"Optimizer: {trainer.optimizer.__class__.__name__}")
print(f"Scheduler: {trainer.lr_scheduler.__class__.__name__}")

Fine-Tuning the Model

Now we can start the fine-tuning process. This will update the model's parameters to improve its performance on our specific task.

print("Starting fine-tuning...")
trainer.train()

print("Fine-tuning completed")
print(f"Total training steps: {trainer.state.global_step}")
print(f"Final training loss: {trainer.state.log_history[-1]['loss']:.4f}")

Evaluating the Fine-Tuned Model

After fine-tuning, we need to evaluate the model's performance on our test dataset to see how well it has learned the task.

eval_results = trainer.evaluate()

print("Evaluation results:")
print(f"Loss: {eval_results['eval_loss']:.4f}")
print(f"Accuracy: {eval_results['eval_accuracy']:.4f}")

Saving the Fine-Tuned Model

To use our fine-tuned model in the future, we need to save it along with its configuration and tokenizer.

output_dir = "./fine_tuned_mistral"
model.save_pretrained(output_dir)
tokenizer.save_pretrained(output_dir)

print(f"Model saved to: {output_dir}")
print("Saved files:")
print("\n".join(os.listdir(output_dir)))

Using the Fine-Tuned Model

Now that we have a fine-tuned model, we can use it to make predictions on new, unseen data.

from transformers import pipeline

classifier = pipeline("sentiment-analysis", model=output_dir, tokenizer=output_dir)

texts = [
    "I absolutely loved this movie! The acting was superb.",
    "This book was a complete waste of time. Terrible plot and characters.",
]

results = classifier(texts)
for text, result in zip(texts, results):
    print(f"Text: {text}")
    print(f"Sentiment: {result['label']}, Score: {result['score']:.4f}\n")

Real-Life Example: Customer Review Analysis

Let's use our fine-tuned model to analyze customer reviews for a product.

customer_reviews = [
    "The product arrived on time and works perfectly. Great purchase!",
    "Disappointing quality. It broke after just a week of use.",
    "Average product, nothing special but gets the job done.",
    "Excellent customer service! They quickly resolved my issue.",
]

results = classifier(customer_reviews)
for review, result in zip(customer_reviews, results):
    sentiment = "Positive" if result['label'] == "LABEL_1" else "Negative"
    print(f"Review: {review}")
    print(f"Sentiment: {sentiment}, Confidence: {result['score']:.4f}\n")

Real-Life Example: Social Media Sentiment Analysis

Another application of our fine-tuned model is analyzing sentiment in social media posts.

social_media_posts = [
    "Just tried the new restaurant downtown. The food was amazing! #YUM",
    "Stuck in traffic for 2 hours. This city's infrastructure is a joke. #Frustrated",
    "Can't believe how beautiful the sunset is tonight. Nature is incredible!",
    "Another day, another data breach. When will companies take our privacy seriously?",
]

results = classifier(social_media_posts)
for post, result in zip(social_media_posts, results):
    sentiment = "Positive" if result['label'] == "LABEL_1" else "Negative"
    print(f"Post: {post}")
    print(f"Sentiment: {sentiment}, Confidence: {result['score']:.4f}\n")

Additional Resources

For more information on fine-tuning language models and the MISTRAL AI model:

  1. Hugging Face Transformers Documentation: https://huggingface.co/docs/transformers/index
  2. "Fine-tuning large language models: Techniques and applications" (arXiv:2303.08804): https://arxiv.org/abs/2303.08804
  3. "A Survey of Large Language Models" (arXiv:2303.18223): https://arxiv.org/abs/2303.18223

More in Fine tuning and alignment

  1. Introduction to Fine-tuning Large Language Models (LLMs) Using Python7 min

    Introduction to Fine-tuning Large Language Models (LLMs)

  2. Fine-Tuning LLMs with LoRA in Python7 min

    Fine-tuning Large Language Models (LLMs) is a crucial technique for adapting pre-trained models to specific tasks or domains. Low-Rank Adaptation (LoRA) is...

  3. Knowledge Distillation in Large Language Models13 min

    Knowledge distillation is a technique used to transfer knowledge from a large, complex model (teacher) to a smaller, more efficient model (student). This...

  4. Fine-Tuning Large Language Models for Business Solutions6 min

    Mixed precision training combines different numerical precisions to optimize memory usage and computational speed while maintaining model accuracy. This...

  5. Limitations of LoRA for Pre-training Large Language Models8 min

    Why LoRA Isn't Suitable for Pre-training LLMs

  6. Entropic Distribution Matching for Supervised LLM Fine-tuning13 min

    Entropic Distribution Matching (EDM) is a novel approach in supervised fine-tuning of Large Language Models (LLMs) that aims to reduce overfitting and...