train.py

"""
FunctionGemma 270M Mobile Actions Finetuning Script

This script finetunes the google/functiongemma-270m-it model on the
google/mobile-actions dataset. It implements the official Gemma Cookbook
recipe within a containerized environment.

Requirements:
    - NVIDIA GPU with Ampere architecture or newer (sm_80+)
    - HF_TOKEN environment variable with Gemma license accepted
    - At least 16GB GPU VRAM

Usage:
    docker compose up --build
    # or
    python train.py
"""

import os
import sys

import torch
from datasets import load_dataset
from huggingface_hub import login
from transformers import AutoModelForCausalLM, AutoTokenizer
from trl import SFTConfig, SFTTrainer

# Model identifier - FunctionGemma 270M instruction-tuned variant
MODEL_ID = 'google/functiongemma-270m-it'

# Dataset identifier - Google Mobile Actions dataset
DATASET_ID = 'google/mobile-actions'


def check_gpu_architecture():
    """
    Validate GPU is available and supports bfloat16 (Ampere or newer).

    FunctionGemma requires bfloat16 precision for stable training. This is only
    supported on NVIDIA GPUs with compute capability 8.0+ (Ampere architecture).

    Raises:
        RuntimeError: If no GPU is available or GPU is too old.
    """
    print("=" * 60)
    print("Checking GPU configuration...")
    print("=" * 60)

    if not torch.cuda.is_available():
        raise RuntimeError(
            "No CUDA GPU detected!\n"
            "This training script requires an NVIDIA GPU with:\n"
            "  - Ampere architecture or newer (RTX 3090, A100, RTX 4090, etc.)\n"
            "  - At least 16GB VRAM\n"
            "  - NVIDIA Container Toolkit installed on host\n\n"
            "Please verify:\n"
            "  1. nvidia-smi works on your host machine\n"
            "  2. NVIDIA Container Toolkit is installed\n"
            "  3. Docker is configured for GPU access"
        )

    # Get GPU compute capability
    major, minor = torch.cuda.get_device_capability()
    compute_capability = f"{major}.{minor}"
    device_name = torch.cuda.get_device_name(0)
    vram_gb = torch.cuda.get_device_properties(0).total_memory / (1024**3)

    print(f"  GPU detected: {device_name}")
    print(f"  Compute capability: {compute_capability}")
    print(f"  VRAM: {vram_gb:.1f} GB")
    print(f"  CUDA version: {torch.version.cuda}")

    # Require compute capability >= 8.0 (Ampere) for bfloat16 support
    if major < 8:
        raise RuntimeError(
            f"GPU compute capability {compute_capability} is too old!\n"
            f"Detected GPU: {device_name}\n\n"
            "FunctionGemma requires Ampere architecture (sm_80) or newer for:\n"
            "  - Native bfloat16 support (required for stable training)\n"
            "  - Efficient gradient checkpointing\n\n"
            "Supported GPUs include:\n"
            "  - NVIDIA A100, A10, A6000 (Ampere)\n"
            "  - NVIDIA RTX 3080, 3090 (Ampere Consumer)\n"
            "  - NVIDIA RTX 4080, 4090 (Ada Lovelace)\n"
            "  - NVIDIA H100 (Hopper)"
        )

    # Warn if VRAM is below recommended
    if vram_gb < 16:
        print(f"\n  WARNING: VRAM ({vram_gb:.1f} GB) is below recommended 16GB.")
        print("  Training may fail with CUDA out of memory errors.")
        print("  Consider reducing batch size if you encounter OOM errors.\n")

    print("  GPU check passed!")
    print()


def check_hf_token():
    """
    Validate HuggingFace token is available and authenticate.

    The HF_TOKEN must:
        1. Be set as an environment variable
        2. Have read access to the model repository
        3. Have accepted the Gemma license agreement

    Raises:
        ValueError: If HF_TOKEN environment variable is not set.

    Returns:
        str: The HuggingFace token.
    """
    print("=" * 60)
    print("Checking HuggingFace authentication...")
    print("=" * 60)

    hf_token = os.environ.get('HF_TOKEN')

    if not hf_token:
        raise ValueError(
            "HF_TOKEN environment variable is required!\n\n"
            "To get your token:\n"
            "  1. Go to https://huggingface.co/settings/tokens\n"
            "  2. Create a new token with 'read' access\n"
            "  3. Accept the Gemma license at:\n"
            "     https://huggingface.co/google/functiongemma-270m-it\n\n"
            "To set the token:\n"
            "  1. Copy .env.example to .env\n"
            "  2. Add your token: HF_TOKEN=hf_xxxxxxxxxx\n"
            "  3. Restart the container"
        )

    # Mask token for logging (show first 4 and last 4 characters)
    masked_token = f"{hf_token[:7]}...{hf_token[-4:]}" if len(hf_token) > 11 else "***"
    print(f"  HF_TOKEN found: {masked_token}")

    # Authenticate with HuggingFace Hub
    print("  Authenticating with HuggingFace Hub...")
    login(token=hf_token)
    print("  Authentication successful!")
    print()

    return hf_token


def run_startup_validations():
    """
    Run all startup validations before training begins.

    This function is called at the start of training to fail fast
    if the environment is not properly configured. It's better to
    catch configuration errors upfront than during training.
    """
    print()
    print("#" * 60)
    print("# FunctionGemma 270M Mobile Actions Finetuning")
    print("#" * 60)
    print()
    print("Running startup validations...")
    print()

    # Check GPU first (required for training)
    check_gpu_architecture()

    # Check HuggingFace authentication
    check_hf_token()

    print("=" * 60)
    print("All startup validations passed!")
    print("=" * 60)
    print()


def load_model_and_tokenizer():
    """
    Load FunctionGemma model and tokenizer with mandatory configurations.

    CRITICAL CONFIGURATION NOTES:
        - attn_implementation='eager' is MANDATORY. Using SDPA or flash_attention_2
          causes NaN gradients during training. This is a known issue with Gemma models.
        - torch_dtype=torch.bfloat16 is REQUIRED. Using float16 produces erratic
          model outputs. bfloat16 requires Ampere (sm_80+) GPU architecture.
        - device_map='auto' enables automatic GPU memory allocation across devices.

    These settings are derived from the official Gemma Cookbook recipe and are
    essential for stable training without gradient issues.

    Returns:
        tuple: (model, tokenizer) - The loaded model and tokenizer instances.

    Raises:
        Exception: If model loading fails (e.g., authentication issues,
            network errors, or insufficient GPU memory).
    """
    print("=" * 60)
    print("Loading model and tokenizer...")
    print("=" * 60)

    print(f"  Model: {MODEL_ID}")
    print("  Configuration:")
    print("    - attn_implementation: eager (CRITICAL: prevents NaN gradients)")
    print("    - torch_dtype: bfloat16 (REQUIRED: prevents erratic outputs)")
    print("    - device_map: auto (automatic GPU allocation)")
    print()

    # Load tokenizer
    print("  Loading tokenizer...")
    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
    print("  Tokenizer loaded successfully!")

    # Load model with mandatory configurations
    # WARNING: Do NOT change these settings without understanding the implications:
    # - eager attention: SDPA/flash_attention cause NaN during training
    # - bfloat16: float16 causes unstable outputs
    print("  Loading model (this may take a moment)...")
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID,
        device_map='auto',
        torch_dtype=torch.bfloat16,
        attn_implementation='eager'  # CRITICAL: SDPA/flash_attention causes NaN
    )
    print("  Model loaded successfully!")

    # Log model device placement
    if hasattr(model, 'hf_device_map'):
        print(f"  Device map: {model.hf_device_map}")

    print()
    return model, tokenizer


def check_existing_dataset():
    """
    Check if the mobile-actions dataset already exists in the HuggingFace cache.

    This function checks for existing cached dataset to offer the user a choice
    between reusing the cache or downloading fresh. When running in a container
    with a mounted volume, this prevents unnecessary re-downloads if the dataset
    was already cached in a previous container run.

    Returns:
        str or None: Path to existing dataset cache if found, None otherwise.
    """
    # HuggingFace datasets cache path pattern
    # The cache directory uses the format: ~/.cache/huggingface/datasets/<org>___<dataset>
    cache_base = os.path.expanduser('~/.cache/huggingface/datasets')
    dataset_cache_path = os.path.join(cache_base, 'google___mobile-actions')

    # Check if the dataset cache directory exists
    if os.path.exists(dataset_cache_path):
        return dataset_cache_path

    return None


def load_dataset_mobile_actions():
    """
    Load the Google Mobile Actions dataset for training.

    The mobile-actions dataset contains 9,650 examples of mobile UI interactions
    formatted as conversational data with tool definitions. Each example includes:
        - tools: List of available function/tool definitions
        - messages: Conversation history with 'developer', 'user', 'assistant' roles
        - datetime: Timestamp of the interaction

    Note: Only a 'train' split is available - no validation split is provided.
    The dataset is already formatted for SFTTrainer with the 'messages' field.

    Dataset loading is controlled by the DATASET_MODE environment variable:
        - DATASET_MODE=cache: Use cached dataset if available, otherwise download (default)
        - DATASET_MODE=fresh: Force re-download even if cache exists

    Returns:
        datasets.Dataset: The loaded training dataset.

    Raises:
        Exception: If dataset loading fails (e.g., network errors,
            authentication issues for private datasets).
    """
    print("=" * 60)
    print("Loading dataset...")
    print("=" * 60)

    print(f"  Dataset: {DATASET_ID}")
    print("  Split: train (only split available)")
    print()

    # Check if dataset already exists in cache
    existing_cache_path = check_existing_dataset()

    # Read DATASET_MODE from environment variable
    # Values: 'cache' (use cached if available), 'fresh' (force fresh download)
    # Default: 'cache' (use cached dataset if available)
    dataset_mode = os.getenv('DATASET_MODE', 'cache').lower()

    if dataset_mode not in ('cache', 'fresh'):
        print(f"  WARNING: Invalid DATASET_MODE='{dataset_mode}'. Using 'fresh'.")
        dataset_mode = 'fresh'

    print(f"  DATASET_MODE: {dataset_mode}")

    if existing_cache_path and dataset_mode == 'cache':
        print("  Cached dataset found!")
        print(f"  Cache location: {existing_cache_path}")
        print("  Using cached dataset...")
        download_mode = 'reuse_cache_if_exists'
    elif existing_cache_path and dataset_mode == 'fresh':
        print("  Cached dataset found, but DATASET_MODE=fresh")
        print("  Re-downloading dataset...")
        download_mode = 'force_redownload'
    else:
        print("  No cached dataset found - downloading...")
        download_mode = 'reuse_cache_if_exists'

    print()

    # Load the mobile-actions dataset
    # Note: Only 'train' split exists - no validation split provided
    print("  Downloading and loading dataset...")
    dataset = load_dataset(
        'google/mobile-actions',
        split='train',
        download_mode=download_mode
    )

    # Log dataset statistics
    print(f"  Dataset loaded successfully!")
    print(f"  Number of examples: {len(dataset):,}")
    print(f"  Features: {list(dataset.features.keys())}")
    print()

    return dataset


def create_trainer(model, tokenizer, dataset):
    """
    Create SFTTrainer with reference hyperparameters from Gemma Cookbook.

    CRITICAL CONFIGURATION NOTES:
        - assistant_only_loss=False: Compute loss on the entire sequence (prompt + response).
          When True, only computes loss on assistant responses.
        - gradient_checkpointing=True: Essential for fitting model in 16GB VRAM.
          Trades compute for memory by recomputing activations during backward pass.
        - bf16=True: Use bfloat16 mixed precision for stable training.

    Reference hyperparameters from Gemma Cookbook:
        - 2 epochs of training
        - Learning rate: 1e-5 with cosine scheduler
        - Gradient accumulation steps: 8
        - Batch size: 4

    Args:
        model: The loaded FunctionGemma model.
        tokenizer: The model tokenizer.
        dataset: The preprocessed training dataset.

    Returns:
        SFTTrainer: Configured trainer ready for training.
    """
    print("=" * 60)
    print("Configuring SFTTrainer...")
    print("=" * 60)

    # Configure training arguments with reference hyperparameters
    # These values are from the official Gemma Cookbook recipe
    training_args = SFTConfig(
        output_dir='./results',

        # Training hyperparameters from reference notebook
        num_train_epochs=2,                    # Reference uses 2 epochs
        per_device_train_batch_size=4,         # Reference uses 4
        gradient_accumulation_steps=8,         # Reference uses 8
        learning_rate=1e-5,                    # Reference uses 1e-5
        lr_scheduler_type='cosine',            # Reference uses cosine scheduler
        warmup_ratio=0.1,                      # 10% warmup

        # Logging and saving
        logging_steps=10,                      # Log every 10 steps
        save_strategy='epoch',                 # Save checkpoint each epoch

        # Precision and memory optimization
        bf16=True,                             # Use bfloat16 (required for Gemma)
        gradient_checkpointing=True,           # CRITICAL: Memory efficiency for 16GB VRAM

        # Loss configuration
        # When False, compute loss on the entire sequence (prompt + response)
        # When True, only compute loss on assistant responses
        assistant_only_loss=False,

        # NEVER use liger kernel with assistant_only_loss - incompatible!
        # use_liger_kernel=True  # DO NOT ENABLE - silently ignores loss masking

        # Disable automatic model push to hub
        push_to_hub=False,
    )

    print("  Training configuration:")
    print(f"    - Output directory: {training_args.output_dir}")
    print(f"    - Number of epochs: {training_args.num_train_epochs}")
    print(f"    - Batch size: {training_args.per_device_train_batch_size}")
    print(f"    - Gradient accumulation steps: {training_args.gradient_accumulation_steps}")
    print(f"    - Effective batch size: {training_args.per_device_train_batch_size * training_args.gradient_accumulation_steps}")
    print(f"    - Learning rate: {training_args.learning_rate}")
    print(f"    - LR scheduler: {training_args.lr_scheduler_type}")
    print(f"    - bf16: {training_args.bf16}")
    print(f"    - Gradient checkpointing: {training_args.gradient_checkpointing}")
    print(f"    - assistant_only_loss: {training_args.assistant_only_loss}")
    print()

    # Create the SFTTrainer
    # Note: Use processing_class=tokenizer (not tokenizer=tokenizer) for TRL 0.27.0+
    print("  Creating trainer...")
    trainer = SFTTrainer(
        model=model,
        args=training_args,
        train_dataset=dataset,
        processing_class=tokenizer,  # Correct parameter name for TRL 0.27.0+
    )

    print("  SFTTrainer configured successfully!")
    print()

    return trainer


def run_training(trainer, model, tokenizer):
    """
    Execute training with progress logging and save the final model.

    This function runs the SFT training loop and handles:
        1. Training execution with automatic progress logging
        2. Checkpoint saving at each epoch (configured in SFTConfig)
        3. Final model and tokenizer saving after training completes
        4. Training metrics reporting

    The trainer is already configured with:
        - logging_steps=10: Logs metrics every 10 training steps
        - save_strategy='epoch': Saves checkpoints at each epoch
        - gradient_checkpointing=True: Memory-efficient training

    Args:
        trainer: Configured SFTTrainer instance.
        model: The FunctionGemma model (used for reference, saved via trainer).
        tokenizer: The tokenizer to save alongside the model.

    Returns:
        transformers.trainer_utils.TrainOutput: Training results including
            global_step, training_loss, and metrics.

    Raises:
        Exception: If training fails (e.g., CUDA OOM, gradient issues).
    """
    print("=" * 60)
    print("Starting training...")
    print("=" * 60)
    print()

    # Calculate and display training info
    num_examples = len(trainer.train_dataset)
    batch_size = trainer.args.per_device_train_batch_size
    grad_accum = trainer.args.gradient_accumulation_steps
    effective_batch_size = batch_size * grad_accum
    steps_per_epoch = num_examples // effective_batch_size
    total_epochs = trainer.args.num_train_epochs
    total_steps = int(steps_per_epoch * total_epochs)

    print("  Training overview:")
    print(f"    - Dataset size: {num_examples:,} examples")
    print(f"    - Batch size: {batch_size}")
    print(f"    - Gradient accumulation: {grad_accum}")
    print(f"    - Effective batch size: {effective_batch_size}")
    print(f"    - Steps per epoch: ~{steps_per_epoch:,}")
    print(f"    - Total epochs: {int(total_epochs)}")
    print(f"    - Total training steps: ~{total_steps:,}")
    print(f"    - Logging every: {trainer.args.logging_steps} steps")
    print(f"    - Checkpoints: saved each epoch to {trainer.args.output_dir}")
    print()
    print("  Training progress will be logged below...")
    print("  (Loss should decrease over time; NaN indicates a configuration issue)")
    print()
    print("-" * 60)

    # Execute training
    # The trainer handles:
    #   - Forward/backward passes with gradient accumulation
    #   - Learning rate scheduling (cosine with warmup)
    #   - Logging to console every logging_steps
    #   - Checkpoint saving at each epoch
    #   - Gradient checkpointing for memory efficiency
    train_result = trainer.train()

    print("-" * 60)
    print()
    print("=" * 60)
    print("Training completed!")
    print("=" * 60)
    print()

    # Report training metrics
    metrics = train_result.metrics
    print("  Training metrics:")
    print(f"    - Total steps: {train_result.global_step:,}")
    print(f"    - Training loss: {metrics.get('train_loss', 'N/A'):.4f}" if isinstance(metrics.get('train_loss'), (int, float)) else f"    - Training loss: {metrics.get('train_loss', 'N/A')}")
    print(f"    - Training runtime: {metrics.get('train_runtime', 0):.1f} seconds")
    print(f"    - Samples per second: {metrics.get('train_samples_per_second', 0):.2f}")
    print(f"    - Steps per second: {metrics.get('train_steps_per_second', 0):.4f}")
    print()

    # Save the final trained model and tokenizer
    print("=" * 60)
    print("Saving final model...")
    print("=" * 60)

    final_model_path = f"{trainer.args.output_dir}/final_model"
    print(f"  Saving to: {final_model_path}")

    # Save model using trainer's save method (handles distributed training correctly)
    trainer.save_model(final_model_path)
    print("  Model saved successfully!")

    # Save tokenizer alongside model for easy loading
    tokenizer.save_pretrained(final_model_path)
    print("  Tokenizer saved successfully!")

    print()
    print(f"  To load the finetuned model:")
    print(f"    model = AutoModelForCausalLM.from_pretrained('{final_model_path}')")
    print(f"    tokenizer = AutoTokenizer.from_pretrained('{final_model_path}')")
    print()

    return train_result


def run_inference_check(model, tokenizer):
    """
    Run inference on test examples to verify the finetuned model works correctly.

    This function tests the finetuned model with 3 mobile action examples to ensure
    it can generate appropriate tool calls for mobile UI interactions. The examples
    cover common mobile actions that match the training dataset: turning on flashlight,
    creating calendar events, and opening WiFi settings.

    The prompts include tool definitions (matching the mobile-actions dataset format)
    so the model can generate proper function calls.

    Args:
        model: The finetuned FunctionGemma model.
        tokenizer: The model tokenizer.

    Note:
        This is a basic sanity check. For production use, implement comprehensive
        evaluation on a held-out test set with automated metrics.
    """
    print("=" * 60)
    print("Running inference check...")
    print("=" * 60)
    print()

    # Define tools matching the mobile-actions dataset format
    # These are the same tool definitions used during training
    tools = [
        {
            "type": "function",
            "function": {
                "name": "turn_on_flashlight",
                "description": "Turns the flashlight on.",
                "parameters": {
                    "type": "object",
                    "properties": {},
                    "required": []
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "turn_off_flashlight",
                "description": "Turns the flashlight off.",
                "parameters": {
                    "type": "object",
                    "properties": {},
                    "required": []
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "create_calendar_event",
                "description": "Creates a new calendar event.",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "title": {
                            "type": "string",
                            "description": "The title of the event."
                        },
                        "datetime": {
                            "type": "string",
                            "description": "The date and time of the event in the format YYYY-MM-DDTHH:MM:SS."
                        }
                    },
                    "required": ["title", "datetime"]
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "open_wifi_settings",
                "description": "Opens the WiFi settings page.",
                "parameters": {
                    "type": "object",
                    "properties": {},
                    "required": []
                }
            }
        },
    ]

    # Define 3 mobile action test examples with proper message format
    # Each example includes a developer system message and user request
    test_examples = [
        {
            "name": "Turn On Flashlight",
            "messages": [
                {"role": "user", "content": "Turn on the flashlight"},
            ],
        },
        {
            "name": "Create Calendar Event",
            "messages": [
                {"role": "user", "content": "Create a meeting called 'Team Standup' tomorrow at 10 AM"},
            ],
        },
        {
            "name": "Open WiFi Settings",
            "messages": [
                {"role": "user", "content": "Open the WiFi settings"},
            ],
        },
    ]

    # Set model to evaluation mode
    model.eval()

    print(f"  Testing {len(test_examples)} mobile action examples...")
    print(f"  Available tools: {[t['function']['name'] for t in tools]}")
    print()

    for i, example in enumerate(test_examples, 1):
        print(f"  Test {i}: {example['name']}")
        print("-" * 40)

        # Format prompt using apply_chat_template with tools
        # This ensures the model sees the tool definitions during inference
        prompt = tokenizer.apply_chat_template(
            example["messages"],
            tools=tools,
            add_generation_prompt=True,
            tokenize=False
        )

        # Tokenize the formatted prompt
        inputs = tokenizer(
            prompt,
            return_tensors="pt",
            add_special_tokens=False
        ).to(model.device)

        # Generate response (function call)
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=128,
                do_sample=False,  # Deterministic for reproducibility
                pad_token_id=tokenizer.pad_token_id or tokenizer.eos_token_id,
            )

        # Extract only the generated response (after the prompt)
        generated_text = tokenizer.decode(
            outputs[0][inputs['input_ids'].shape[1]:],
            skip_special_tokens=False
        )

        print(f"  User prompt: {example['messages'][0]['content']}")
        print(f"  Generated tool call: {generated_text.strip()}")
        print()

    print("=" * 60)
    print("Inference check completed!")
    print("=" * 60)
    print()


def check_existing_weights():
    """
    Check if finetuned model weights already exist in the results volume.

    This function checks for existing model weights to skip redundant training.
    When running in a container with a mounted volume, this prevents retraining
    if the model was already finetuned in a previous container run.

    Returns:
        str or None: Path to existing model if found, None otherwise.
    """
    final_model_path = './results/final_model'

    # Check if the final model directory exists with required files
    if os.path.exists(final_model_path):
        # Check for essential model files
        config_exists = os.path.exists(os.path.join(final_model_path, 'config.json'))
        model_exists = (
            os.path.exists(os.path.join(final_model_path, 'model.safetensors')) or
            os.path.exists(os.path.join(final_model_path, 'pytorch_model.bin'))
        )

        if config_exists and model_exists:
            return final_model_path

    return None


def load_finetuned_model(model_path):
    """
    Load a previously finetuned model from disk.

    Args:
        model_path: Path to the saved model directory.

    Returns:
        tuple: (model, tokenizer) - The loaded model and tokenizer instances.
    """
    print("=" * 60)
    print("Loading existing finetuned model...")
    print("=" * 60)

    print(f"  Model path: {model_path}")
    print("  Configuration:")
    print("    - attn_implementation: eager (CRITICAL: prevents NaN gradients)")
    print("    - torch_dtype: bfloat16 (REQUIRED: prevents erratic outputs)")
    print("    - device_map: auto (automatic GPU allocation)")
    print()

    # Load tokenizer from saved model
    print("  Loading tokenizer...")
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    print("  Tokenizer loaded successfully!")

    # Load model with mandatory configurations
    print("  Loading model...")
    model = AutoModelForCausalLM.from_pretrained(
        model_path,
        device_map='auto',
        torch_dtype=torch.bfloat16,
        attn_implementation='eager'
    )
    print("  Model loaded successfully!")

    if hasattr(model, 'hf_device_map'):
        print(f"  Device map: {model.hf_device_map}")

    print()
    return model, tokenizer


def main():
    """
    Main training function.

    This function orchestrates the complete training pipeline:
        1. Run startup validations
        2. Check if finetuned weights already exist (skip training if found)
        3. Load model and tokenizer
        4. Load and preprocess dataset
        5. Configure SFTTrainer
        6. Run training and save model
        7. Run inference check on test examples
    """
    # Run all startup checks
    run_startup_validations()

    # Check if finetuned model already exists in the volume
    existing_model_path = check_existing_weights()

    # Read WEIGHTS_MODE from environment variable
    # Values: 'cache' (use cached weights if available), 'fresh' (force retraining)
    # Default: 'cache' (use cached weights if available)
    weights_mode = os.getenv('WEIGHTS_MODE', 'cache').lower()

    if weights_mode not in ('cache', 'fresh'):
        print(f"  WARNING: Invalid WEIGHTS_MODE='{weights_mode}'. Using 'fresh'.")
        weights_mode = 'fresh'

    print(f"  WEIGHTS_MODE: {weights_mode}")
    print()

    if existing_model_path:
        print("=" * 60)
        print("Existing finetuned model found!")
        print("=" * 60)
        print()
        print(f"  Found weights at: {existing_model_path}")
        print()

        if weights_mode == 'cache':
            print("  WEIGHTS_MODE=cache - Using cached weights (skipping training)")
            print()
        else:
            print("  WEIGHTS_MODE=fresh - Retraining from scratch")
            print()
            # Set to None to trigger training pipeline
            existing_model_path = None

    if existing_model_path:
        print("  Loading existing model instead of training.")
        print()

        # Load the existing finetuned model
        model, tokenizer = load_finetuned_model(existing_model_path)

        # Run inference check on the existing model
        run_inference_check(model, tokenizer)

        # Final summary for existing model
        print("#" * 60)
        print("# FunctionGemma - Using Existing Model")
        print("#" * 60)
        print()
        print("  Summary:")
        print(f"    - Loaded existing model from: {existing_model_path}")
        print("    - Training was skipped (weights already exist)")
        print()
        print("  To retrain, set WEIGHTS_MODE=fresh in your environment.")
        print()

    else:
        # No existing model found - run full training pipeline
        print("=" * 60)
        print("No existing model found - starting training...")
        print("=" * 60)
        print()

        # Load model and tokenizer with mandatory configurations
        model, tokenizer = load_model_and_tokenizer()

        # Load the mobile-actions dataset
        dataset = load_dataset_mobile_actions()

        # Configure SFTTrainer with reference hyperparameters
        trainer = create_trainer(model, tokenizer, dataset)

        # Run training and save the finetuned model
        train_result = run_training(trainer, model, tokenizer)

        # Run inference check on test examples
        run_inference_check(model, tokenizer)

        # Final summary
        print("#" * 60)
        print("# FunctionGemma Finetuning Complete!")
        print("#" * 60)
        print()
        print("  Summary:")
        print(f"    - Training steps completed: {train_result.global_step:,}")
        print(f"    - Final training loss: {train_result.metrics.get('train_loss', 'N/A')}")
        print(f"    - Checkpoints saved to: ./results/")
        print(f"    - Final model saved to: ./results/final_model/")
        print()
        print("  Next steps:")
        print("    1. Verify the model works with inference testing")
        print("    2. Evaluate on held-out examples if available")
        print("    3. Optionally push to HuggingFace Hub")
        print()
        print("  Note: On next container start with WEIGHTS_MODE=cache, training")
        print("  will be skipped as the model weights are persisted in the volume.")
        print()


if __name__ == '__main__':
    main()