ref: refactor before chekout

src/music_engine/matcher.py

@@ -1,67 +1,83 @@
+import joblib
 import numpy as np
 import pandas as pd
 from pathlib import Path
-import joblib
 
 class MusicMatcher:
     def __init__(self, db_path: Path | str, model_path: Path | str):
-        # Загружаем твою новую, обогащенную базу
+        # Загрузка базы данных музыкальных произведений
         self.music_db = pd.read_csv(db_path)
         self.acoustic_features = ['energy', 'flux', 'centroid', 'pitch', 'hnr', 'zcr']
         
-        # Удаляем строки, где нет акустических фич
-        self.music_db = self.music_db.dropna(subset=['valence', 'arousal'] + self.acoustic_features)
+        # Удаление записей с пропущенными целевыми или акустическими признаками
+        target_columns = ['valence', 'arousal'] + self.acoustic_features
+        self.music_db = self.music_db.dropna(subset=target_columns)
         
-        # Нормализуем акустику от 0 до 1, чтобы сравнивать с ответом LLM
+        # Масштабирование акустических параметров к диапазону [0, 1]
         self.norm_db = self.music_db.copy()
         for feat in self.acoustic_features:
-            f_min, f_max = self.norm_db[feat].min(), self.norm_db[feat].max()
+            f_min = self.norm_db[feat].min()
+            f_max = self.norm_db[feat].max()
             if f_max > f_min:
                 self.norm_db[f"norm_{feat}"] = (self.norm_db[feat] - f_min) / (f_max - f_min)
             else:
                 self.norm_db[f"norm_{feat}"] = 0.0
 
+        # Определение путей к аудиофайлам и загрузка модели регрессии
         self.audio_dir = Path(db_path).parent / "DEAM_audio" / "MEMD_audio"
-        self.regressor = joblib.load(model_path) if Path(model_path).exists() else None
+        
+        if Path(model_path).exists():
+            self.regressor = joblib.load(model_path)
+        else:
+            self.regressor = None
 
-    def predict_va(self, embedding: np.ndarray):
-        if self.regressor:
-            prediction = self.regressor.predict(embedding.reshape(1, -1))[0]
-            return np.clip(prediction[0], 1.0, 9.0), np.clip(prediction[1], 1.0, 9.0)
-        return 5.0, 5.0
+    def predict_va(self, embedding: np.ndarray) -> tuple[float, float]:
+        # Прогнозирование координат Valence/Arousal по визуальному эмбеддингу
+        if not self.regressor:
+            return 5.0, 5.0
+            
+        raw_prediction = self.regressor.predict(embedding.reshape(1, -1))[0]
+        valence_pred = np.clip(raw_prediction[0], 1.0, 9.0)
+        arousal_pred = np.clip(raw_prediction[1], 1.0, 9.0)
+        
+        return float(valence_pred), float(arousal_pred)
 
-    def get_audio_path(self, song_id):
-        if not self.audio_dir.exists(): return None
+    def get_audio_path(self, song_id: int | float | str) -> Path | None:
+        # Поиск физического пути к аудиофайлу в зависимости от расширения
+        if not self.audio_dir.exists():
+            return None
+            
         clean_id = str(int(float(song_id)))
         for ext in ['.mp3', '.wav']:
             path = self.audio_dir / f"{clean_id}{ext}"
-            if path.exists(): return path
+            if path.exists():
+                return path
         return None
 
-    def find_nearest_tracks(self, target_v: float, target_a: float, llm_profile: dict = None, top_k: int = 5):
-        # 1. Эмоциональная дистанция (как и раньше)
-        emo_dist = np.sqrt(
-            1.0 * (self.norm_db['valence'] - target_v)**2 + 
-            2.5 * (self.norm_db['arousal'] - target_a)**2
-        )
-        self.norm_db['emo_distance'] = emo_dist
+    def find_nearest_tracks(self, target_v: float, target_a: float, llm_profile: dict = None, top_k: int = 5) -> pd.DataFrame:
+        # Расчет евклидова расстояния в эмоциональном пространстве Рассела
+        v_dist = (self.norm_db['valence'] - target_v) ** 2
+        a_dist = (self.norm_db['arousal'] - target_a) ** 2
         
-        # Если LLM не дала ответ, сортируем только по эмоциям
+        # Взвешенное расстояние с приоритетом оси активации (Arousal)
+        self.norm_db['emo_distance'] = np.sqrt(1.0 * v_dist + 2.5 * a_dist)
+        
+        # Ранжирование только по эмоциональному критерию при отсутствии профиля LLM
         if not llm_profile:
             self.norm_db['final_score'] = self.norm_db['emo_distance']
             return self.norm_db.sort_values(by='final_score').head(top_k)
 
-        # 2. Акустическая дистанция (сравниваем треки с запросом LLM)
+        # Расчет отклонений по вектору акустических параметров LLM
         acoustic_penalty = np.zeros(len(self.norm_db))
         for feat in self.acoustic_features:
             if feat in llm_profile:
                 target_val = llm_profile[feat]
                 acoustic_penalty += np.abs(self.norm_db[f"norm_{feat}"] - target_val)
 
-        # Усредняем штраф
+        # Нормирование акустической дистанции
         self.norm_db['acoustic_distance'] = acoustic_penalty / len(self.acoustic_features)
 
-        # 3. Финальный Score (Смесь Эмоций и Акустики). Коэф 4.0 делает акустику важной!
+        # Вычисление интегральной метрики соответствия (мультимодальный скоринг)
         self.norm_db['final_score'] = self.norm_db['emo_distance'] + (self.norm_db['acoustic_distance'] * 4.0)
 
         return self.norm_db.sort_values(by='final_score').head(top_k)