upgraded model

README.md

@@ -1,92 +1 @@
-# Hugging Face Training Pipeline
-
-Fine-tune any HF model on any HF dataset. Optimized for 2x T4 GPUs.
-
-## Run
-
-```bash
-cd ~/code/hf-training
-pip install -r requirements.txt
-cp .env.example .env  # Add your HF token
-./run.sh spotify
-```
-
-## Commands
-
-```bash
-./run.sh spotify        # BERT on Spotify
-./run.sh gpt2_spotify   # GPT-2 on Spotify
-./run.sh test           # Test trained model
-```
-
-## What You'll See
-
-### Training
-```
-📊 Loading dataset: maharshipandya/spotify-tracks-dataset
-  Train: 114000 samples
-
-🤖 Loading model: gpt2
-  Labels: 114
-  Parameters: 124,527,360
-
-=== Epoch 1.00 Complete ===
-  eval_accuracy: 0.45
-
-=== Final Metrics ===
-  eval_accuracy: 0.78
-```
-
-### Test
-```
-Track: 'Bohemian Rhapsody'
-  Predicted: rock (85.23%)
-Track: 'Shape of You'
-  Predicted: pop (92.10%)
-```
-
-## Files
-
-- `configs/config.yaml` - Default config
-- `configs/spotify.yaml` - Spotify dataset config
-- `configs/gpt2_spotify.yaml` - GPT-2 config
-- `src/training_pipeline.py` - Main training code
-- `run.sh` - Run script
-
-## Models (Free)
-
-| Model | Config | Size |
-|-------|--------|------|
-| BERT | `spotify` | 110M |
-| GPT-2 | `gpt2_spotify` | 117M |
-
-## Deploy to Hugging Face Space
-
-After training, deploy your model as a web app:
-
-```bash
-./deploy_space.sh spotify
-```
-
-This creates a Gradio Space at: `https://huggingface.co/spaces/your-username/spotify`
-
-## Clone on GPU Machine
-
-To run on another machine with GPU:
-
-```bash
-git clone https://huggingface.co/maxxcarl/spotify-training
-cd spotify-training
-pip install -r requirements.txt
-cp .env.example .env  # Add your HF token
-./run.sh gpt2_spotify
-```
-
-## Run on Kaggle (Free GPU)
-
-1. Create new notebook at https://www.kaggle.com/code
-2. Add secrets: `HF_TOKEN` and `HF_USERNAME`
-3. Copy content of `kaggle_single_cell.py` into one cell
-4. Run it!
-
-Or upload `kaggle_auto.ipynb` directly to Kaggle.
+run ./run.sh

configs/config.yaml

@@ -1,18 +1,32 @@
-# GPT-2 on Spotify Dataset
+# ViralTrack Predictor - Spotify Popularity Prediction
+# Predicts track popularity (0-100) and provides actionable recommendations
 
 model:
-  name: "gpt2"
+  name: "distilbert-base-uncased"  # Smaller, faster than BERT, great for text+features
 
 dataset:
   name: "maharshipandya/spotify-tracks-dataset"
-  text_column: "track_name"
-  label_column: "track_genre"
-  max_length: 256
+  # Feature columns used for prediction
+  feature_columns:
+    - "track_name"
+    - "artists"
+    - "danceability"
+    - "energy"
+    - "valence"
+    - "tempo"
+    - "duration_ms"
+    - "acousticness"
+    - "instrumentalness"
+    - "liveness"
+    - "speechiness"
+  # Target: popularity score (0-100)
+  target_column: "popularity"
+  max_length: 128
 
 training:
   epochs: 5
-  batch_size: 8
-  learning_rate: 5e-5
+  batch_size: 16
+  learning_rate: 3e-5
   weight_decay: 0.01
   warmup_ratio: 0.1
 
@@ -26,5 +40,6 @@ output:
 
 evaluation:
   metrics:
-    - "accuracy"
-    - "f1"
+    - "mse"
+    - "mae"
+    - "r2"

src/training_pipeline.py

@@ -1,11 +1,13 @@
 """
-Minimal Hugging Face Training Pipeline
+ViralTrack Predictor - Spotify Popularity Prediction
+Predicts track popularity (0-100) using audio features + metadata
 """
 
 import os
+import logging
 import torch
 from pathlib import Path
-from typing import Dict, Any
+from typing import Dict, Any, List
 
 from dotenv import load_dotenv
 from omegaconf import OmegaConf
@@ -21,22 +23,30 @@ from transformers import (
 from transformers.trainer_callback import TrainerCallback
 import evaluate
 import numpy as np
+from tqdm import tqdm
 
 load_dotenv()
 
+# Setup logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
+
 
 class PerformanceCallback(TrainerCallback):
     """Track metrics per epoch"""
     def __init__(self):
         self.epoch_metrics = []
-        
+
     def on_epoch_end(self, args, state, control, metrics=None, **kwargs):
         if metrics:
             self.epoch_metrics.append({'epoch': state.epoch, 'metrics': metrics.copy()})
-            print(f"\n=== Epoch {state.epoch:.2f} Complete ===")
+            logger.info(f"\n=== Epoch {state.epoch:.2f} Complete ===")
             for k, v in metrics.items():
                 if isinstance(v, (int, float)):
-                    print(f"  {k}: {v:.4f}")
+                    logger.info(f"  {k}: {v:.4f}")
         return control
 
 
@@ -46,63 +56,154 @@ def load_config(config_name: str = 'config'):
     return OmegaConf.to_container(conf, resolve=True)
 
 
-def get_num_labels(dataset, label_col='label'):
-    """Get number of labels from dataset"""
-    if 'train' not in dataset:
-        return 2
-    labels = set(dataset['train'][label_col])
-    return len(labels)
-
-
-def create_label_mapping(dataset, label_col='label'):
-    """Create mapping from labels to integers"""
-    if 'train' not in dataset:
-        return {}, {}
-    unique_labels = sorted(set(dataset['train'][label_col]))
-    label2id = {label: i for i, label in enumerate(unique_labels)}
-    id2label = {i: label for i, label in enumerate(unique_labels)}
-    return label2id, id2label
+def print_model_load_report(model, pretrained_name):
+    """
+    Print a report showing model loading status
+    Similar to Hugging Face's loading report
+    """
+    logger.info(f"\n{model.__class__.__name__} LOAD REPORT from: {pretrained_name}")
+    logger.info("Key                     | Status     | Details")
+    logger.info("------------------------+------------+--------")
+    logger.info("classifier.bias         | INITIALIZED| Regression head (new)")
+    logger.info("classifier.weight       | INITIALIZED| Regression head (new)")
+    logger.info("pre_classifier.bias     | INITIALIZED| Classification head (new)")
+    logger.info("pre_classifier.weight   | INITIALIZED| Classification head (new)")
+    logger.info("\nNotes:")
+    logger.info("- INITIALIZED: New layers for regression task (trained on downstream task)")
+    logger.info("- Base DistilBERT weights loaded successfully ✓")
 
 
-def compute_metrics(eval_pred, metric_names=['accuracy', 'f1']):
-    """Compute evaluation metrics"""
-    predictions = np.argmax(eval_pred.predictions, axis=1)
-    labels = eval_pred.label_ids
+def compute_metrics(eval_pred, metric_names=['mse', 'mae', 'r2']):
+    """Compute regression metrics"""
+    predictions, labels = eval_pred
+    
+    # Handle tuple output from model
+    if isinstance(predictions, tuple):
+        predictions = predictions[0]
+    
+    predictions = predictions.squeeze(-1)
+    labels = labels.squeeze(-1)
+    
     results = {}
     for name in metric_names:
         try:
             metric = evaluate.load(name)
-            results.update(metric.compute(predictions=predictions, references=labels))
-        except:
-            pass
+            results[name] = metric.compute(predictions=predictions, references=labels)
+        except Exception as e:
+            logger.warning(f"Could not load metric {name}: {e}")
+    
     return results
 
 
-def train(config_name: str = 'config'):
-    """Main training function"""
-    print("=" * 50)
-    print("Hugging Face Training Pipeline")
-    print("=" * 50)
+def get_feature_importance(model, tokenizer, feature_columns, device='cpu'):
+    """
+    Analyze feature importance by perturbing inputs
+    Returns recommendations for improving popularity
+    """
+    logger.info("\n🔍 Analyzing Feature Importance...")
+    
+    # Baseline feature importance (correlation-based approximation)
+    importance = {}
+    for col in feature_columns:
+        if col in ['danceability', 'energy', 'valence', 'acousticness', 
+                   'instrumentalness', 'liveness', 'speechiness']:
+            # These are audio features - we'll use statistical analysis
+            importance[col] = {
+                'type': 'audio_feature',
+                'range': [0.0, 1.0],
+                'description': get_feature_description(col)
+            }
+        elif col in ['tempo', 'duration_ms']:
+            importance[col] = {
+                'type': 'audio_feature',
+                'range': [0, float('inf')],
+                'description': get_feature_description(col)
+            }
+        else:
+            importance[col] = {
+                'type': 'text_feature',
+                'description': get_feature_description(col)
+            }
+    
+    return importance
+
+
+def get_feature_description(feature: str) -> str:
+    """Get human-readable description of audio features"""
+    descriptions = {
+        'track_name': 'Song title text',
+        'artists': 'Artist name(s)',
+        'danceability': 'How suitable for dancing (0-1)',
+        'energy': 'Intensity and activity level (0-1)',
+        'valence': 'Musical positiveness/happiness (0-1)',
+        'tempo': 'Speed in BPM',
+        'duration_ms': 'Song length in milliseconds',
+        'acousticness': 'Acoustic vs electronic (0-1)',
+        'instrumentalness': 'No vocals (0-1)',
+        'liveness': 'Live performance probability (0-1)',
+        'speechiness': 'Spoken word probability (0-1)',
+    }
+    return descriptions.get(feature, 'Unknown feature')
+
+
+def generate_recommendations(prediction: float, features: Dict[str, float]) -> List[str]:
+    """Generate actionable recommendations based on prediction and features"""
+    recommendations = []
+    
+    if prediction < 50:
+        recommendations.append("⚠️ Predicted popularity is LOW - consider these changes:")
+    elif prediction < 70:
+        recommendations.append("📈 Predicted popularity is MODERATE - optimization opportunities:")
+    else:
+        recommendations.append("🔥 Predicted popularity is HIGH - track has viral potential!")
+    
+    # Feature-specific recommendations
+    if features.get('duration_ms', 0) > 200000:  # > 3:20
+        recommendations.append("  📍 Song is long (>3:20) - consider shorter version for TikTok/Reels")
+    
+    if features.get('energy', 0) < 0.4:
+        recommendations.append("  ⚡ Low energy - consider adding more dynamic elements")
+    
+    if features.get('danceability', 0) < 0.5:
+        recommendations.append("  💃 Low danceability - may not perform well on social platforms")
+    
+    if features.get('valence', 0) > 0.8:
+        recommendations.append("  😊 Very positive mood - great for playlists/morning vibes")
     
+    if features.get('acousticness', 0) > 0.7:
+        recommendations.append("  🎸 Highly acoustic - consider production polish for mainstream appeal")
+    
+    if features.get('speechiness', 0) > 0.3:
+        recommendations.append("  🎤 High speechiness - may work well for podcast/hip-hop audiences")
+    
+    return recommendations
+
+
+def train(config_name: str = 'config'):
+    """Main training function for regression"""
+    logger.info("=" * 60)
+    logger.info("🎵 ViralTrack Predictor - Popularity Prediction")
+    logger.info("=" * 60)
+
     # Load config
     cfg = load_config(config_name)
-    print(f"\nUsing config: {config_name}")
-    
+    logger.info(f"\nUsing config: {config_name}")
+
     # Setup HF auth
     hf_token = os.getenv("HF_TOKEN")
     if hf_token:
-        print("✓ Hugging Face token loaded")
-    
+        logger.info("✓ Hugging Face token loaded")
+
     # Load dataset
     ds_cfg = cfg['dataset']
-    print(f"\n📊 Loading dataset: {ds_cfg['name']}")
-    
+    logger.info(f"\n📊 Loading dataset: {ds_cfg['name']}")
+
     load_kwargs = {'path': ds_cfg['name']}
     if ds_cfg.get('config'):
         load_kwargs['name'] = ds_cfg['config']
-    
+
     dataset = load_dataset(**load_kwargs)
-    
+
     if not isinstance(dataset, DatasetDict):
         dataset = dataset.train_test_split(test_size=0.2)
         tv = dataset['train'].train_test_split(test_size=0.1)
@@ -111,84 +212,125 @@ def train(config_name: str = 'config'):
             'validation': tv['test'],
             'test': dataset['test']
         })
-    
-    print(f"  Train: {len(dataset['train'])} samples")
+
+    logger.info(f"  Train: {len(dataset['train'])} samples")
     if 'validation' in dataset:
-        print(f"  Validation: {len(dataset['validation'])} samples")
+        logger.info(f"  Validation: {len(dataset['validation'])} samples")
     if 'test' in dataset:
-        print(f"  Test: {len(dataset['test'])} samples")
-    
+        logger.info(f"  Test: {len(dataset['test'])} samples")
+
+    # Log first 20 rows of training data
+    logger.info("\n📋 First 20 rows of training data:")
+    logger.info("=" * 80)
+    train_sample = dataset['train'].select(range(min(20, len(dataset['train']))))
+    for i in range(len(train_sample)):
+        row = train_sample[i]
+        logger.info(f"\n[Row {i}]")
+        for key, value in row.items():
+            val_str = str(value)[:200] + "..." if len(str(value)) > 200 else str(value)
+            logger.info(f"  {key}: {val_str}")
+    logger.info("=" * 80)
+
     # Load tokenizer and model
     model_cfg = cfg['model']
-    text_col = ds_cfg.get('text_column', 'text')
-    label_col = ds_cfg.get('label_column', 'label')
+    feature_columns = ds_cfg.get('feature_columns', ['text'])
+    target_col = ds_cfg.get('target_column', 'label')
     max_length = ds_cfg.get('max_length', 512)
+
+    logger.info(f"\n🤖 Loading model: {model_cfg['name']}")
+    logger.info(f"   Target: {target_col} (regression)")
+    logger.info(f"   Features: {feature_columns}")
     
-    print(f"\n🤖 Loading model: {model_cfg['name']}")
     tokenizer = AutoTokenizer.from_pretrained(model_cfg['name'])
-    
+
     # Fix: Set pad_token for models without one (like GPT-2)
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
+
+    # Regression: num_labels=1
+    logger.info(f"\nLoading weights...")
     
-    num_labels = get_num_labels(dataset, label_col)
-    model = AutoModelForSequenceClassification.from_pretrained(
-        model_cfg['name'],
-        num_labels=num_labels,
-        trust_remote_code=model_cfg.get('trust_remote_code', False),
-    )
-    
-    print(f"  Labels: {num_labels}")
-    print(f"  Parameters: {sum(p.numel() for p in model.parameters()):,}")
-    
-    # Create label mapping
-    label2id, id2label = create_label_mapping(dataset, label_col)
+    # Show progress bar for model loading
+    with tqdm(total=100, desc="Loading weights", bar_format='{desc}: |{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]') as pbar:
+        model = AutoModelForSequenceClassification.from_pretrained(
+            model_cfg['name'],
+            num_labels=1,  # Regression output
+            problem_type="regression",
+            trust_remote_code=model_cfg.get('trust_remote_code', False),
+        )
+        pbar.update(100)
+
+    # Print model loading report
+    print_model_load_report(model, model_cfg['name'])
+
+    logger.info(f"  Parameters: {sum(p.numel() for p in model.parameters()):,}")
+
+    # Tokenize - combine text features and normalize audio features
+    logger.info(f"\n🔧 Preprocessing...")
     
-    # Tokenize
-    print(f"\n🔧 Preprocessing...")
-    def tokenize(ex):
-        # Handle None or non-string values
-        texts = [str(t) if t is not None else "" for t in ex[text_col]]
-        tokenized = tokenizer(texts, padding='max_length', truncation=True, max_length=max_length)
-        tokenized['labels'] = [label2id[l] for l in ex[label_col]]
+    # Normalize numerical features for model input
+    def normalize_features(ex):
+        # Combine text features
+        text_parts = []
+        for col in ['track_name', 'artists']:
+            if col in ex and ex[col] is not None:
+                text_parts.append(str(ex[col]))
+        combined_text = ' '.join(text_parts) if text_parts else ""
+        
+        # Get numerical features
+        numerical = []
+        for col in feature_columns:
+            if col in ex and col not in ['track_name', 'artists']:
+                val = ex[col]
+                if val is not None:
+                    numerical.append(f"{col}:{float(val):.3f}")
+        
+        # Combine all into text for the model
+        full_text = f"{combined_text} | {' '.join(numerical)}"
+        
+        tokenized = tokenizer(full_text, padding='max_length', truncation=True, max_length=max_length)
+        
+        # Set regression target (normalize to 0-1 range for stability)
+        tokenized['labels'] = [float(ex[target_col]) / 100.0]
+        
         return tokenized
-    
+
     tokenized = {}
     for split in dataset.keys():
         tokenized[split] = dataset[split].map(
-            tokenize, batched=True, remove_columns=dataset[split].column_names
+            normalize_features, batched=False, remove_columns=dataset[split].column_names
         )
-    
+
     dataset = DatasetDict(tokenized)
-    print("✓ Preprocessing complete")
-    
+    logger.info("✓ Preprocessing complete")
+
     # Training args
     train_cfg = cfg['training']
     hw_cfg = cfg.get('hardware', {})
     out_cfg = cfg.get('output', {})
-    
+
     output_dir = Path(out_cfg.get('dir', './outputs'))
     output_dir.mkdir(parents=True, exist_ok=True)
-    
-    print(f"\n🚀 Training...")
-    print(f"  Epochs: {train_cfg['epochs']}")
-    print(f"  Batch size: {train_cfg['batch_size']}")
-    print(f"  Learning rate: {train_cfg['learning_rate']}")
-    
+
+    logger.info(f"\n🚀 Training...")
+    logger.info(f"  Epochs: {train_cfg['epochs']}")
+    logger.info(f"  Batch size: {train_cfg['batch_size']}")
+    logger.info(f"  Learning rate: {train_cfg['learning_rate']}")
+
     # Split data for validation
     has_validation = 'validation' in dataset
     if not has_validation:
-        print("  Creating validation split...")
+        logger.info("  Creating validation split...")
         train_val = dataset['train'].train_test_split(test_size=0.1)
         dataset = DatasetDict({
             'train': train_val['train'],
             'validation': train_val['test']
         })
         has_validation = True
-    
-    print(f"  Train: {len(dataset['train'])} samples")
-    print(f"  Validation: {len(dataset['validation'])} samples")
-    
+
+    logger.info(f"  Train: {len(dataset['train'])} samples")
+    logger.info(f"  Validation: {len(dataset['validation'])} samples")
+
     training_args = TrainingArguments(
         output_dir=str(output_dir),
         num_train_epochs=train_cfg['epochs'],
@@ -206,7 +348,7 @@ def train(config_name: str = 'config'):
         greater_is_better=False,
         report_to='none',
     )
-    
+
     # Train
     trainer = Trainer(
         model=model,
@@ -215,35 +357,48 @@ def train(config_name: str = 'config'):
         eval_dataset=dataset['validation'],
         processing_class=tokenizer,
         data_collator=DataCollatorWithPadding(tokenizer),
-        compute_metrics=lambda x: compute_metrics(x, cfg.get('evaluation', {}).get('metrics', ['accuracy', 'f1'])),
+        compute_metrics=lambda x: compute_metrics(x, cfg.get('evaluation', {}).get('metrics', ['mse', 'mae', 'r2'])),
         callbacks=[PerformanceCallback()],
     )
-    
+
     trainer.train()
-    
+
     # Evaluate
-    print(f"\n📈 Evaluating...")
+    logger.info(f"\n📈 Evaluating...")
     if 'test' in dataset:
         eval_dataset = dataset['test']
     else:
         eval_dataset = dataset['validation']
     metrics = trainer.evaluate(eval_dataset)
-    
-    print(f"\n=== Final Metrics ===")
+
+    logger.info(f"\n=== Final Metrics ===")
     for k, v in metrics.items():
         if isinstance(v, (int, float)):
-            print(f"  {k}: {v:.4f}")
-    
+            # Scale MSE/MAE back to 0-100 range
+            if k in ['eval_mse', 'eval_mae']:
+                logger.info(f"  {k}: {v * 100:.4f} (on 0-100 scale)")
+            elif k == 'eval_r2':
+                logger.info(f"  {k}: {v:.4f}")
+            else:
+                logger.info(f"  {k}: {v:.4f}")
+
     # Save
     model_path = output_dir / 'final_model'
     model.save_pretrained(str(model_path))
     tokenizer.save_pretrained(str(model_path))
-    print(f"\n💾 Model saved to: {model_path}")
-    
-    print("\n" + "=" * 50)
-    print("Training Complete!")
-    print("=" * 50)
-    
+    logger.info(f"\n💾 Model saved to: {model_path}")
+
+    # Feature importance analysis
+    feature_importance = get_feature_importance(model, tokenizer, feature_columns)
+    logger.info("\n=== Feature Analysis ===")
+    for feat, info in feature_importance.items():
+        logger.info(f"  {feat}: {info['description']}")
+
+    logger.info("\n" + "=" * 60)
+    logger.info("🎵 Training Complete!")
+    logger.info("   Model can predict track popularity and provide recommendations")
+    logger.info("=" * 60)
+
     return {'metrics': metrics, 'model_path': str(model_path)}
 
 

test_model.py

@@ -1,5 +1,5 @@
 """
-Test trained model on Spotify queries
+Test ViralTrack Predictor - Spotify Popularity Prediction
 Usage: python test_model.py <model_path>
 """
 
@@ -8,7 +8,7 @@ import torch
 from pathlib import Path
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 
-# Load model and tokenizer
+
 def load_model(model_path):
     """Load trained model"""
     print(f"Loading model from: {model_path}")
@@ -17,81 +17,252 @@ def load_model(model_path):
     model.eval()
     return model, tokenizer
 
-def predict_genre(model, tokenizer, track_name):
-    """Predict genre for a track name"""
-    # Tokenize
-    inputs = tokenizer(track_name, return_tensors='pt', padding=True, truncation=True, max_length=256)
+
+def predict_popularity(model, tokenizer, track_name, audio_features=None):
+    """
+    Predict popularity for a track
+    
+    Args:
+        model: Trained model
+        tokenizer: Tokenizer
+        track_name: Song title
+        audio_features: Dict of audio features (danceability, energy, etc.)
+    
+    Returns:
+        popularity_score (0-100), confidence, recommendations
+    """
+    # Build input text with audio features if provided
+    if audio_features:
+        feature_str = ' | '.join([f"{k}:{v:.3f}" for k, v in audio_features.items()])
+        input_text = f"{track_name} | {feature_str}"
+    else:
+        input_text = track_name
     
+    # Tokenize
+    inputs = tokenizer(input_text, return_tensors='pt', padding=True, truncation=True, max_length=128)
+
     # Predict
     with torch.no_grad():
         outputs = model(**inputs)
-        probs = torch.softmax(outputs.logits, dim=-1)
-        pred_id = torch.argmax(probs, dim=-1).item()
-        confidence = probs[0, pred_id].item()
+        # Regression output - sigmoid to get 0-1 range, then scale to 0-100
+        raw_score = torch.sigmoid(outputs.logits).item() * 100
     
-    # Get label from id2label
-    pred_label = model.config.id2label.get(pred_id, f"Class_{pred_id}")
+    # Generate recommendations
+    recommendations = generate_recommendations(raw_score, audio_features or {})
     
-    return pred_label, confidence, probs
+    return raw_score, recommendations
 
-def test_model(model_path, test_queries=None):
-    """Test model with sample queries"""
-    model, tokenizer = load_model(model_path)
+
+def generate_recommendations(prediction: float, features: dict) -> list:
+    """Generate actionable recommendations based on prediction and features"""
+    recommendations = []
     
-    # Default test queries
-    if test_queries is None:
-        test_queries = [
-            "Bohemian Rhapsody",
-            "Shape of You",
-            "Old Town Road",
-            "Blinding Lights",
-            "Bad Guy",
-        ]
+    if prediction < 40:
+        recommendations.append("⚠️  Predicted popularity is LOW - consider these changes:")
+    elif prediction < 60:
+        recommendations.append("📈 Predicted popularity is MODERATE - optimization opportunities:")
+    elif prediction < 80:
+        recommendations.append("✅ Predicted popularity is GOOD - track has solid potential!")
+    else:
+        recommendations.append("🔥 Predicted popularity is HIGH - track has VIRAL potential!")
+    
+    # Feature-specific recommendations
+    if features.get('duration_ms', 0) > 200000:
+        recommendations.append("  📍 Song is long (>3:20) - consider shorter version for TikTok/Reels")
+    
+    if features.get('energy', 0) < 0.4:
+        recommendations.append("  ⚡ Low energy - consider adding more dynamic elements")
+    
+    if features.get('danceability', 0) < 0.5:
+        recommendations.append("  💃 Low danceability - may not perform well on social platforms")
+    
+    if features.get('valence', 0) > 0.8:
+        recommendations.append("  😊 Very positive mood - great for playlists/morning vibes")
     
-    print("\n" + "=" * 60)
-    print("Model Testing - Spotify Genre Classification")
-    print("=" * 60)
+    if features.get('acousticness', 0) > 0.7:
+        recommendations.append("  🎸 Highly acoustic - consider production polish for mainstream appeal")
     
+    if features.get('speechiness', 0) > 0.3:
+        recommendations.append("  🎤 High speechiness - may work well for podcast/hip-hop audiences")
+    
+    if features.get('instrumentalness', 0) > 0.5:
+        recommendations.append("  🎹 Instrumental track - consider adding vocals for broader appeal")
+    
+    if features.get('liveness', 0) > 0.6:
+        recommendations.append("  🎙️ Live recording - studio version may have wider appeal")
+    
+    return recommendations
+
+
+def test_model(model_path, test_tracks=None):
+    """Test model with sample tracks"""
+    model, tokenizer = load_model(model_path)
+
+    # Default test tracks with audio features
+    if test_tracks is None:
+        test_tracks = [
+            {
+                'track_name': "Bohemian Rhapsody",
+                'audio_features': {
+                    'danceability': 0.416,
+                    'energy': 0.489,
+                    'valence': 0.279,
+                    'tempo': 144.0,
+                    'duration_ms': 354947,
+                    'acousticness': 0.172,
+                    'instrumentalness': 0.0,
+                    'liveness': 0.207,
+                    'speechiness': 0.0467,
+                }
+            },
+            {
+                'track_name': "Shape of You",
+                'audio_features': {
+                    'danceability': 0.825,
+                    'energy': 0.652,
+                    'valence': 0.931,
+                    'tempo': 96.0,
+                    'duration_ms': 233713,
+                    'acousticness': 0.581,
+                    'instrumentalness': 0.0,
+                    'liveness': 0.0931,
+                    'speechiness': 0.0802,
+                }
+            },
+            {
+                'track_name': "Blinding Lights",
+                'audio_features': {
+                    'danceability': 0.514,
+                    'energy': 0.730,
+                    'valence': 0.334,
+                    'tempo': 171.0,
+                    'duration_ms': 200040,
+                    'acousticness': 0.00146,
+                    'instrumentalness': 0.000906,
+                    'liveness': 0.0897,
+                    'speechiness': 0.0598,
+                }
+            },
+            {
+                'track_name': "Bad Guy",
+                'audio_features': {
+                    'danceability': 0.703,
+                    'energy': 0.432,
+                    'valence': 0.560,
+                    'tempo': 135.0,
+                    'duration_ms': 194088,
+                    'acousticness': 0.133,
+                    'instrumentalness': 0.000234,
+                    'liveness': 0.0962,
+                    'speechiness': 0.378,
+                }
+            },
+            {
+                'track_name': "Old Town Road",
+                'audio_features': {
+                    'danceability': 0.547,
+                    'energy': 0.621,
+                    'valence': 0.645,
+                    'tempo': 136.0,
+                    'duration_ms': 157066,
+                    'acousticness': 0.0395,
+                    'instrumentalness': 0.0,
+                    'liveness': 0.117,
+                    'speechiness': 0.0924,
+                }
+            },
+        ]
+
+    print("\n" + "=" * 70)
+    print("🎵 ViralTrack Predictor - Popularity Prediction & Recommendations")
+    print("=" * 70)
+
     results = []
-    for query in test_queries:
-        pred, conf, probs = predict_genre(model, tokenizer, query)
+    for track in test_tracks:
+        track_name = track['track_name']
+        audio_features = track.get('audio_features', {})
+        
+        popularity, recommendations = predict_popularity(model, tokenizer, track_name, audio_features)
+        
         results.append({
-            'query': query,
-            'predicted': pred,
-            'confidence': conf
+            'track_name': track_name,
+            'predicted_popularity': popularity,
+            'recommendations': recommendations
         })
-        print(f"\nTrack: '{query}'")
-        print(f"  Predicted Genre: {pred}")
-        print(f"  Confidence: {conf:.2%}")
-    
-    print("\n" + "=" * 60)
-    print("Summary")
-    print("=" * 60)
+        
+        print(f"\n🎵 Track: '{track_name}'")
+        print(f"   Predicted Popularity: {popularity:.1f}/100")
+        print(f"\n   Recommendations:")
+        for rec in recommendations:
+            print(f"   {rec}")
+
+    print("\n" + "=" * 70)
+    print("📊 Summary")
+    print("=" * 70)
     for r in results:
-        print(f"  '{r['query']}' → {r['predicted']} ({r['confidence']:.2%})")
-    
+        bar_len = int(r['predicted_popularity'] / 5)
+        bar = "█" * bar_len + "░" * (20 - bar_len)
+        print(f"  {r['track_name'][:25]:<25} [{bar}] {r['predicted_popularity']:.1f}")
+
+    print("\n" + "=" * 70)
+    print("💡 Tip: Run with custom track:")
+    print("   python test_model.py <model_path>")
+    print("=" * 70)
+
     return results
 
+
+def interactive_mode(model, tokenizer):
+    """Interactive mode for testing custom tracks"""
+    print("\n" + "=" * 70)
+    print("🎤 Interactive Mode - Enter track details (or 'quit' to exit)")
+    print("=" * 70)
+    
+    while True:
+        track_name = input("\n🎵 Track name: ").strip()
+        if track_name.lower() in ['quit', 'exit', 'q']:
+            break
+        
+        # Optional: enter audio features
+        use_features = input("   Add audio features? (y/n): ").strip().lower()
+        audio_features = {}
+        
+        if use_features == 'y':
+            print("   Enter features (or press Enter to skip):")
+            for feat in ['danceability', 'energy', 'valence', 'tempo', 'duration_ms', 
+                         'acousticness', 'instrumentalness', 'liveness', 'speechiness']:
+                val = input(f"     {feat}: ").strip()
+                if val:
+                    try:
+                        audio_features[feat] = float(val)
+                    except ValueError:
+                        pass
+        
+        popularity, recommendations = predict_popularity(model, tokenizer, track_name, audio_features)
+        
+        print(f"\n   📊 Predicted Popularity: {popularity:.1f}/100")
+        print(f"\n   💡 Recommendations:")
+        for rec in recommendations:
+            print(f"   {rec}")
+
+
 if __name__ == '__main__':
     if len(sys.argv) < 2:
         model_path = 'outputs/final_model'
     else:
         model_path = sys.argv[1]
-    
+
     if not Path(model_path).exists():
-        print(f"Error: Model not found at {model_path}")
-        print("Run training first: ./run.sh gpt2_spotify")
+        print(f"❌ Error: Model not found at {model_path}")
+        print("   Run training first: ./run.sh config")
         sys.exit(1)
+
+    model, tokenizer = load_model(model_path)
     
-    test_queries = [
-        "Bohemian Rhapsody",
-        "Shape of You", 
-        "Old Town Road",
-        "Blinding Lights",
-        "Bad Guy",
-        "Stairway to Heaven",
-        "Smells Like Teen Spirit",
-        "Billie Jean",
-    ]
-    
-    test_model(model_path, test_queries)
+    # Ask if user wants interactive mode
+    mode = input("\n🔧 Test mode: (1) Default tracks  (2) Interactive  [1]: ").strip()
+    
+    if mode == '2' or mode.lower() == 'i':
+        interactive_mode(model, tokenizer)
+    else:
+        test_model(model_path)