fix: add return type annotations and docstrings for clarity in Wav2Vec2MusicClassifier and AudioDataset

app/training/wav2vec2_classifier.py

@@ -53,7 +53,8 @@ class Wav2Vec2MusicClassifier(nn.Module):
     learn task-specific temporal patterns.
     """
 
-    def __init__(self, config: Wav2Vec2Config | None = None):
+    def __init__(self, config: Wav2Vec2Config | None = None) -> None:
+        """Initialize wav2vec2 classifier with frozen CNN encoder."""
         super().__init__()
         self.config = config or Wav2Vec2Config()
 
@@ -109,16 +110,18 @@ class AudioDataset(Dataset):
         labels: list[int],
         sample_rate: int = 16000,
         max_sec: float = 30.0,
-    ):
+    ) -> None:
+        """Initialize audio dataset with file paths and labels."""
         self.file_paths = file_paths
         self.labels = labels
         self.sample_rate = sample_rate
         self.max_samples = int(max_sec * sample_rate)
 
-    def __len__(self):
+    def __len__(self) -> int:
+        """Return dataset size."""
         return len(self.file_paths)
 
-    def __getitem__(self, idx):
+    def __getitem__(self, idx: int) -> tuple[torch.Tensor, int]:
         import librosa
 
         path = self.file_paths[idx]
@@ -136,7 +139,9 @@ class AudioDataset(Dataset):
         return torch.tensor(y, dtype=torch.float32), label
 
 
-def collate_fn(batch):
+def collate_fn(
+    batch: list[tuple[torch.Tensor, int]],
+) -> tuple[torch.Tensor, torch.Tensor]:
     """Collate audio tensors and labels."""
     audios, labels = zip(*batch)
     audios = torch.stack(audios)