#!/bin/bash

# Train Crystal Classifier with Nano Model
# This script trains the smallest available transformer model for ultra-fast experimentation

# Set experiment name
RUN_NAME="crystal_nano_model_$(date +%Y%m%d_%H%M%S)"

# Set data path (modify this to point to your data)
DATA_PATH="data/basic"

# Use nano model (smallest available)
MODEL_SIZE="nano"

# Training parameters (optimized for nano model)
BATCH_SIZE=64  # Can use much larger batch size with tiny model
MAX_LENGTH=128  # Shorter sequences for faster training
LEARNING_RATE=5e-5  # Higher learning rate for smaller model
NUM_EPOCHS=20  # More epochs since model is very small

echo "Starting crystal classification training with NANO model..."
echo "Run name: $RUN_NAME"
echo "Data path: $DATA_PATH"
echo "Model size: $MODEL_SIZE (microsoft/DialoGPT-tiny ~33M parameters) - SMALLEST AVAILABLE"

# Run the training
python crystal_classifier.py \
    --run-name "$RUN_NAME" \
    --data-path "$DATA_PATH" \
    --model-size "$MODEL_SIZE" \
    --batch-size $BATCH_SIZE \
    --max-length $MAX_LENGTH \
    --learning-rate $LEARNING_RATE \
    --num-epochs $NUM_EPOCHS \
    --use-amp

echo "Training completed!"
echo "Check wandb for training logs and metrics" 