#!/bin/bash

# Train Crystal Classifier with Tiny Model
# This script trains a very small transformer model for quick experimentation

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

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

# Use tiny model
MODEL_SIZE="tiny"

# Training parameters (adjusted for tiny model)
BATCH_SIZE=32  # Can use larger batch size with smaller model
MAX_LENGTH=256  # Shorter sequences for faster training
LEARNING_RATE=3e-5  # Slightly higher learning rate
NUM_EPOCHS=15  # More epochs since model is smaller

echo "Starting crystal classification training with TINY model..."
echo "Run name: $RUN_NAME"
echo "Data path: $DATA_PATH"
echo "Model size: $MODEL_SIZE (distilbert-base-uncased ~66M parameters)"

# 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" 