Fine-tuning Gemma 3-4B with Unsloth on Modal: Production-Ready Vision-Language Training¶

from modal import App, Image as ModalImage, Volume, Secret

# Create the Modal app - this is our project namespace
app = App("Finetuned_Gemma_3_4b_it")

# Create persistent storage - everything goes here
# Models, datasets, checkpoints, evaluation results - all in one volume
exp_volume = Volume.from_name("Finetuned_Gemma_3_4b_it", create_if_missing=True)

# Mount the volume at /data in all our containers
VOLUME_CONFIG = {
    "/data": exp_volume,  # Single volume for the entire experiment
}

# Load secrets for Hugging Face and Weights & Biases
# This injects HUGGINGFACE_TOKEN and WANDB_API_KEY as environment variables
huggingface_secret = Secret.from_name("secrets-hf-wandb")

# Time constants
HOURS = 60 * 60  # Makes timeouts more readable

# Model configuration
BASE_MODEL_NAME = "unsloth/gemma-3-4b-it"  # Unsloth's optimized Gemma
WANDB_PROJECT_DEFAULT = "GemmaFinetuning"   # W&B project name
OUTPUT_DIR_DEFAULT = "/data/Finetuned_Gemma_3_4b_it"  # Where to save checkpoints

CUDA_VERSION = "12.8.1"     # Latest CUDA version
CUDA_FLAVOR = "devel"        # "devel" includes nvcc compiler
CUDA_OS = "ubuntu24.04"      # Ubuntu 24.04 LTS
CUDA_TAG = f"{CUDA_VERSION}-{CUDA_FLAVOR}-{CUDA_OS}"

FINETUNING_GPU_IMAGE = (
    # Start with NVIDIA's official CUDA image
    ModalImage.from_registry(f"nvidia/cuda:{CUDA_TAG}", add_python="3.12")

    # Install system dependencies
    # git: for cloning repos if needed
    # build-essential: gcc, make, etc. for compiling Python extensions
    .apt_install("git", "build-essential")

    # Install PyTorch first (required by most other packages)
    # Using uv for faster installs (it's like pip but 10-100x faster)
    .uv_pip_install(["torch", "torchvision", "torchaudio"])

    # Now install the ML ecosystem
    .uv_pip_install([
        # === Unsloth core ===
        "unsloth",              # The star of the show - optimized training
        "unsloth_zoo",          # Pre-configured models

        # === Quantization and efficiency ===
        "bitsandbytes",         # 8-bit optimizers, quantization
        "accelerate",           # Multi-GPU support, mixed precision
        "xformers",             # Memory-efficient attention
        "peft",                 # LoRA and other parameter-efficient methods
        "trl",                  # Transformer Reinforcement Learning
        "triton",               # GPU kernel language (used by flash-attn)
        "cut_cross_entropy",    # Optimized loss computation

        # === Transformers ecosystem ===
        "transformers",         # Hugging Face transformers
        "timm",                 # Vision model utilities

        # === Training tools ===
        "wandb",                # Experiment tracking (highly recommend!)
        "weave",                # W&B's LLM eval framework
        "deepspeed",            # For multi-GPU training (optional here)

        # === Evaluation metrics ===
        "nltk",                 # NLP toolkit
        "rouge_score",          # ROUGE metrics
        "bert_score",           # BERTScore
        "jiwer",                # Word/Character Error Rate
        "scikit-learn",         # General ML utilities

        # === Utilities ===
        "pillow",               # Image processing
        "opencv-python-headless",  # More image processing
        "gradio",               # Quick UI demos
        "hf_transfer",          # Faster Hugging Face downloads
    ])

    # Set environment variables
    .env({
        # Enable fast multi-threaded downloads from Hugging Face
        # This can be 5-10x faster for large models!
        "HF_HUB_ENABLE_HF_TRANSFER": "1",

        # Cache everything in the volume (so it persists)
        # This means we download models once, use them forever
        "HF_HOME": "/data/.cache",
    })
)

@app.function(
    image=FINETUNING_GPU_IMAGE,     # Our big image with all dependencies
    volumes=VOLUME_CONFIG,           # Mount /data volume
    secrets=[huggingface_secret],   # Inject HF token
    timeout=24 * HOURS,              # Give it up to 24 hours (large datasets)
    # Notice: No GPU! This runs on CPU to save money
)
def download_datasets(
    dataset_name: str = "unsloth/LaTeX_OCR",  # HuggingFace dataset ID
    split: str = "train",                      # Which split to download
    cache_dir: str = "/data/.cache",           # Where to cache it
):
    """
    Download and cache a dataset from Hugging Face.

    Runs on CPU (no GPU wasted on downloading files).
    Dataset gets cached in the volume, so we only download once.
    """
    from datasets import load_dataset
    import os

    # Set HF token from our secret
    # Modal injects HUGGINGFACE_TOKEN from the secret we passed in
    os.environ["HF_TOKEN"] = os.environ["HUGGINGFACE_TOKEN"]

    # Make sure cache directory exists
    os.makedirs(cache_dir, exist_ok=True)

    print(f"Downloading {dataset_name} ({split} split)...")
    print(f"Cache dir: {cache_dir}")

    # Download the dataset
    # cache_dir tells it to save in our volume (persists across runs)
    dataset = load_dataset(dataset_name, split=split, cache_dir=cache_dir)

    # Print some info
    print("\n✓ Dataset loaded successfully!")
    print(f"  - Name: {dataset_name}")
    print(f"  - Split: {split}")
    print(f"  - Number of samples: {len(dataset)}")
    print(f"  - Features: {dataset.features}")

    # CRITICAL: Commit changes to the volume
    # This persists the downloaded data
    exp_volume.commit()

    # Return metadata
    return {
        "status": "completed",
        "dataset_name": dataset_name,
        "num_samples": len(dataset),
    }

@app.function(
    image=FINETUNING_GPU_IMAGE,
    gpu="l40s:1",                   # Use a cheap GPU (L40S is ~$1/hr)
    volumes=VOLUME_CONFIG,           # Mount our volume
    secrets=[huggingface_secret],   # Need HF token for model access
    timeout=24 * HOURS,
)
def download_models(
    model_name: str = BASE_MODEL_NAME,      # "unsloth/gemma-3-4b-it"
    cache_dir: str = "/data/.cache",        # Cache in volume
):
    """
    Download and cache the base model using Unsloth's FastVisionModel.

    Why L40S GPU? Some models need a GPU just to load (for safety checks, etc.)
    L40S is cheaper than A100, perfect for this one-time download.
    """
    from unsloth import FastVisionModel
    import os
    import torch

    # Set HF token
    os.environ["HF_TOKEN"] = os.environ["HUGGINGFACE_TOKEN"]

    print(f"Downloading model: {model_name}")
    print(f"Cache dir: {cache_dir}")

    # Load the model with Unsloth's optimized loader
    # This downloads and caches the model weights
    model, processor = FastVisionModel.from_pretrained(
        model_name,
        load_in_4bit=False,                    # Full precision for now
        use_gradient_checkpointing="unsloth",  # Unsloth's optimized checkpointing
        max_seq_length=8000,                    # Max context length
        dtype=torch.bfloat16,                   # Use bfloat16 (good balance)
    )

    print(f"\n✓ Model downloaded and cached!")
    print(f"  - Model: {model_name}")
    print(f"  - Parameters: {sum(p.numel() for p in model.parameters()):,}")
    print(f"  - Cache: {cache_dir}")

    # Commit to volume
    exp_volume.commit()

    return {
        "status": "completed",
        "model_name": model_name,
        "cache_dir": cache_dir,
    }

TRAIN_GPU = "a100-80gb"    # For 4B vision models, A100-80GB is ideal
NUM_GPUS = 1                # Unsloth is optimized for single-GPU
TRAINING_GPU_CONFIG = f"{TRAIN_GPU}:{NUM_GPUS}"

@app.function(
    image=FINETUNING_GPU_IMAGE,
    volumes=VOLUME_CONFIG,
    secrets=[huggingface_secret, Secret.from_dotenv()],  # Both Modal secrets and .env
    gpu=TRAINING_GPU_CONFIG,                              # "a100-80gb:1"
    timeout=24 * HOURS,                                   # Long timeout for big datasets
)
def fine_tune_unsloth(
    # Model and dataset config
    model_path: str = BASE_MODEL_NAME,                 # Which model to fine-tune
    dataset_name: str = "unsloth/LaTeX_OCR",          # Which dataset to use
    dataset_split: str = "train",                      # Which split
    output_dir: str = OUTPUT_DIR_DEFAULT,              # Where to save checkpoints
    hub_id: str = None,                                # Push to HF Hub? (optional)
    max_samples: int = None,                           # Limit dataset (for testing)

    # LoRA hyperparameters
    lora_r: int = 32,                                  # LoRA rank (higher = more capacity)
    lora_alpha: int = 64,                              # LoRA scaling (usually 2x rank)
    lora_dropout: float = 0.0,                         # Dropout in LoRA layers

    # Training hyperparameters
    per_device_train_batch_size: int = 4,              # Batch size per GPU
    gradient_accumulation_steps: int = 4,              # Effective batch = 4 * 4 = 16
    num_train_epochs: int = 1,                         # How many epochs
    learning_rate: float = 3e-4,                       # Learning rate
    warmup_ratio: float = 0.2,                         # Warmup 20% of steps
    max_seq_length: int = 8000,                        # Max tokens per sample

    # Checkpointing
    save_strategy: str = "steps",                      # Save by steps (not epochs)
    save_steps: int = 250,                             # Save every 250 steps
    save_total_limit: int = 20,                        # Keep only 20 checkpoints
    logging_steps: int = 10,                           # Log every 10 steps

    # Weights & Biases
    wandb_project: str = WANDB_PROJECT_DEFAULT,        # W&B project name
    wandb_run_name: str = None,                        # W&B run name (auto-generated)
):
    """
    Fine-tune Gemma 3-4B vision model with LoRA using Unsloth.

    This is based on Unsloth's Colab notebook but productionized for Modal.
    """
    from unsloth import FastVisionModel, get_chat_template
    from unsloth.trainer import UnslothVisionDataCollator
    from datasets import load_dataset
    from trl import SFTTrainer, SFTConfig
    import torch
    import os
    from datetime import datetime

    # === Environment Setup ===
    print("=" * 80)
    print("SETTING UP TRAINING ENVIRONMENT")
    print("=" * 80)

    # Set up authentication tokens
    os.environ["HF_TOKEN"] = os.environ["HUGGINGFACE_TOKEN"]
    os.environ["WANDB_API_KEY"] = os.environ["WANDB_API_KEY"]
    os.environ["WANDB_PROJECT"] = wandb_project

    # Auto-generate W&B run name if not provided
    # Format: finetune_gemma-3-4b-it_20250110_143022
    if wandb_run_name is None:
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        model_short = model_path.split("/")[-1]  # Extract "gemma-3-4b-it" from path
        wandb_run_name = f"finetune_{model_short}_{timestamp}"

    os.environ["WANDB_RUN_NAME"] = wandb_run_name

    # Memory optimization: only use GPU 0
    # (In single-GPU setup, this prevents memory fragmentation)
    os.environ["CUDA_VISIBLE_DEVICES"] = "0"

    # Disable torch compile (can cause issues with some models)
    torch._dynamo.config.disable = True

    print(f"Model: {model_path}")
    print(f"Dataset: {dataset_name} ({dataset_split})")
    print(f"Output: {output_dir}")
    print(f"W&B: {wandb_project}/{wandb_run_name}")
    print("")

    # === Load Model with LoRA ===
    print("=" * 80)
    print("LOADING MODEL AND ADDING LORA ADAPTERS")
    print("=" * 80)

    # Load base model
    # Unsloth's FastVisionModel is a drop-in replacement for HF's model
    # but with optimized kernels and memory usage
    model, processor = FastVisionModel.from_pretrained(
        model_path,
        load_in_4bit=False,                    # Use full precision (more accurate)
        use_gradient_checkpointing="unsloth",  # Unsloth's optimized checkpointing
        max_seq_length=max_seq_length,         # Context window
        dtype=torch.bfloat16,                   # bfloat16 is great for training
    )

    # Add LoRA adapters
    # LoRA (Low-Rank Adaptation) trains small adapter layers instead of the full model
    # This is WAY more efficient - we only train ~1% of parameters!
    model = FastVisionModel.get_peft_model(
        model,
        finetune_vision_layers=False,      # Keep vision encoder frozen
        finetune_language_layers=True,     # Train the language model part
        finetune_attention_modules=True,   # Add LoRA to attention
        finetune_mlp_modules=True,         # Add LoRA to MLPs

        # LoRA config
        r=lora_r,                          # Rank (32 is a good default)
        lora_alpha=lora_alpha,             # Scaling (usually 2x rank)
        lora_dropout=lora_dropout,         # Dropout (0.0 often works fine)
        bias="none",                       # Don't train bias terms
        random_state=3407,                 # For reproducibility
        target_modules="all-linear",       # Apply to all linear layers
        modules_to_save=["lm_head", "embed_tokens"],  # Also train these
    )

    # Set up chat template for the model
    # This formats inputs/outputs correctly for Gemma
    processor = get_chat_template(processor, "gemma-3")

    print(f"✓ Model loaded with LoRA adapters")
    print(f"  - Base model: {model_path}")
    print(f"  - LoRA rank: {lora_r}")
    print(f"  - Trainable params: ~1-2% of total")
    print("")

    # === Load and Prepare Dataset ===
    print("=" * 80)
    print("LOADING DATASET")
    print("=" * 80)

    # Load dataset from cache (downloaded in Stage 1)
    dataset = load_dataset(dataset_name, split=dataset_split)

    # Limit dataset size if specified (useful for testing)
    if max_samples is not None and max_samples > 0:
        dataset = dataset.select(range(min(max_samples, len(dataset))))
        print(f"⚠️  Limited to {len(dataset)} samples for testing")

    print(f"✓ Dataset loaded: {len(dataset)} samples")
    print("")

    # === Format Dataset ===
    # Convert dataset to chat format that Gemma expects
    # Each sample has an image and corresponding LaTeX code

    instruction = "Write the LaTeX representation for this image."

    def convert_to_conversation(sample):
        """
        Convert a dataset sample to chat format.

        Input sample has:
          - "image": PIL Image of equation
          - "text": LaTeX code for that equation

        Output format:
          - User message: instruction + image
          - Assistant message: LaTeX code
        """
        conversation = [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": instruction},
                    {"type": "image", "image": sample["image"]},
                ],
            },
            {
                "role": "assistant",
                "content": [{"type": "text", "text": sample["text"]}],
            },
        ]
        return {"messages": conversation}

    # Convert all samples
    print("Converting dataset to chat format...")
    converted_dataset = [convert_to_conversation(sample) for sample in dataset]
    print(f"✓ Converted {len(converted_dataset)} samples")
    print("")

    # === Training Setup ===
    print("=" * 80)
    print("STARTING TRAINING")
    print("=" * 80)

    # Enable training mode (sets up gradient computation)
    FastVisionModel.for_training(model)

    # Create trainer
    # SFTTrainer is from TRL library - supervised fine-tuning trainer
    trainer = SFTTrainer(
        model=model,
        train_dataset=converted_dataset,
        processing_class=processor.tokenizer,

        # Data collator handles batching images + text
        # Unsloth's collator is optimized for vision-language models
        data_collator=UnslothVisionDataCollator(
            model=model,
            processor=processor
        ),

        # Training arguments
        args=SFTConfig(
            # === Batch size config ===
            per_device_train_batch_size=per_device_train_batch_size,
            gradient_accumulation_steps=gradient_accumulation_steps,
            # Effective batch size = 4 * 4 = 16

            # === Learning rate schedule ===
            warmup_ratio=warmup_ratio,         # Warm up for 20% of training
            num_train_epochs=num_train_epochs,
            learning_rate=learning_rate,
            lr_scheduler_type="linear",        # Linear decay after warmup

            # === Logging ===
            logging_steps=logging_steps,       # Log every 10 steps
            report_to="wandb",                 # Log to W&B

            # === Checkpointing ===
            save_strategy=save_strategy,       # Save by steps
            save_steps=save_steps,             # Every 250 steps
            save_total_limit=save_total_limit, # Keep only 20 checkpoints
            output_dir=output_dir,             # Where to save

            # === Optimization ===
            gradient_checkpointing=True,       # Trade compute for memory
            gradient_checkpointing_kwargs={"use_reentrant": False},
            max_grad_norm=0.3,                 # Gradient clipping
            optim="adamw_torch_fused",         # Fastest AdamW implementation
            weight_decay=0.01,                 # L2 regularization

            # === Precision ===
            bf16=True,                         # Use bfloat16 (faster + stable)
            tf32=False,                        # Don't need TF32

            # === Vision-specific settings ===
            remove_unused_columns=False,       # Keep all columns (need images!)
            dataset_text_field="",             # We handle formatting ourselves
            dataset_kwargs={"skip_prepare_dataset": True},
            max_length=max_seq_length,
        ),
    )

    print(f"Training config:")
    print(f"  - Effective batch size: {per_device_train_batch_size * gradient_accumulation_steps}")
    print(f"  - Learning rate: {learning_rate}")
    print(f"  - Epochs: {num_train_epochs}")
    print(f"  - Total steps: ~{len(converted_dataset) // (per_device_train_batch_size * gradient_accumulation_steps) * num_train_epochs}")
    print("")

    # === TRAIN! ===
    print("🚀 Starting training...")
    print("=" * 80)
    trainer_stats = trainer.train()
    print("=" * 80)
    print("✓ Training completed!")
    print("")

    # === Save Model ===
    print("=" * 80)
    print("SAVING MODEL")
    print("=" * 80)

    # Create output directories
    final_weights_dir = os.path.join(output_dir, "final_weights")  # Merged model
    final_lora_dir = os.path.join(output_dir, "final_lora")        # LoRA adapters only

    os.makedirs(final_weights_dir, exist_ok=True)
    os.makedirs(final_lora_dir, exist_ok=True)

    # Save LoRA adapters (small, ~100MB)
    print("Saving LoRA adapters...")
    model.save_pretrained(final_lora_dir)
    processor.save_pretrained(final_lora_dir)
    print(f"  ✓ LoRA adapters saved to {final_lora_dir}")

    # Optionally push LoRA to Hugging Face Hub
    if hub_id:
        print(f"Pushing LoRA to Hub: {hub_id}_lora")
        model.push_to_hub(
            f"{hub_id}_lora",
            token=os.environ["HUGGINGFACE_TOKEN"]
        )
        processor.push_to_hub(
            f"{hub_id}_lora",
            token=os.environ["HUGGINGFACE_TOKEN"]
        )
        print(f"  ✓ Pushed to {hub_id}_lora")

    # Save merged model (base + LoRA combined, ready to deploy)
    print("Saving merged model (this takes a few minutes)...")
    model.save_pretrained_merged(
        final_weights_dir,
        processor,
        save_method="merged_16bit"  # Save in 16-bit precision
    )
    print(f"  ✓ Merged model saved to {final_weights_dir}")

    # Optionally push merged model to Hub
    if hub_id:
        print(f"Pushing merged model to Hub: {hub_id}")
        model.push_to_hub_merged(
            hub_id,
            processor,
            token=os.environ["HUGGINGFACE_TOKEN"],
            save_method="merged_16bit"
        )
        print(f"  ✓ Pushed to {hub_id}")

    # CRITICAL: Commit everything to the volume
    # This persists checkpoints, final models, everything
    print("\nCommitting to volume...")
    exp_volume.commit()
    print("✓ Volume committed")

    print("")
    print("=" * 80)
    print("🎉 FINE-TUNING COMPLETE!")
    print("=" * 80)
    print(f"LoRA adapters: {final_lora_dir}")
    print(f"Merged model: {final_weights_dir}")
    if hub_id:
        print(f"Hugging Face: {hub_id} and {hub_id}_lora")
    print("")

    return {
        "status": "completed",
        "output_dir": output_dir,
        "lora_dir": final_lora_dir,
        "merged_dir": final_weights_dir,
        "hub_id": hub_id,
    }

@app.function(
    image=FINETUNING_GPU_IMAGE,
    volumes=VOLUME_CONFIG,
    gpu=TRAINING_GPU_CONFIG,              # Need same GPU as training
    secrets=[huggingface_secret, Secret.from_dotenv()],
    timeout=2 * HOURS,                    # Merging takes ~10-30 minutes
)
def export_model(
    lora_model_path: str = f"{OUTPUT_DIR_DEFAULT}",  # Where LoRA adapters are
    output_path: str = None,                          # Where to save merged model
    hub_model_id: str = None,                         # Optional: push to HF Hub
    push_to_hub: bool = True,                         # Whether to push
):
    """
    Export LoRA adapters and merge them with base model.

    Why? Two reasons:
    1. Merged models are easier to deploy (no need to load base + adapters separately)
    2. Merged models can be quantized for faster inference
    """
    from unsloth import FastVisionModel
    import os

    # Set HF token for pushing to Hub
    os.environ["HF_TOKEN"] = os.environ["HUGGINGFACE_TOKEN"]

    print("=" * 80)
    print("LOADING LORA MODEL AND MERGING")
    print("=" * 80)
    print(f"LoRA path: {lora_model_path}")

    # Load the LoRA model
    # This loads base model + LoRA adapters
    model, processor = FastVisionModel.from_pretrained(
        lora_model_path,           # Path to LoRA adapters
        load_in_4bit=False,        # Load in full precision
    )

    # Prepare for inference
    # This merges the LoRA weights into the base model
    FastVisionModel.for_inference(model)

    print("✓ Model loaded and LoRA weights merged")
    print("")

    # === Save or Push ===
    if push_to_hub and hub_model_id:
        # Push merged model to Hugging Face Hub
        print(f"Pushing merged model to Hub: {hub_model_id}")
        model.push_to_hub_merged(
            hub_model_id,
            processor,
            token=os.environ["HUGGINGFACE_TOKEN"],
            save_method="merged_16bit",  # Save in 16-bit (good balance)
        )
        print(f"✓ Pushed to https://huggingface.co/{hub_model_id}")
    else:
        # Save locally to volume
        if output_path is None:
            output_path = f"{lora_model_path}_merged"

        print(f"Saving merged model to: {output_path}")
        model.save_pretrained_merged(
            output_path,
            processor,
            save_method="merged_16bit"
        )
        print(f"✓ Saved to {output_path}")

    # Commit to volume
    exp_volume.commit()

    print("")
    print("=" * 80)
    print("✓ EXPORT COMPLETE!")
    print("=" * 80)

    return {
        "status": "completed",
        "lora_path": lora_model_path,
        "merged_path": output_path if not push_to_hub else hub_model_id,
    }

VLLM_CUDA_VERSION = "12.8.1"
VLLM_CUDA_TAG = f"{VLLM_CUDA_VERSION}-devel-ubuntu24.04"

VLLM_GPU_IMAGE = (
    # Start with CUDA base
    ModalImage.from_registry(f"nvidia/cuda:{VLLM_CUDA_TAG}", add_python="3.12")

    # Install system dependencies for vLLM
    .apt_install("libopenmpi-dev", "libnuma-dev")  # For distributed inference

    # Upgrade pip and install uv
    .run_commands("pip install --upgrade pip")
    .run_commands("pip install uv")

    # Install vLLM (latest version)
    .run_commands("uv pip install vllm -U --system")

    # Install supporting packages
    .pip_install(
        "datasets",                       # For eval/testing
        "pillow",                         # Image handling
        "huggingface_hub[hf_transfer]",  # Fast model downloads
        "requests",                       # HTTP requests
        "numpy",                          # Numerical ops
    )

    # Install flash-attention (required for vLLM)
    # Must be installed separately with --no-build-isolation
    .run_commands(
        "uv pip install 'flash-attn>=2.7.1,<=2.8.0' --no-build-isolation --system"
    )

    # Enable fast HF downloads
    .env({"HF_HUB_ENABLE_HF_TRANSFER": "1"})
)

# Which model to serve (path on volume)
DEFAULT_SERVE_MODEL = "/data/Finetuned_Gemma_3_4b_it/final_weights"

# GPU for serving (can be different from training!)
SERVE_GPU = "L40S"        # L40S is great for inference (~$1/hr)
SERVE_NUM_GPUS = 1
SERVE_GPU_CONFIG = f"{SERVE_GPU}:{SERVE_NUM_GPUS}"

VLLM_PORT = 8000          # Internal port

@app.function(
    image=VLLM_GPU_IMAGE,
    gpu=SERVE_GPU_CONFIG,                # L40S for serving
    scaledown_window=3 * 60,             # Scale to 0 after 3 min idle (saves $$$)
    secrets=[huggingface_secret],        # Need HF token
    volumes=VOLUME_CONFIG,                # Mount our volume (has the model)
    max_containers=2,                     # Auto-scale up to 2 instances
    timeout=24 * HOURS,
)
@modal.concurrent(max_inputs=50)         # Handle 50 concurrent requests per instance
@modal.web_server(port=8000, startup_timeout=5 * 60)  # Expose as web server
def serve_vllm():
    """
    Serve the fine-tuned model using vLLM.

    This creates an OpenAI-compatible API endpoint that:
    - Auto-scales from 0 to max_containers based on load
    - Shuts down after 3 minutes of inactivity (cost optimization!)
    - Handles up to 50 concurrent requests per container
    """
    import subprocess
    import os

    # Set HF token (might need to download model files)
    os.environ["HF_TOKEN"] = os.environ["HUGGINGFACE_TOKEN"]

    print("=" * 80)
    print("STARTING VLLM SERVER")
    print("=" * 80)
    print(f"Model: {DEFAULT_SERVE_MODEL}")
    print(f"Port: {VLLM_PORT}")
    print(f"GPU: {SERVE_GPU_CONFIG}")
    print("")

    # Build vLLM command
    cmd = [
        "vllm", "serve",                           # vLLM serve command
        "--uvicorn-log-level=info",                # Logging level
        DEFAULT_SERVE_MODEL,                       # Path to model
        "--host", "0.0.0.0",                       # Listen on all interfaces
        "--port", str(VLLM_PORT),                  # Port to serve on
        "--enforce-eager",                         # Faster startup (skip torch.compile)
        "--tensor-parallel-size", str(SERVE_NUM_GPUS),  # How many GPUs to use
        "--gpu-memory-utilization", "0.4",         # Use 40% of GPU memory (be conservative)
        "--trust-remote-code",                     # Allow custom model code
    ]

    print(f"Command: {' '.join(cmd)}")
    print("")
    print("🚀 Starting vLLM server...")
    print("=" * 80)

    # Start vLLM in background
    # Popen returns immediately, server keeps running
    subprocess.Popen(" ".join(cmd), shell=True)

from openai import OpenAI
import base64

# Create client pointing to your Modal endpoint
client = OpenAI(
    base_url="https://your-endpoint.modal.run/v1",  # Your Modal URL + /v1
    api_key="EMPTY"  # Modal doesn't require API key (it's behind Modal auth)
)

# Encode image to base64
with open("equation.jpg", "rb") as f:
    image_b64 = base64.b64encode(f.read()).decode()

# Make request (just like OpenAI!)
response = client.chat.completions.create(
    model="/data/Finetuned_Gemma_3_4b_it/final_weights",  # Model path
    messages=[
        {
            "role": "user",
            "content": [
                # Send image as base64
                {
                    "type": "image_url",
                    "image_url": {"url": f"data:image/jpeg;base64,{image_b64}"}
                },
                # Send text prompt
                {
                    "type": "text",
                    "text": "Write the LaTeX representation for this image."
                },
            ],
        },
    ],
    temperature=0.1,      # Low temp for deterministic output
    max_tokens=512,       # Max length of response
)

# Print the LaTeX code
print(response.choices[0].message.content)

import requests
import base64

# Your Modal endpoint
url = "https://your-endpoint.modal.run/v1/chat/completions"

# Load and encode image
with open("test_equation.jpg", "rb") as f:
    img_b64 = base64.b64encode(f.read()).decode()

# Make request
response = requests.post(
    url,
    json={
        "model": "/data/Finetuned_Gemma_3_4b_it/final_weights",
        "messages": [
            {
                "role": "user",
                "content": [
                    {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{img_b64}"}},
                    {"type": "text", "text": "Write the LaTeX representation for this image."}
                ]
            }
        ],
        "temperature": 0.1,
        "max_tokens": 512
    }
)

print(response.json()["choices"][0]["message"]["content"])

EVAL_IMAGE = (
    # Lightweight Debian base (no CUDA needed for eval)
    ModalImage.debian_slim(python_version="3.12")

    # Install evaluation dependencies
    .pip_install(
        "openai",                         # To call our vLLM endpoint
        "datasets",                       # Load test dataset
        "pillow",                         # Image processing
        "numpy",                          # Numerical ops
        "jiwer",                          # Word/Character Error Rate metrics
        "nltk",                           # NLP utilities
        "tqdm",                           # Progress bars
        "huggingface_hub[hf_transfer]",  # Fast dataset downloads
    )

    # Enable fast downloads
    .env({"HF_HUB_ENABLE_HF_TRANSFER": "1"})
)

@app.function(
    image=EVAL_IMAGE,                # Lightweight CPU image
    volumes=VOLUME_CONFIG,            # Access cached datasets
    secrets=[huggingface_secret],    # HF token for datasets
    timeout=2 * HOURS,                # Eval can take a while
    # No GPU! Runs on CPU
)
def evaluate_model(
    endpoint_url: str = None,                                  # vLLM endpoint (auto-detected)
    model_name: str = "/data/Finetuned_Gemma_3_4b_it/final_weights",
    dataset_name: str = "unsloth/LaTeX_OCR",                  # Test dataset
    dataset_split: str = "test",                              # Use test split
    max_samples: int = 100,                                    # How many to evaluate
    max_parallel_requests: int = 8,                           # Concurrent requests
    temperature: float = 0.1,                                 # Low temp for consistency
    max_tokens: int = 512,                                    # Max response length
):
    """
    Evaluate the fine-tuned model on LaTeX OCR test set.

    Metrics:
    - Exact Match Accuracy: % of perfect predictions
    - Character Error Rate (CER): Edit distance at character level
    - Word Error Rate (WER): Edit distance at word level
    """
    from concurrent.futures import ThreadPoolExecutor, as_completed
    from openai import OpenAI
    from datasets import load_dataset
    from jiwer import wer, cer
    from tqdm import tqdm
    import base64
    from io import BytesIO
    import os

    print("=" * 80)
    print("EVALUATING MODEL")
    print("=" * 80)

    # === Get endpoint URL ===
    if endpoint_url is None:
        # Auto-retrieve the vLLM endpoint URL
        print("Auto-detecting vLLM endpoint...")
        endpoint_url = serve_vllm.get_web_url().rstrip("/") + "/v1"

    print(f"Endpoint: {endpoint_url}")
    print(f"Model: {model_name}")
    print(f"Dataset: {dataset_name} ({dataset_split})")
    print(f"Max samples: {max_samples}")
    print("")

    # === Load test dataset ===
    print("Loading test dataset...")
    os.environ["HF_TOKEN"] = os.environ["HUGGINGFACE_TOKEN"]

    dataset = load_dataset(dataset_name, split=dataset_split)

    # Limit to max_samples
    if max_samples and max_samples < len(dataset):
        dataset = dataset.select(range(max_samples))

    print(f"✓ Loaded {len(dataset)} samples")
    print("")

    # === Set up OpenAI client ===
    client = OpenAI(
        base_url=endpoint_url,
        api_key="EMPTY"  # Modal doesn't require API key
    )

    # === Helper function to encode images ===
    def encode_image_to_base64(image):
        """Convert PIL Image to base64 string."""
        buffered = BytesIO()
        image.save(buffered, format="PNG")
        img_bytes = buffered.getvalue()
        return base64.b64encode(img_bytes).decode()

    # === Run inference on all samples (in parallel) ===
    def run_inference(sample, idx):
        """
        Run inference on a single sample.

        Returns:
            dict with "prediction" and "ground_truth"
        """
        try:
            # Encode image
            image_base64 = encode_image_to_base64(sample["image"])

            # Call API
            response = client.chat.completions.create(
                model=model_name,
                messages=[
                    {
                        "role": "user",
                        "content": [
                            {
                                "type": "image_url",
                                "image_url": {"url": f"data:image/png;base64,{image_base64}"}
                            },
                            {
                                "type": "text",
                                "text": "Write the LaTeX representation for this image."
                            },
                        ],
                    },
                ],
                temperature=temperature,
                max_tokens=max_tokens,
            )

            # Extract prediction
            prediction = response.choices[0].message.content.strip()
            ground_truth = sample["text"].strip()

            return {
                "prediction": prediction,
                "ground_truth": ground_truth,
            }
        except Exception as e:
            print(f"Error on sample {idx}: {e}")
            return {
                "prediction": "",
                "ground_truth": sample["text"].strip(),
            }

    # Run evaluation with parallel requests
    print(f"Running inference on {len(dataset)} samples...")
    print(f"Parallelism: {max_parallel_requests} concurrent requests")
    print("")

    results = []
    with ThreadPoolExecutor(max_workers=max_parallel_requests) as executor:
        # Submit all tasks
        futures = [
            executor.submit(run_inference, dataset[i], i)
            for i in range(len(dataset))
        ]

        # Collect results with progress bar
        for future in tqdm(as_completed(futures), total=len(dataset), desc="Evaluating"):
            results.append(future.result())

    # === Calculate metrics ===
    print("")
    print("=" * 80)
    print("CALCULATING METRICS")
    print("=" * 80)

    predictions = [r["prediction"] for r in results]
    ground_truths = [r["ground_truth"] for r in results]

    # Exact match accuracy
    exact_matches = sum(p == g for p, g in zip(predictions, ground_truths))
    exact_match_accuracy = exact_matches / len(results)

    # Character Error Rate (CER)
    # Lower is better, 0 = perfect
    character_error_rate = cer(ground_truths, predictions)

    # Word Error Rate (WER)
    # Lower is better, 0 = perfect
    word_error_rate = wer(ground_truths, predictions)

    # === Print results ===
    print("")
    print("📊 EVALUATION RESULTS")
    print("=" * 80)
    print(f"Samples evaluated: {len(results)}")
    print(f"")
    print(f"Exact Match Accuracy:  {exact_match_accuracy:.2%}  ({exact_matches}/{len(results)})")
    print(f"Character Error Rate:  {character_error_rate:.2%}  (lower is better)")
    print(f"Word Error Rate:       {word_error_rate:.2%}  (lower is better)")
    print("=" * 80)

    # === Print example predictions ===
    print("")
    print("📝 EXAMPLE PREDICTIONS (first 5)")
    print("=" * 80)
    for i in range(min(5, len(results))):
        print(f"\nSample {i+1}:")
        print(f"  Ground Truth: {results[i]['ground_truth']}")
        print(f"  Prediction:   {results[i]['prediction']}")
        print(f"  Match: {'✓' if results[i]['prediction'] == results[i]['ground_truth'] else '✗'}")
    print("=" * 80)

    # Save full results to volume
    results_file = f"/data/Finetuned_Gemma_3_4b_it/eval_results_{dataset_split}.json"
    import json
    with open(results_file, "w") as f:
        json.dump({
            "metrics": {
                "exact_match_accuracy": exact_match_accuracy,
                "character_error_rate": character_error_rate,
                "word_error_rate": word_error_rate,
            },
            "num_samples": len(results),
            "examples": results[:20],  # Save first 20 examples
        }, f, indent=2)

    exp_volume.commit()
    print(f"\n✓ Full results saved to {results_file}")

    return {
        "status": "completed",
        "metrics": {
            "exact_match_accuracy": exact_match_accuracy,
            "character_error_rate": character_error_rate,
            "word_error_rate": word_error_rate,
        },
        "num_samples": len(results),
        "examples": results[:10],  # Return first 10 examples
    }

# In your application
from openai import OpenAI

client = OpenAI(
    base_url="https://your-endpoint.modal.run/v1",
    api_key="EMPTY"
)

# Your app can now read LaTeX from images!