How to Fine-Tune Llama 3.2 Vision On a Custom Dataset?

# Install core dependencies
pip install unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git
pip install --no-deps "xformers<0.0.27" trl peft accelerate bitsandbytes
pip install datasets transformers torch torchvision pillow

# Set up environment
import torch
from unsloth import FastVisionModel
from datasets import load_dataset
import os
os.environ["WANDB_DISABLED"] = "true"  # Disable wandb logging

# Load Unsloth's 4-bit quantized Llama 3.2 Vision model
print("Loading Llama 3.2 Vision 11B model...")
model, tokenizer = FastVisionModel.from_pretrained(
    "unsloth/Llama-3.2-11B-Vision-Instruct-unsloth-bnb-4bit",
    load_in_4bit=True,
    use_gradient_checkpointing="unsloth"
)

print("Model loaded successfully!")
print(f"GPU memory used: {torch.cuda.memory_allocated()/1024**3:.2f} GB")

# Enable training mode
FastVisionModel.for_training(model)

# Configure LoRA (Low-Rank Adaptation) parameters for efficient fine-tuning
model = FastVisionModel.get_peft_model(
    model,
    finetune_vision_layers=True,      # Train vision components
    finetune_language_layers=True,    # Train language components  
    finetune_attention_modules=True,  # Train attention layers
    finetune_mlp_modules=True,        # Train MLP layers
    r=16,                            # LoRA rank
    lora_alpha=16,                   # LoRA scaling parameter
    lora_dropout=0,                  # No dropout for stability
    bias="none",                     # No bias training
    random_state=3407,               # For reproducibility
    use_rslora=False,                # Standard LoRA
    loftq_config=None,               # No quantization-aware training
)

print("Model configured for fine-tuning!")

# Load DocVQA dataset from Hugging Face
print("Loading DocVQA dataset...")
dataset = load_dataset("lmms-lab/DocVQA")

# Explore dataset structure
print("Dataset structure:")
print(f"Training samples: {len(dataset['train'])}")
print(f"Validation samples: {len(dataset['val'])}")

# Look at a sample to understand the data format
sample = dataset['train'][0]
print("\nSample data:")
print(f"Question: {sample['question']}")
print(f"Answers: {sample['answers']}")
print(f"Question types: {sample['question_types']}")

# Display image information
print(f"Image size: {sample['image'].size}")
print(f"Image mode: {sample['image'].mode}")

def format_docvqa_sample(sample):
    """Convert DocVQA samples to Llama 3.2 Vision chat format"""
    
    # Create conversation format that Llama expects
    conversation = [
        {
            "role": "user",
            "content": [
                {"type": "image"},
                {"type": "text", "text": sample['question']}
            ]
        },
        {
            "role": "assistant", 
            "content": sample['answers'][0]  # Use first answer as target
        }
    ]
    
    return {
        "messages": conversation,
        "image": sample['image']
    }

# Apply formatting to entire dataset
print("Formatting dataset for training...")
formatted_dataset = dataset.map(format_docvqa_sample)

# Create training and evaluation subsets
# Using smaller subsets for faster training - adjust based on your needs
train_dataset = formatted_dataset['train'].select(range(1000))  
eval_dataset = formatted_dataset['val'].select(range(200))

print(f"Training samples: {len(train_dataset)}")
print(f"Evaluation samples: {len(eval_dataset)}")

from transformers import TrainingArguments
from trl import SFTTrainer
from unsloth.trainer import UnslothVisionDataCollator

# Configure training parameters
training_args = TrainingArguments(
    # Core training settings
    per_device_train_batch_size=1,        # Small batch size for memory efficiency
    per_device_eval_batch_size=1,         # Small eval batch size
    gradient_accumulation_steps=8,        # Accumulate gradients for effective batch size of 8
    warmup_steps=10,                      # Warm up learning rate
    num_train_epochs=1,                   # One epoch for this demo
    max_steps=100,                        # Limit steps for demo - increase for full training
    
    # Learning rate and optimization
    learning_rate=2e-4,                   # Conservative learning rate
    optim="adamw_8bit",                   # Memory-efficient optimizer
    weight_decay=0.01,                    # Prevent overfitting
    lr_scheduler_type="linear",           # Linear learning rate decay
    
    # Precision settings for memory efficiency
    fp16=not torch.cuda.is_bf16_supported(),  # Use fp16 if bf16 not available
    bf16=torch.cuda.is_bf16_supported(),       # Use bf16 if available
    
    # Logging and evaluation
    logging_steps=5,                      # Log every 5 steps
    eval_strategy="steps",                # Evaluate during training
    eval_steps=25,                        # Evaluate every 25 steps
    save_steps=50,                        # Save checkpoint every 50 steps
    save_total_limit=2,                   # Keep only 2 checkpoints
    
    # Output and reproducibility
    output_dir="./llama-vision-docvqa",   # Output directory
    seed=3407,                            # Random seed
    data_seed=3407,                       # Data seed
    
    # Memory optimization
    dataloader_pin_memory=False,          # Disable for memory efficiency
    remove_unused_columns=False,          # Keep all columns
)

print("Training configuration set up successfully!")

def evaluate_model_performance(model, tokenizer, eval_samples, prefix=""):
    """Test model performance on DocVQA samples"""
    results = []
    
    print(f"\n{prefix}EVALUATION STARTING...")
    print("=" * 60)
    
    for i, sample in enumerate(eval_samples[:10]):  # Test on 10 samples
        # Prepare the input in chat format
        messages = [
            {
                "role": "user",
                "content": [
                    {"type": "image"},
                    {"type": "text", "text": sample['question']}
                ]
            }
        ]
        
        # Apply chat template
        input_text = tokenizer.apply_chat_template(
            messages, 
            tokenize=False, 
            add_generation_prompt=True
        )
        
        # Tokenize input
        inputs = tokenizer(
            sample['image'], 
            input_text, 
            return_tensors="pt"
        ).to(model.device)
        
        # Generate response
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=100,
                temperature=0.1,        # Low temperature for consistent results
                do_sample=True,
                pad_token_id=tokenizer.eos_token_id
            )
        
        # Decode the response
        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        predicted_answer = response.split("assistant")[-1].strip()
        
        # Store results
        results.append({
            'question': sample['question'],
            'ground_truth': sample['answers'],
            'predicted': predicted_answer,
            'question_type': sample['question_types'][0] if sample['question_types'] else 'unknown'
        })
        
        # Print detailed results
        print(f{prefix}Sample {i+1}:)
        print(fQuestion: {sample['question']})
        print(fExpected: {sample['answers'][0]})
        print(fPredicted: {predicted_answer})
        print(fType: {sample['question_types']})
        print("-" * 50)
    
    return results

# Evaluate the base model before fine-tuning
print("Testing base model performance...")
pre_results = evaluate_model_performance(
    model, tokenizer, eval_dataset, "PRE-TRAINING "
)

# Initialize the trainer
trainer = SFTTrainer(
    model=model,
    tokenizer=tokenizer,
    data_collator=UnslothVisionDataCollator(model, tokenizer),
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    args=training_args,
)

# Check GPU memory before training
print(f"GPU memory before training: {torch.cuda.memory_allocated()/1024**3:.2f} GB")

# Start training
print("\nStarting fine-tuning...")
print("This may take 30-60 minutes depending on your hardware.")
print("=" * 60)

trainer.train()

print("\nFine-tuning completed!")
print(f"GPU memory after training: {torch.cuda.memory_allocated()/1024**3:.2f} GB")

print("\nTesting fine-tuned model performance...")
post_results = evaluate_model_performance(
    model, tokenizer, eval_dataset, "POST-TRAINING "
)

# Calculate accuracy improvements
def calculate_accuracy(results):
    """Calculate accuracy based on exact match"""
    correct = 0
    total = len(results)
    
    for result in results:
        predicted = result['predicted'].lower().strip()
        ground_truth = [ans.lower().strip() for ans in result['ground_truth']]
        
        # Check if prediction matches any acceptable answer
        if any(predicted == gt for gt in ground_truth):
            correct += 1
    
    return correct / total * 100

# Compare before and after
pre_accuracy = calculate_accuracy(pre_results)
post_accuracy = calculate_accuracy(post_results)

print(f"\n" + "=" * 60)
print("PERFORMANCE COMPARISON")
print("=" * 60)
print(f"Before fine-tuning: {pre_accuracy:.1f}% accuracy")
print(f"After fine-tuning:  {post_accuracy:.1f}% accuracy")
print(f"Improvement:        +{post_accuracy - pre_accuracy:.1f} percentage points")

if post_accuracy > pre_accuracy:
    print(f"🎉 Success! The model improved by {post_accuracy - pre_accuracy:.1f}%")
else:
    print("⚠️  Consider training longer or adjusting hyperparameters")

def analyze_by_question_type(pre_results, post_results):
    """Compare performance by question type"""
    
    # Collect all question types
    question_types = set()
    for result in pre_results + post_results:
        question_types.add(result['question_type'])
    
    print("\n" + "=" * 60)
    print("PERFORMANCE BY QUESTION TYPE")
    print("=" * 60)
    
    for q_type in sorted(question_types):
        # Filter results by question type
        pre_type = [r for r in pre_results if r['question_type'] == q_type]
        post_type = [r for r in post_results if r['question_type'] == q_type]
        
        if pre_type and post_type:
            pre_acc = calculate_accuracy(pre_type)
            post_acc = calculate_accuracy(post_type)
            improvement = post_acc - pre_acc
            
            print(f"{q_type.upper()}:")
            print(f"  Before: {pre_acc:.1f}%")
            print(f"  After:  {post_acc:.1f}%")
            if improvement > 0:
                print(f"  Gain:   +{improvement:.1f}% ✅")
            else:
                print(f"  Change: {improvement:.1f}% ⚠️")
            print()

# Run the analysis
analyze_by_question_type(pre_results, post_results)

Question:

What is the total amount due?

Image:

Invoice.jpg

Question:

How many items are listed in the table?

Image:

Invoice with 5 food items at resort line items but the last line is hand written after the printed text.

Question:

What is the seller name?

Image:

Purchase order form hand written with one flow

def load_finetuned_model():
    """Load the fine-tuned model for inference"""
    model, tokenizer = FastVisionModel.from_pretrained(
        "llama-vision-docvqa-merged",
        load_in_4bit=True,
        device_map="auto"
    )
    return model, tokenizer

def analyze_document(image_path, question, model, tokenizer):
    """Analyze any document with your fine-tuned model"""
    from PIL import Image
    
    # Load the image
    image = Image.open(image_path)
    
    # Prepare the chat format
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "image"},
                {"type": "text", "text": question}
            ]
        }
    ]
    
    # Apply chat template
    input_text = tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    
    # Tokenize
    inputs = tokenizer(image, input_text, return_tensors="pt").to(model.device)
    
    # Generate answer
    with torch.no_grad():
        outputs = model.generate(**inputs, max_new_tokens=150)
    
    # Extract the answer
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    answer = response.split("assistant")[-1].strip()
    
    return answer

# Example usage
print("\nTesting production model...")
finetuned_model, finetuned_tokenizer = load_finetuned_model()

# Test with a new invoice
answer = analyze_document(
    "new_invoice.jpg", 
    "What is the total amount due?", 
    finetuned_model, 
    finetuned_tokenizer
)
print(f"Answer: {answer}")