Fine-Tuning Mistral Small 3.1 for Emotion Recognition in Social Media

In the evolving landscape of Natural Language Processing, the shift from massive models to Small Language Models (SLMs) has gained significant momentum. While Large Language Models (LLMs) like GPT-4 or Claude 3.5 Sonnet—accessible via n1n.ai—are incredibly capable, they are often overkill for specific classification tasks. For developers and enterprises, fine-tuning an SLM like Mistral Small 3.1 offers a cost-effective, high-speed alternative for specialized tasks like emotion recognition in social media text. This tutorial explores how to fine-tune Mistral Small 3.1 on an imbalanced dataset containing 15 distinct emotional labels.

Why Mistral Small 3.1 for Emotion Recognition?

Mistral Small 3.1 strikes a perfect balance between parameter count and reasoning capability. In social media communication, emotions are often nuanced, sarcastic, or layered with slang. A generic model might identify 'anger,' but a fine-tuned SLM can distinguish between 'frustration,' 'outrage,' and 'annoyance.' By using an aggregator like n1n.ai, developers can benchmark their fine-tuned local models against industry leaders to ensure performance parity.

The Challenge: Imbalanced Data and 15 Labels

Real-world social media data is notoriously imbalanced. In a 15-emotion dataset (e.g., Joy, Sadness, Anger, Fear, Surprise, Disgust, Trust, Anticipation, Love, Optimism, Pessimism, etc.), 'Joy' and 'Anger' might dominate, while 'Anticipation' remains a minority class. If we train without addressing this, the model will develop a bias toward the majority classes.

Step 1: Environment Setup and Dependencies

To begin, you will need a Python environment with CUDA support. We will use the Hugging Face transformers library, peft for Parameter-Efficient Fine-Tuning, and bitsandbytes for quantization.

# Install necessary libraries
!pip install -q -U transformers peft bitsandbytes datasets accelerate

Step 2: Data Preprocessing for Imbalance

When dealing with 15 emotions, we must calculate class weights to penalize the model more for misclassifying minority classes. This is crucial for maintaining high F1-scores across all categories.

import pandas as pd
import numpy as np
from sklearn.utils.class_weight import compute_class_weight

# Load your dataset
df = pd.read_csv("social_media_emotions.csv")
labels = df['emotion_id'].values

# Calculate weights
weights = compute_class_weight(
    class_weight='balanced',
    classes=np.unique(labels),
    y=labels
)
class_weights = torch.tensor(weights, dtype=torch.float)

Step 3: Model Loading with 4-bit Quantization

To run Mistral Small 3.1 on consumer-grade GPUs, we use QLoRA (Quantized Low-Rank Adaptation). This reduces memory usage significantly without a major hit to accuracy.

from transformers import AutoModelForSequenceClassification, AutoTokenizer, BitsAndBytesConfig
import torch

model_id = "mistralai/Mistral-Small-v3.1"

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForSequenceClassification.from_pretrained(
    model_id,
    num_labels=15,
    quantization_config=bnb_config,
    device_map="auto"
)

Step 4: Configuring LoRA Adapters

LoRA allows us to train only a tiny fraction of the model's parameters. For Mistral Small, we target the query, key, and value projections in the attention layers.

from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training

model = prepare_model_for_kbit_training(model)

lora_config = LoraConfig(
    r=16,
    lora_alpha=32,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
    lora_dropout=0.05,
    bias="none",
    task_type="SEQ_CLS"
)

model = get_peft_model(model, lora_config)

Step 5: Custom Loss Function for Class Imbalance

Standard Trainer uses CrossEntropyLoss. To handle our 15 imbalanced emotions, we override the compute_loss method in a custom Trainer class.

from transformers import Trainer
import torch.nn as nn

class WeightedTrainer(Trainer):
    def compute_loss(self, model, inputs, return_outputs=False):
        labels = inputs.pop("labels")
        outputs = model(**inputs)
        logits = outputs.get("logits")
        loss_fct = nn.CrossEntropyLoss(weight=class_weights.to(model.device))
        loss = loss_fct(logits.view(-1, self.model.config.num_labels), labels.view(-1))
        return (loss, outputs) if return_outputs else loss

Step 6: Training and Evaluation

Set your training arguments. Focus on a low learning rate (e.g., 2e-4) and use a cosine learning rate scheduler to prevent catastrophic forgetting.

from transformers import TrainingArguments

training_args = TrainingArguments(
    output_dir="./mistral-emotion-results",
    learning_rate=2e-4,
    per_device_train_batch_size=8,
    num_train_epochs=3,
    weight_decay=0.01,
    evaluation_strategy="epoch",
    save_strategy="epoch",
    load_best_model_at_end=True,
    fp16=True,
    logging_steps=10,
)

trainer = WeightedTrainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_train,
    eval_dataset=tokenized_val,
    tokenizer=tokenizer,
)

trainer.train()

Pro Tips for High-Accuracy Emotion Detection

1. Data Augmentation: For the rarest emotions, use an LLM API from n1n.ai to generate synthetic social media posts. This helps the model see more examples of minority classes.
2. Context Windows: Social media posts are short. Ensure your max_length is optimized (usually 128 or 256 tokens) to save memory.
3. F1-Score: Never rely on accuracy for imbalanced sets. Use the Macro-F1 score to evaluate how well the model performs on the 15th emotion as compared to the 1st.

Conclusion

Fine-tuning Mistral Small 3.1 transforms a general-purpose model into a specialized engine capable of navigating the complexities of human emotion. By addressing data imbalance with weighted loss functions and utilizing QLoRA for efficiency, developers can deploy powerful sentiment tools at a fraction of the cost of larger models. For those who need to scale their AI infrastructure or access multiple model providers through a single interface, n1n.ai provides the necessary API aggregation to streamline development.

Get a free API key at n1n.ai.