import os import sys import argparse import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader, random_split, ConcatDataset import torchvision.transforms as transforms from tqdm import tqdm try: from huggingface_hub import HfApi except ImportError: HfApi = None # Add parent directory to path so we can import the model from predictor.py sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) sys.path.append(os.getcwd()) # Fallback for Colab execution from core.predictor import SkribblPredictorModel from dataset import SkribblDataset, TUBerlinDataset from core.embedding_utils import load_glove_embeddings def load_vocab(vocab_path): with open(vocab_path, 'r') as f: return [line.strip() for line in f if line.strip()] def main(): # 2. Set up the argument parser parser = argparse.ArgumentParser(description="Train Skribbl AI Predictor") base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) parser.add_argument('--data_dir', type=str, default=os.path.join(base_dir, "data", "quickdraw_data"), help='Path to the QuickDraw dataset directory') parser.add_argument('--tu_berlin_dir', type=str, default=os.path.join(base_dir, "data", "tu_berlin_data"), help='Path to the TU Berlin dataset directory') parser.add_argument('--output_dir', type=str, default=os.path.join(base_dir, "weights"), help='Directory to save the trained model weights') parser.add_argument('--batch_size', type=int, default=128, help='Batch size for training') parser.add_argument('--epochs', type=int, default=100, help='Number of epochs to train') parser.add_argument('--lr', type=float, default=1e-4, help='Learning rate') parser.add_argument('--hf_repo_id', type=str, default=None, help='Optional: Hugging Face repo ID to upload model to (e.g. SteveaWong/AI-Drawing-Predictor)') # Use parse_known_args so Jupyter internal args don't crash the script args, _ = parser.parse_known_args() # 1. Configuration device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') vocab_path = os.path.join(base_dir, "data", "vocab.txt") data_dir = args.data_dir tu_berlin_dir = args.tu_berlin_dir output_dir = args.output_dir batch_size = args.batch_size num_epochs = args.epochs learning_rate = args.lr hf_repo_id = args.hf_repo_id # Create output directory if it doesn't exist os.makedirs(output_dir, exist_ok=True) print(f"Using device: {device}") vocab = load_vocab(vocab_path) print(f"Loaded {len(vocab)} vocab words.") # Load real embeddings glove_path = os.path.join(base_dir, "data", "glove.6B.300d.txt") word_embeddings = load_glove_embeddings(vocab, glove_path) # 2. Data Loading transform = transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) qd_dataset = SkribblDataset(data_dir=data_dir, vocab=vocab, embeddings=word_embeddings, transform=transform, max_samples_per_class=1000) tu_dataset = TUBerlinDataset(data_dir=tu_berlin_dir, vocab=vocab, embeddings=word_embeddings, transform=transform) datasets_to_concat = [] if len(qd_dataset) > 0: datasets_to_concat.append(qd_dataset) print(f"Loaded {len(qd_dataset)} samples from QuickDraw (Progressive).") if len(tu_dataset) > 0: datasets_to_concat.append(tu_dataset) print(f"Loaded {len(tu_dataset)} samples from TU Berlin (100% Complete).") if not datasets_to_concat: print("No data found! Please ensure datasets are downloaded.") return full_dataset = ConcatDataset(datasets_to_concat) print(f"Total merged samples loaded: {len(full_dataset)}") # Train/Val split train_size = int(0.8 * len(full_dataset)) val_size = len(full_dataset) - train_size train_dataset, val_dataset = random_split(full_dataset, [train_size, val_size]) train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=2) val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=2) # 3. Model & Loss model = SkribblPredictorModel(embedding_dim=300).to(device) # CosineEmbeddingLoss compares two vectors and pushes them together (y=1) or apart (y=-1) criterion = nn.CosineEmbeddingLoss() # Added weight decay (L2 regularization) to prevent overfitting optimizer = optim.Adam(model.parameters(), lr=learning_rate, weight_decay=1e-4) # 4. Training Loop best_val_loss = float('inf') for epoch in range(num_epochs): model.train() running_loss = 0.0 for i, (images, lengths, target_embeddings) in enumerate(tqdm(train_loader, desc=f"Epoch {epoch+1}/{num_epochs}")): images = images.to(device) lengths = lengths.to(device) target_embeddings = target_embeddings.to(device) optimizer.zero_grad() # Forward pass outputs = model(images, lengths) # Loss: target tensor y=1 implies we want cosine similarity between outputs and target_embeddings to be 1 y = torch.ones(outputs.size(0)).to(device) loss = criterion(outputs, target_embeddings, y) # Backward and optimize loss.backward() optimizer.step() running_loss += loss.item() # Validation model.eval() val_loss = 0.0 with torch.no_grad(): for images, lengths, target_embeddings in val_loader: images = images.to(device) lengths = lengths.to(device) target_embeddings = target_embeddings.to(device) outputs = model(images, lengths) y = torch.ones(outputs.size(0)).to(device) loss = criterion(outputs, target_embeddings, y) val_loss += loss.item() val_loss /= len(val_loader) print(f"Epoch [{epoch+1}/{num_epochs}] Validation Loss: {val_loss:.4f}") # Save best model if val_loss < best_val_loss: best_val_loss = val_loss model_path = os.path.join(output_dir, "skribbl_model.pth") torch.save(model.state_dict(), model_path) print(f"Saved best model to {model_path}") # Optionally upload to Hugging Face during training to avoid losing progress if the notebook restarts if hf_repo_id: if HfApi is None: print("huggingface_hub is not installed! Cannot upload model. Run: pip install huggingface_hub") else: if os.path.exists(model_path): print(f"Uploading checkpoint to Hugging Face repo: {hf_repo_id}...") try: api = HfApi() repo_type = "space" if "Predictor" in hf_repo_id else "model" api.upload_file( path_or_fileobj=model_path, path_in_repo="weights/skribbl_model.pth", repo_id=hf_repo_id, repo_type=repo_type ) print("✅ Successfully backed up checkpoint to Hugging Face!") except Exception as e: print(f"❌ Failed to upload checkpoint: {e}") if __name__ == '__main__': main()