train.py

#!/usr/bin/env python3
"""
Train the Tiny Conversational AI Model
Usage: python3 train.py [output_path]
"""

import sys
from tiny_chatbot import TinyChatbot, train, prepare_data
from tiny_chatbot import VOCAB_SIZE, EMBED_DIM, NUM_HEADS, FF_DIM, NUM_LAYERS, MAX_LEN

def main():
    output_path = sys.argv[1] if len(sys.argv) > 1 else 'chatbot_model.pkl'
    
    print("=" * 50)
    print("Training Tiny Conversational AI")
    print("=" * 50)
    
    # Create model
    print("\n1. Creating model architecture...")
    model = TinyChatbot(VOCAB_SIZE, EMBED_DIM, NUM_HEADS, FF_DIM, NUM_LAYERS, MAX_LEN)
    print(f"   - Embedding dim: {EMBED_DIM}")
    print(f"   - Attention heads: {NUM_HEADS}")
    print(f"   - Transformer layers: {NUM_LAYERS}")
    
    # Prepare data
    print("\n2. Preparing training data...")
    data = prepare_data()
    print(f"   - Training samples: {len(data)}")
    
    # Train
    print("\n3. Training model...")
    train(model, data, epochs=200)
    
    # Save model
    model.save(output_path)
    print(f"\n4. Model saved to '{output_path}'")
    print("=" * 50)

if __name__ == "__main__":
    main()