new optimised code

README.md

@@ -1 +1,23 @@
-run ./run.sh
+run ./run.sh
+
+
+
+      1 # Full training (default from config.yaml)
+      2 ./run.sh
+      3 
+      4 # Fast test cycle (90% faster)
+      5 ./run.sh --epochs 1 --batch_size 64 --num_samples 1000
+      6 
+      7 # Custom configurations
+      8 ./run.sh --epochs 3 --batch_size 32
+      9 ./run.sh --num_samples 5000
+     10 ./run.sh --epochs 1  # Just test with 1 epoch
+
+    Speed Comparison
+
+
+    ┌──────────────────────────────────────────────────────┬──────────────────┐
+    │ Command                                              │ Time Estimate    │
+    ├──────────────────────────────────────────────────────┼──────────────────┤
+    │ ./run.sh                                             │ ~10 hours (full) │
+    │ ./run.sh --epochs 1 --batch_size 64 --num_samples 1000 │ ~5-10 minutes ⚡ │

configs/config.yaml

@@ -22,6 +22,7 @@ dataset:
   # Target: popularity score (0-100)
   target_column: "popularity"
   max_length: 128
+  # num_samples: 1000  # Uncomment to use subset (for faster testing)
 
 training:
   epochs: 5
@@ -34,7 +35,7 @@ hardware:
   mixed_precision: "fp16"
 
 output:
-  dir: "./outputs"
+  dir: "./model"
   save_strategy: "epoch"
   logging_steps: 10
 

run.sh

@@ -1,23 +1,51 @@
 #!/bin/bash
-# Usage: ./run.sh
+# Usage: ./run.sh [--epochs N] [--batch_size N] [--num_samples N]
+# Example: ./run.sh --epochs 1 --batch_size 64 --num_samples 1000
 
 set -e
 
 GREEN='\033[0;32m'
-YELLOW='\033[1;33m'
 BLUE='\033[0;34m'
+YELLOW='\033[1;33m'
 NC='\033[0m'
 
 cd "$(dirname "${BASH_SOURCE[0]}")"
 
+# Default values
+EPOCHS=""
+BATCH_SIZE=""
+NUM_SAMPLES=""
+
+# Parse arguments
+while [[ $# -gt 0 ]]; do
+    case $1 in
+        --epochs)
+            EPOCHS="--epochs $2"
+            shift 2
+            ;;
+        --batch_size)
+            BATCH_SIZE="--batch_size $2"
+            shift 2
+            ;;
+        --num_samples)
+            NUM_SAMPLES="--num_samples $2"
+            shift 2
+            ;;
+        *)
+            echo -e "${YELLOW}Unknown option: $1${NC}"
+            exit 1
+            ;;
+    esac
+done
+
 # Install uv if not exists
 if ! command -v uv &> /dev/null; then
-    echo -e "${YELLOW}Installing uv...${NC}"
+    echo -e "${GREEN}Installing uv...${NC}"
     curl -LsSf https://astral.sh/uv/install.sh | sh
 fi
 
-# Sync dependencies (creates .venv if needed)
-echo -e "${GREEN}✓ Syncing dependencies with uv...${NC}"
+# Sync dependencies
+echo -e "${GREEN}✓ Syncing dependencies...${NC}"
 uv sync
 
 # Activate venv
@@ -28,32 +56,43 @@ if [ ! -f ".env" ]; then
     cp .env.example .env
 fi
 
-# Check GPU
-echo ""
-echo -e "${YELLOW}Checking GPU availability...${NC}"
+# Quick GPU check
 python3 -c "
 import torch
 if torch.cuda.is_available():
-    print(f'✓ CUDA available: {torch.cuda.device_count()} GPU(s)')
-    for i in range(torch.cuda.device_count()):
-        print(f'  GPU {i}: {torch.cuda.get_device_name(i)}')
+    print(f'✓ GPU: {torch.cuda.get_device_name(0)}')
 else:
-    print('⚠ No CUDA available - training will use CPU (slower)')
+    print('⚠ CPU mode (slower)')
 "
 
-echo -e "${BLUE}========================================${NC}"
-echo -e "${BLUE}  Training GPT-2 on Spotify Dataset${NC}"
-echo -e "${BLUE}========================================${NC}"
+# Build training command
+TRAIN_CMD="python3 src/training_pipeline.py config"
+if [ -n "$EPOCHS" ]; then
+    TRAIN_CMD="$TRAIN_CMD $EPOCHS"
+fi
+if [ -n "$BATCH_SIZE" ]; then
+    TRAIN_CMD="$TRAIN_CMD $BATCH_SIZE"
+fi
+if [ -n "$NUM_SAMPLES" ]; then
+    TRAIN_CMD="$TRAIN_CMD $NUM_SAMPLES"
+fi
+
+echo -e "${BLUE}╔════════════════════════════════════════╗${NC}"
+echo -e "${BLUE}║  Training GPT-2 on Spotify Dataset     ║${NC}"
+echo -e "${BLUE}╚════════════════════════════════════════╝${NC}"
+echo ""
+echo -e "${YELLOW}Command: $TRAIN_CMD${NC}"
+echo ""
 
-python3 src/training_pipeline.py config
+eval $TRAIN_CMD
 
 echo ""
 echo -e "${GREEN}✓ Training Complete!${NC}"
-echo -e "Model saved to: ${YELLOW}outputs/final_model${NC}"
+echo -e "  Model: ${BLUE}./model${NC}"
 echo ""
 
-echo -e "${BLUE}========================================${NC}"
-echo -e "${BLUE}  Testing Model${NC}"
-echo -e "${BLUE}========================================${NC}"
+echo -e "${BLUE}╔════════════════════════════════════════╗${NC}"
+echo -e "${BLUE}║  Testing Model                         ║${NC}"
+echo -e "${BLUE}╚════════════════════════════════════════╝${NC}"
 
-python3 test_model.py outputs/final_model
+python3 test_model.py ./model

src/training_pipeline.py

@@ -25,28 +25,43 @@ import evaluate
 import numpy as np
 from tqdm import tqdm
 
+# Suppress HTTP logs from transformers/datasets
+logging.getLogger("filelock").setLevel(logging.ERROR)
+logging.getLogger("urllib3").setLevel(logging.ERROR)
+logging.getLogger("huggingface_hub").setLevel(logging.ERROR)
+logging.getLogger("datasets").setLevel(logging.ERROR)
+
 load_dotenv()
 
-# Setup logging
+# Setup logging - cleaner format
 logging.basicConfig(
-    level=logging.INFO,
-    format='%(asctime)s - %(levelname)s - %(message)s'
+    level=logging.ERROR,
+    format='%(message)s',
+    handlers=[logging.StreamHandler()]
 )
 logger = logging.getLogger(__name__)
 
 
 class PerformanceCallback(TrainerCallback):
-    """Track metrics per epoch"""
+    """Track metrics per epoch with clean output"""
     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()})
-            logger.info(f"\n=== Epoch {state.epoch:.2f} Complete ===")
-            for k, v in metrics.items():
-                if isinstance(v, (int, float)):
-                    logger.info(f"  {k}: {v:.4f}")
+            # Clean epoch summary
+            print(f"\n{'='*50}")
+            print(f"✅ Epoch {state.epoch:.0f}/{args.num_train_epochs:.0f} Complete")
+            print(f"{'='*50}")
+            key_metrics = ['loss', 'mae', 'r2']
+            for k in key_metrics:
+                full_key = f'eval_{k}' if k != 'loss' else k
+                if full_key in metrics:
+                    val = metrics[full_key]
+                    if isinstance(val, (int, float)):
+                        print(f"  {k.upper():<15} {val:.4f}")
+            print(f"{'='*50}\n")
         return control
 
 
@@ -57,41 +72,29 @@ def load_config(config_name: str = 'config'):
 
 
 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 ✓")
+    """Clean model loading status"""
+    print(f"\n📦 Model: {model.__class__.__name__}")
+    print(f"   Source: {pretrained_name}")
+    print(f"   Params: {sum(p.numel() for p in model.parameters()):,}")
+    print(f"   ✓ Base weights loaded")
+    print(f"   ✓ New regression head added\n")
 
 
 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[name] = metric.compute(predictions=predictions, references=labels)
-        except Exception as e:
-            logger.warning(f"Could not load metric {name}: {e}")
-    
+        metric = evaluate.load(name)
+        results[name] = metric.compute(predictions=predictions, references=labels)
+
     return results
 
 
@@ -179,24 +182,32 @@ def generate_recommendations(prediction: float, features: Dict[str, float]) -> L
     return recommendations
 
 
-def train(config_name: str = 'config'):
+def train(config_name: str = 'config', epochs: int = None, batch_size: int = None, num_samples: int = None):
     """Main training function for regression"""
-    logger.info("=" * 60)
-    logger.info("🎵 ViralTrack Predictor - Popularity Prediction")
-    logger.info("=" * 60)
+    print(f"\n{'🎵'*30}")
+    print("  VIRALTRACK PREDICTOR - Spotify Popularity Prediction")
+    print(f"{'🎵'*30}\n")
 
     # Load config
     cfg = load_config(config_name)
-    logger.info(f"\nUsing config: {config_name}")
+    print(f"📋 Config: {config_name}\n")
+
+    # Override config with CLI args if provided
+    if epochs is not None:
+        cfg['training']['epochs'] = epochs
+    if batch_size is not None:
+        cfg['training']['batch_size'] = batch_size
+    if num_samples is not None:
+        cfg['dataset']['num_samples'] = num_samples
 
     # Setup HF auth
     hf_token = os.getenv("HF_TOKEN")
     if hf_token:
-        logger.info("✓ Hugging Face token loaded")
+        print("✓ Hugging Face token loaded\n")
 
     # Load dataset
     ds_cfg = cfg['dataset']
-    logger.info(f"\n📊 Loading dataset: {ds_cfg['name']}")
+    print(f"📊 Dataset: {ds_cfg['name']}")
 
     load_kwargs = {'path': ds_cfg['name']}
     if ds_cfg.get('config'):
@@ -213,23 +224,23 @@ def train(config_name: str = 'config'):
             'test': dataset['test']
         })
 
-    logger.info(f"  Train: {len(dataset['train'])} samples")
+    # Subsample if requested
+    num_samples = ds_cfg.get('num_samples')
+    if num_samples is not None:
+        print(f"⚡ Using subset: {num_samples} samples (for faster testing)")
+        if len(dataset['train']) > num_samples:
+            dataset['train'] = dataset['train'].select(range(num_samples))
+        if 'validation' in dataset and len(dataset['validation']) > num_samples // 10:
+            dataset['validation'] = dataset['validation'].select(range(min(num_samples // 10, len(dataset['validation']))))
+        if 'test' in dataset and len(dataset['test']) > num_samples // 10:
+            dataset['test'] = dataset['test'].select(range(min(num_samples // 10, len(dataset['test']))))
+
+    print(f"   ├─ Train:      {len(dataset['train']):,} samples")
     if 'validation' in dataset:
-        logger.info(f"  Validation: {len(dataset['validation'])} samples")
+        print(f"   ├─ Validation: {len(dataset['validation']):,} samples")
     if 'test' in dataset:
-        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)
+        print(f"   └─ Test:       {len(dataset['test']):,} samples")
+    print()
 
     # Load tokenizer and model
     model_cfg = cfg['model']
@@ -237,38 +248,33 @@ def train(config_name: str = 'config'):
     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"🤖 Model: {model_cfg['name']}")
+    print(f"   Target: {target_col} (regression)")
+    print(f"   Features: {len(feature_columns)} columns\n")
+
+    print("⏳ Loading tokenizer...")
     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...")
-    
-    # Show progress bar for model loading
+    print("⏳ Loading model weights...\n")
+
     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
+            num_labels=1,
             problem_type="regression",
             trust_remote_code=model_cfg.get('trust_remote_code', False),
+            ignore_mismatched_sizes=True,
         )
         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...")
-    
-    # Normalize numerical features for model input
+    print("🔧 Preprocessing data...")
+
     def normalize_features(ex):
         # Combine text features
         text_parts = []
@@ -302,7 +308,7 @@ def train(config_name: str = 'config'):
         )
 
     dataset = DatasetDict(tokenized)
-    logger.info("✓ Preprocessing complete")
+    print("✓ Preprocessing complete\n")
 
     # Training args
     train_cfg = cfg['training']
@@ -312,15 +318,19 @@ def train(config_name: str = 'config'):
     output_dir = Path(out_cfg.get('dir', './outputs'))
     output_dir.mkdir(parents=True, exist_ok=True)
 
-    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']}")
+    print(f"{'='*50}")
+    print("🚀 TRAINING CONFIGURATION")
+    print(f"{'='*50}")
+    print(f"  Epochs:        {train_cfg['epochs']}")
+    print(f"  Batch size:    {train_cfg['batch_size']}")
+    print(f"  Learning rate: {train_cfg['learning_rate']}")
+    print(f"  Output dir:    {output_dir}")
+    print(f"{'='*50}\n")
 
     # Split data for validation
     has_validation = 'validation' in dataset
     if not has_validation:
-        logger.info("  Creating validation split...")
+        print("  Creating validation split...")
         train_val = dataset['train'].train_test_split(test_size=0.1)
         dataset = DatasetDict({
             'train': train_val['train'],
@@ -328,8 +338,8 @@ def train(config_name: str = 'config'):
         })
         has_validation = True
 
-    logger.info(f"  Train: {len(dataset['train'])} samples")
-    logger.info(f"  Validation: {len(dataset['validation'])} samples")
+    print(f"📈 Training: {len(dataset['train']):,} samples")
+    print(f"   Validating: {len(dataset['validation']):,} samples\n")
 
     training_args = TrainingArguments(
         output_dir=str(output_dir),
@@ -347,9 +357,11 @@ def train(config_name: str = 'config'):
         metric_for_best_model='loss',
         greater_is_better=False,
         report_to='none',
+        disable_tqdm=False,
     )
 
     # Train
+    print("⏳ Starting training...\n")
     trainer = Trainer(
         model=model,
         args=training_args,
@@ -364,45 +376,57 @@ def train(config_name: str = 'config'):
     trainer.train()
 
     # Evaluate
-    logger.info(f"\n📈 Evaluating...")
+    print(f"\n{'='*50}")
+    print("📈 EVALUATION")
+    print(f"{'='*50}")
     if 'test' in dataset:
         eval_dataset = dataset['test']
     else:
         eval_dataset = dataset['validation']
     metrics = trainer.evaluate(eval_dataset)
 
-    logger.info(f"\n=== Final Metrics ===")
+    print(f"\n=== Final Metrics ===")
     for k, v in metrics.items():
         if isinstance(v, (int, float)):
-            # 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)")
+                print(f"  {k:<15} {v * 100:.4f} (on 0-100 scale)")
             elif k == 'eval_r2':
-                logger.info(f"  {k}: {v:.4f}")
+                print(f"  {k:<15} {v:.4f}")
             else:
-                logger.info(f"  {k}: {v:.4f}")
+                print(f"  {k:<15} {v:.4f}")
+    print(f"{'='*50}\n")
 
     # Save
-    model_path = output_dir / 'final_model'
+    model_path = output_dir  # Save directly to output_dir (e.g., ./model)
     model.save_pretrained(str(model_path))
     tokenizer.save_pretrained(str(model_path))
-    logger.info(f"\n💾 Model saved to: {model_path}")
+    print(f"💾 Model saved to: {model_path}\n")
 
     # Feature importance analysis
     feature_importance = get_feature_importance(model, tokenizer, feature_columns)
-    logger.info("\n=== Feature Analysis ===")
+    print(f"{'='*50}")
+    print("📊 FEATURE ANALYSIS")
+    print(f"{'='*50}")
     for feat, info in feature_importance.items():
-        logger.info(f"  {feat}: {info['description']}")
+        print(f"  {feat}: {info['description']}")
+    print(f"{'='*50}\n")
 
-    logger.info("\n" + "=" * 60)
-    logger.info("🎵 Training Complete!")
-    logger.info("   Model can predict track popularity and provide recommendations")
-    logger.info("=" * 60)
+    print(f"{'🎵'*30}")
+    print("  ✅ TRAINING COMPLETE!")
+    print(f"{'🎵'*30}")
+    print("   Model can predict track popularity and provide recommendations\n")
 
     return {'metrics': metrics, 'model_path': str(model_path)}
 
 
 if __name__ == '__main__':
-    import sys
-    config = sys.argv[1] if len(sys.argv) > 1 else 'config'
-    train(config)
+    import argparse
+    
+    parser = argparse.ArgumentParser(description='Train ViralTrack Predictor')
+    parser.add_argument('config', nargs='?', default='config', help='Config file name (default: config)')
+    parser.add_argument('--epochs', type=int, default=None, help='Number of training epochs')
+    parser.add_argument('--batch_size', type=int, default=None, help='Training batch size')
+    parser.add_argument('--num_samples', type=int, default=None, help='Number of samples to use (for faster testing)')
+    
+    args = parser.parse_args()
+    train(args.config, epochs=args.epochs, batch_size=args.batch_size, num_samples=args.num_samples)

test_model.py

@@ -18,38 +18,54 @@ def load_model(model_path):
     return model, tokenizer
 
 
-def predict_popularity(model, tokenizer, track_name, audio_features=None):
+def predict_popularity(model, tokenizer, track_name, artists="", audio_features=None):
     """
     Predict popularity for a track
-    
+
     Args:
         model: Trained model
         tokenizer: Tokenizer
         track_name: Song title
+        artists: Artist name(s)
         audio_features: Dict of audio features (danceability, energy, etc.)
-    
+
     Returns:
-        popularity_score (0-100), confidence, recommendations
+        popularity_score (0-100), recommendations
     """
-    # Build input text with audio features if provided
+    # Build input text - SAME FORMAT AS TRAINING
+    text_parts = [track_name]
+    if artists:
+        text_parts.append(artists)
+    combined_text = ' '.join(text_parts)
+
+    # Add audio features in same format as training
+    numerical = []
     if audio_features:
-        feature_str = ' | '.join([f"{k}:{v:.3f}" for k, v in audio_features.items()])
-        input_text = f"{track_name} | {feature_str}"
+        for col, val in audio_features.items():
+            if col not in ['track_name', 'artists']:
+                numerical.append(f"{col}:{float(val):.3f}")
+
+    # Combine all into text for the model (matches training preprocessing)
+    if numerical:
+        input_text = f"{combined_text} | {' '.join(numerical)}"
     else:
-        input_text = track_name
-    
-    # Tokenize
+        input_text = combined_text
+
+    # Tokenize - SAME PARAMETERS AS TRAINING (max_length=128)
     inputs = tokenizer(input_text, return_tensors='pt', padding=True, truncation=True, max_length=128)
 
     # Predict
     with torch.no_grad():
         outputs = model(**inputs)
-        # Regression output - sigmoid to get 0-1 range, then scale to 0-100
-        raw_score = torch.sigmoid(outputs.logits).item() * 100
-    
+        # Regression output - already scaled to 0-1 during training, scale back to 0-100
+        raw_score = outputs.logits.item() * 100
+
+    # Clamp to valid range
+    raw_score = max(0, min(100, raw_score))
+
     # Generate recommendations
     recommendations = generate_recommendations(raw_score, audio_features or {})
-    
+
     return raw_score, recommendations
 
 
@@ -103,6 +119,7 @@ def test_model(model_path, test_tracks=None):
         test_tracks = [
             {
                 'track_name': "Bohemian Rhapsody",
+                'artists': "Queen",
                 'audio_features': {
                     'danceability': 0.416,
                     'energy': 0.489,
@@ -117,6 +134,7 @@ def test_model(model_path, test_tracks=None):
             },
             {
                 'track_name': "Shape of You",
+                'artists': "Ed Sheeran",
                 'audio_features': {
                     'danceability': 0.825,
                     'energy': 0.652,
@@ -131,6 +149,7 @@ def test_model(model_path, test_tracks=None):
             },
             {
                 'track_name': "Blinding Lights",
+                'artists': "The Weeknd",
                 'audio_features': {
                     'danceability': 0.514,
                     'energy': 0.730,
@@ -145,6 +164,7 @@ def test_model(model_path, test_tracks=None):
             },
             {
                 'track_name': "Bad Guy",
+                'artists': "Billie Eilish",
                 'audio_features': {
                     'danceability': 0.703,
                     'energy': 0.432,
@@ -159,6 +179,7 @@ def test_model(model_path, test_tracks=None):
             },
             {
                 'track_name': "Old Town Road",
+                'artists': "Lil Nas X",
                 'audio_features': {
                     'danceability': 0.547,
                     'energy': 0.621,
@@ -180,9 +201,10 @@ def test_model(model_path, test_tracks=None):
     results = []
     for track in test_tracks:
         track_name = track['track_name']
+        artists = track.get('artists', '')
         audio_features = track.get('audio_features', {})
-        
-        popularity, recommendations = predict_popularity(model, tokenizer, track_name, audio_features)
+
+        popularity, recommendations = predict_popularity(model, tokenizer, track_name, artists, audio_features)
         
         results.append({
             'track_name': track_name,
@@ -217,19 +239,21 @@ def interactive_mode(model, tokenizer):
     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
-        
+
+        artists = input("   Artists: ").strip()
+
         # 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', 
+            for feat in ['danceability', 'energy', 'valence', 'tempo', 'duration_ms',
                          'acousticness', 'instrumentalness', 'liveness', 'speechiness']:
                 val = input(f"     {feat}: ").strip()
                 if val:
@@ -237,8 +261,8 @@ def interactive_mode(model, tokenizer):
                         audio_features[feat] = float(val)
                     except ValueError:
                         pass
-        
-        popularity, recommendations = predict_popularity(model, tokenizer, track_name, audio_features)
+
+        popularity, recommendations = predict_popularity(model, tokenizer, track_name, artists, audio_features)
         
         print(f"\n   📊 Predicted Popularity: {popularity:.1f}/100")
         print(f"\n   💡 Recommendations:")
@@ -248,7 +272,7 @@ def interactive_mode(model, tokenizer):
 
 if __name__ == '__main__':
     if len(sys.argv) < 2:
-        model_path = 'outputs/final_model'
+        model_path = 'model'  # Default: ./model (current directory)
     else:
         model_path = sys.argv[1]