""" ShortSmith v2 - Trained Hype Scorer Uses the MLP model trained on Mr. HiSum dataset to score segments. Falls back to heuristic scoring if weights not available. """ import os from pathlib import Path from typing import Optional, List, Tuple import numpy as np import torch import torch.nn as nn from utils.logger import get_logger logger = get_logger("scoring.trained_scorer") class HypeScorerMLP(nn.Module): """ 2-layer MLP for hype scoring. Must match the architecture from training notebook. """ def __init__( self, visual_dim: int = 512, audio_dim: int = 13, hidden_dim: int = 256, dropout: float = 0.3, ): super().__init__() self.visual_dim = visual_dim self.audio_dim = audio_dim input_dim = visual_dim + audio_dim self.network = nn.Sequential( # Layer 1 nn.Linear(input_dim, hidden_dim), nn.BatchNorm1d(hidden_dim), nn.ReLU(), nn.Dropout(dropout), # Layer 2 nn.Linear(hidden_dim, hidden_dim // 2), nn.BatchNorm1d(hidden_dim // 2), nn.ReLU(), nn.Dropout(dropout), # Output layer nn.Linear(hidden_dim // 2, 1), ) def forward(self, features: torch.Tensor) -> torch.Tensor: """Forward pass with concatenated features.""" return self.network(features) class TrainedHypeScorer: """ Trained neural network hype scorer. Uses MLP trained on Mr. HiSum "Most Replayed" data. """ # Default weights path relative to project root DEFAULT_WEIGHTS_PATH = "weights/hype_scorer_weights.pt" def __init__( self, weights_path: Optional[str] = None, device: Optional[str] = None, visual_dim: int = 512, audio_dim: int = 13, ): """ Initialize trained scorer. Args: weights_path: Path to trained weights (.pt file) device: Device to run on (cuda/cpu/mps) visual_dim: Visual feature dimension audio_dim: Audio feature dimension """ self.visual_dim = visual_dim self.audio_dim = audio_dim self.model = None self.device = device or self._get_device() # Find weights file if weights_path is None: # Look in common locations candidates = [ self.DEFAULT_WEIGHTS_PATH, "hype_scorer_weights.pt", "weights/hype_scorer_weights.pt", os.path.join(os.path.dirname(__file__), "..", "weights", "hype_scorer_weights.pt"), ] for candidate in candidates: if os.path.exists(candidate): weights_path = candidate break if weights_path and os.path.exists(weights_path): self._load_model(weights_path) else: logger.warning( f"Trained weights not found. TrainedHypeScorer will use fallback scoring. " f"To use trained model, place weights at: {self.DEFAULT_WEIGHTS_PATH}" ) def _get_device(self) -> str: """Detect best available device.""" if torch.cuda.is_available(): return "cuda" elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): return "mps" return "cpu" def _load_model(self, weights_path: str) -> None: """Load trained model weights.""" try: logger.info(f"Loading trained hype scorer from {weights_path}") # Initialize model self.model = HypeScorerMLP( visual_dim=self.visual_dim, audio_dim=self.audio_dim, ) # Load weights state_dict = torch.load(weights_path, map_location=self.device) # Handle different save formats if isinstance(state_dict, dict) and "model_state_dict" in state_dict: state_dict = state_dict["model_state_dict"] self.model.load_state_dict(state_dict) self.model.to(self.device) self.model.eval() logger.info(f"✓ Trained hype scorer loaded successfully on {self.device}") except Exception as e: logger.error(f"Failed to load trained model: {e}") self.model = None @property def is_available(self) -> bool: """Check if trained model is loaded.""" return self.model is not None @torch.no_grad() def score( self, visual_features: np.ndarray, audio_features: np.ndarray, ) -> float: """ Score a single segment. Args: visual_features: Visual feature vector (visual_dim,) audio_features: Audio feature vector (audio_dim,) Returns: Hype score (0-1) """ if not self.is_available: return self._fallback_score(visual_features, audio_features) # Prepare input features = np.concatenate([visual_features, audio_features]) tensor = torch.tensor(features, dtype=torch.float32).unsqueeze(0).to(self.device) # Forward pass raw_score = self.model(tensor) # Normalize to 0-1 with sigmoid score = torch.sigmoid(raw_score).item() return score @torch.no_grad() def score_batch( self, visual_features: np.ndarray, audio_features: np.ndarray, ) -> np.ndarray: """ Score multiple segments in batch. Args: visual_features: Visual features (N, visual_dim) audio_features: Audio features (N, audio_dim) Returns: Array of hype scores (N,) """ if not self.is_available: return np.array([ self._fallback_score(visual_features[i], audio_features[i]) for i in range(len(visual_features)) ]) # Prepare batch input features = np.concatenate([visual_features, audio_features], axis=1) tensor = torch.tensor(features, dtype=torch.float32).to(self.device) # Forward pass raw_scores = self.model(tensor) # Normalize to 0-1 scores = torch.sigmoid(raw_scores).squeeze().cpu().numpy() return scores def _fallback_score( self, visual_features: np.ndarray, audio_features: np.ndarray, ) -> float: """ Fallback heuristic scoring when model not available. Uses similar logic to training data generation. """ # Visual contribution (mean of first 50 dims if available) visual_len = min(50, len(visual_features)) visual_score = np.mean(visual_features[:visual_len]) * 0.5 + 0.5 visual_score = np.clip(visual_score, 0, 1) # Audio contribution if len(audio_features) >= 8: audio_score = ( audio_features[0] * 0.4 + # RMS energy audio_features[5] * 0.3 + # Spectral flux (if available) audio_features[7] * 0.3 # Onset strength (if available) ) * 0.5 + 0.5 else: audio_score = np.mean(audio_features) * 0.5 + 0.5 audio_score = np.clip(audio_score, 0, 1) # Combined return float(0.5 * visual_score + 0.5 * audio_score) def compare_segments( self, visual_a: np.ndarray, audio_a: np.ndarray, visual_b: np.ndarray, audio_b: np.ndarray, ) -> int: """ Compare two segments. Returns: 1 if A is more engaging, -1 if B is more engaging, 0 if equal """ score_a = self.score(visual_a, audio_a) score_b = self.score(visual_b, audio_b) if score_a > score_b + 0.05: return 1 elif score_b > score_a + 0.05: return -1 return 0 # Singleton instance for easy access _trained_scorer: Optional[TrainedHypeScorer] = None def get_trained_scorer( weights_path: Optional[str] = None, force_reload: bool = False, ) -> TrainedHypeScorer: """ Get singleton trained scorer instance. Args: weights_path: Optional path to weights file force_reload: Force reload even if already loaded Returns: TrainedHypeScorer instance """ global _trained_scorer if _trained_scorer is None or force_reload: _trained_scorer = TrainedHypeScorer(weights_path=weights_path) return _trained_scorer __all__ = ["TrainedHypeScorer", "HypeScorerMLP", "get_trained_scorer"]