lklk

run.sh

@@ -94,5 +94,6 @@ echo ""
 echo -e "${BLUE}╔════════════════════════════════════════╗${NC}"
 echo -e "${BLUE}║  Testing Model                         ║${NC}"
 echo -e "${BLUE}╚════════════════════════════════════════╝${NC}"
+echo ""
 
 python3 test_model.py ./model

src/training_pipeline.py

@@ -21,7 +21,6 @@ from transformers import (
     DataCollatorWithPadding,
 )
 from transformers.trainer_callback import TrainerCallback
-import evaluate
 import numpy as np
 from tqdm import tqdm
 
@@ -30,6 +29,8 @@ 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)
+logging.getLogger("transformers").setLevel(logging.ERROR)
+logging.getLogger("torch").setLevel(logging.ERROR)
 
 load_dotenv()
 
@@ -71,17 +72,10 @@ def load_config(config_name: str = 'config'):
     return OmegaConf.to_container(conf, resolve=True)
 
 
-def print_model_load_report(model, pretrained_name):
-    """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"""
+    """Compute regression metrics using scikit-learn"""
+    from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
+    
     predictions, labels = eval_pred
 
     if isinstance(predictions, tuple):
@@ -90,10 +84,11 @@ def compute_metrics(eval_pred, metric_names=['mse', 'mae', 'r2']):
     predictions = predictions.squeeze(-1)
     labels = labels.squeeze(-1)
 
-    results = {}
-    for name in metric_names:
-        metric = evaluate.load(name)
-        results[name] = metric.compute(predictions=predictions, references=labels)
+    results = {
+        'mse': mean_squared_error(labels, predictions),
+        'mae': mean_absolute_error(labels, predictions),
+        'r2': r2_score(labels, predictions),
+    }
 
     return results
 
@@ -270,7 +265,10 @@ def train(config_name: str = 'config', epochs: int = None, batch_size: int = Non
         )
         pbar.update(100)
 
-    print_model_load_report(model, model_cfg['name'])
+    print(f"\n📦 Model: {model.__class__.__name__}")
+    print(f"   Source: {model_cfg['name']}")
+    print(f"   Params: {sum(p.numel() for p in model.parameters()):,}")
+    print(f"   ✓ Ready for training\n")
 
     # Tokenize - combine text features and normalize audio features
     print("🔧 Preprocessing data...")
@@ -358,6 +356,7 @@ def train(config_name: str = 'config', epochs: int = None, batch_size: int = Non
         greater_is_better=False,
         report_to='none',
         disable_tqdm=False,
+        dataloader_pin_memory=False,
     )
 
     # Train

test_model.py

@@ -4,17 +4,23 @@ Usage: python test_model.py <model_path>
 """
 
 import sys
+import logging
 import torch
 from pathlib import Path
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 
+# Suppress logs
+logging.getLogger("transformers").setLevel(logging.ERROR)
+logging.getLogger("torch").setLevel(logging.ERROR)
+
 
 def load_model(model_path):
     """Load trained model"""
-    print(f"Loading model from: {model_path}")
+    print(f"⏳ Loading model from: {model_path}")
     tokenizer = AutoTokenizer.from_pretrained(model_path)
     model = AutoModelForSequenceClassification.from_pretrained(model_path)
     model.eval()
+    print(f"✓ Model loaded\n")
     return model, tokenizer
 
 
@@ -282,11 +288,6 @@ if __name__ == '__main__':
         sys.exit(1)
 
     model, tokenizer = load_model(model_path)
-    
-    # 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)
+
+    # Run default test tracks automatically
+    test_model(model_path)