Upload train.py with huggingface_hub

train.py

@@ -0,0 +1,178 @@
+"""
+Training script for ML Pipeline on Kaggle with HF Spaces storage.
+
+This script:
+1. Downloads training data from Hugging Face Datasets
+2. Trains a model using GPU acceleration
+3. Pushes the trained model to Hugging Face Model Hub
+"""
+
+import os
+from config import MODEL_REPO_ID, DATASET_REPO_ID, validate_config
+
+# Validate configuration before proceeding
+validate_config()
+
+from datasets import load_dataset
+from transformers import (
+    AutoTokenizer,
+    AutoModelForSequenceClassification,
+    TrainingArguments,
+    Trainer,
+)
+from sklearn.metrics import accuracy_score, f1_score
+import torch
+
+
+def load_training_data():
+    """Load dataset from Hugging Face."""
+    print(f"Loading dataset...")
+    
+    # Option 1: Load from HF Datasets Hub (public datasets)
+    # Using Spotify Songs dataset: https://huggingface.co/datasets/gem1925/spotify_songs
+    try:
+        print("Loading Spotify songs dataset...")
+        dataset = load_dataset("gem1925/spotify_songs", split="train")
+        print(f"✓ Loaded Spotify dataset: {len(dataset)} samples")
+        
+        # Convert to text classification format
+        from datasets import Dataset
+        # Use song name + artist as text, energy level as label
+        texts = [f"{row['song_name']} by {row['artist']}" for row in dataset]
+        labels = [1 if row['energy'] > 0.5 else 0 for row in dataset]
+        
+        # Create new dataset with text/label format
+        new_dataset = Dataset.from_dict({"text": texts, "label": labels})
+        new_dataset = new_dataset.train_test_split(test_size=0.2)
+        return new_dataset
+        
+    except Exception as e:
+        print(f"Could not load Spotify dataset: {e}")
+        print("Falling back to IMDB reviews dataset...")
+        
+        # Fallback to IMDB
+        try:
+            dataset = load_dataset("imdb")
+            print(f"✓ Loaded IMDB dataset")
+            return dataset
+        except Exception as e2:
+            print(f"Could not load IMDB: {e2}")
+            print("Using sample dataset for demonstration...")
+            from datasets import Dataset
+            sample_data = {
+                "text": [
+                    "I love this product! It works great.",
+                    "Terrible experience, would not recommend.",
+                    "Amazing quality and fast shipping.",
+                    "Waste of money, broke after one day.",
+                ] * 100,
+                "label": [1, 0, 1, 0] * 100,
+            }
+            dataset = Dataset.from_dict(sample_data).train_test_split(test_size=0.2)
+            return dataset
+    
+    return dataset
+
+
+def compute_metrics(eval_pred):
+    """Compute evaluation metrics."""
+    predictions, labels = eval_pred
+    predictions = predictions.argmax(axis=1)
+    return {
+        "accuracy": accuracy_score(labels, predictions),
+        "f1": f1_score(labels, predictions),
+    }
+
+
+def main():
+    """Main training function."""
+    print("=" * 50)
+    print("🚀 ML Training Pipeline - Starting")
+    print("=" * 50)
+    
+    # Check GPU availability
+    if torch.cuda.is_available():
+        print(f"✓ GPU Available: {torch.cuda.get_device_name(0)}")
+    else:
+        print("⚠ No GPU detected, training on CPU")
+    
+    # Load data
+    print("\n📊 Loading training data...")
+    dataset = load_training_data()
+    print(f"✓ Dataset loaded: {len(dataset['train'])} training samples")
+    
+    # Load tokenizer and model
+    print("\n🤖 Loading model and tokenizer...")
+    model_name = "distilbert-base-uncased"
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForSequenceClassification.from_pretrained(
+        model_name, 
+        num_labels=2
+    )
+    
+    # Tokenize data
+    print("\n📝 Tokenizing data...")
+    
+    def tokenize(batch):
+        return tokenizer(
+            batch["text"], 
+            padding="max_length", 
+            truncation=True, 
+            max_length=128
+        )
+    
+    tokenized_dataset = dataset.map(tokenize, batched=True)
+    tokenized_dataset = tokenized_dataset.rename_column("label", "labels")
+    tokenized_dataset.set_format(
+        type="torch", 
+        columns=["input_ids", "attention_mask", "labels"]
+    )
+    
+    # Training arguments
+    training_args = TrainingArguments(
+        output_dir="./results",
+        num_train_epochs=3,
+        per_device_train_batch_size=16,
+        per_device_eval_batch_size=32,
+        warmup_steps=500,
+        weight_decay=0.01,
+        logging_dir="./logs",
+        logging_steps=100,
+        eval_strategy="epoch",
+        save_strategy="epoch",
+        load_best_model_at_end=True,
+        push_to_hub=True,
+        hub_model_id=MODEL_REPO_ID,
+        hub_token=os.getenv("HF_TOKEN"),
+    )
+    
+    # Initialize trainer
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=tokenized_dataset["train"],
+        eval_dataset=tokenized_dataset["test"],
+        compute_metrics=compute_metrics,
+    )
+    
+    # Train
+    print("\n🔥 Starting training...")
+    trainer.train()
+    
+    # Evaluate
+    print("\n📈 Evaluating model...")
+    results = trainer.evaluate()
+    print(f"✓ Evaluation results: {results}")
+    
+    # Push to HF Hub
+    print(f"\n💾 Pushing model to Hugging Face: {MODEL_REPO_ID}")
+    trainer.push_to_hub()
+    print(f"✓ Model successfully pushed to: https://huggingface.co/{MODEL_REPO_ID}")
+    
+    print("\n" + "=" * 50)
+    print("✅ Training Complete!")
+    print("=" * 50)
+
+
+if __name__ == "__main__":
+    main()