Fix documentation inconsistencies, type annotations, and CI configuration

.env.example

@@ -1,2 +1,4 @@
 HUGGINGFACE_TOKEN="your_huggingface_token_here"
 WANDB_API_KEY="your_wandb_api_key_here"
+WANDB_PROJECT="your_wandb_project_name_here"
+WANDB_ENTITY="your_wandb_username_here"

.github/workflows/python-package.yml

@@ -15,7 +15,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: ["3.10", "3.11", "3.12"]
+        python-version: ["3.12"]
 
     steps:
       - uses: actions/checkout@v3
@@ -27,6 +27,7 @@ jobs:
       - name: Install uv
         run: |
           curl -LsSf https://astral.sh/uv/install.sh | sh
+          echo "$HOME/.local/bin" >> $GITHUB_PATH
       - name: Run pre-commit hook
         run: |
           uv run pre-commit run -a

.gitignore

@@ -172,3 +172,7 @@ cython_debug/
 
 # MacOS
 .DS_Store
+
+# Training outputs
+logs/
+wandb/

README.md

@@ -248,14 +248,15 @@ class Config(BaseModel):
     seed: int = 42
 
     # Data
-    test_split: float = 0.05
-    batch_size: int = 4
+    val_split: float = 0.05
+    batch_size: int = 8
 
     # Training
     max_epochs: int = 200
     early_stopping_patience: int = 10
-    learning_rate: float = 1e-4
+    learning_rate: float = 5e-5
     min_learning_rate: float = 1e-6
+    lr_scheduler: Literal["onecycle", "constant"] = "constant"
     weight_decay: float = 1e-2
     accumulate_grad_batches: int = 1
     gradient_clip_val: float = 1.0
@@ -273,8 +274,9 @@ class Config(BaseModel):
 - **base_model**: HuggingFace model ID for the pretrained encoder
 - **num_frames**: Number of video frames per sample (5 frames = 0.2s at 25fps)
 - **negative_fraction**: Proportion of negative samples (0.5 = 50% out-of-sync)
-- **batch_size**: Adjust based on GPU memory (4 works well for most GPUs)
-- **learning_rate**: Initial learning rate with OneCycleLR scheduler
+- **batch_size**: Adjust based on GPU memory (8 is the default, works well for most GPUs)
+- **learning_rate**: Initial learning rate (5e-5 by default)
+- **lr_scheduler**: Learning rate scheduler type ("constant" or "onecycle")
 
 ## Training Details
 
@@ -288,8 +290,8 @@ class Config(BaseModel):
 ### Optimization
 
 - **Optimizer**: AdamW with weight decay
-- **Scheduler**: OneCycleLR with cosine annealing
-  - 10% warmup period
+- **Scheduler**: Configurable (constant by default, or OneCycleLR with cosine annealing)
+  - OneCycleLR: 10% warmup period, cosine annealing
   - Peak learning rate: `config.learning_rate`
   - Final learning rate: `config.min_learning_rate`
 

pyproject.toml

@@ -58,4 +58,3 @@ exclude = [".venv"]
 
 [project.scripts]
 train = "syncnet.scripts.train:main"
-compile = "syncnet.scripts.compile:main"

src/syncnet/datamodule.py

@@ -5,6 +5,7 @@
 collation, audio-visual preprocessing, and data augmentation.
 """
 
+import logging
 import random
 
 import torch
@@ -17,6 +18,8 @@
 from syncnet.config import Config
 from syncnet.datasets import Batch
 
+logger = logging.getLogger(__name__)
+
 
 class SyncNetDataModule(LightningDataModule):
     """PyTorch Lightning DataModule for SyncNet training and evaluation.
@@ -186,7 +189,7 @@ def pad_collate_fn(
                 audio_segments.append(input_values["input_values"][0])
                 labels.append(label)
             except Exception as e:
-                print(f"Error processing sample: {e}")
+                logger.warning(f"Error processing sample: {e}")
                 continue
 
         return Batch(

src/syncnet/lightning_module.py

@@ -5,6 +5,7 @@
 Lightning framework.
 """
 
+import logging
 from pathlib import Path
 
 import torch
@@ -18,6 +19,8 @@
 from syncnet.datasets import Batch
 from syncnet.modeling.model import SyncNet, SyncNetConfig
 
+logger = logging.getLogger(__name__)
+
 
 class SyncNetLightningModule(LightningModule):
     """PyTorch Lightning Module for training SyncNet models.
@@ -80,7 +83,7 @@ def __init__(
         self.processor = PeAudioVideoProcessor.from_pretrained(config.base_model)
         self.lowest_val_loss = float("inf")
 
-    def training_step(self, batch: Batch, batch_idx: int) -> None:
+    def training_step(self, batch: Batch, batch_idx: int) -> torch.Tensor:
         """Execute a single training step.
 
         Performs forward pass through the model, computes loss, and logs metrics.
@@ -100,7 +103,7 @@ def training_step(self, batch: Batch, batch_idx: int) -> None:
         self.log("train_loss", self.train_loss(loss), prog_bar=True)
         return loss
 
-    def validation_step(self, batch: Batch, batch_idx: int) -> None:
+    def validation_step(self, batch: Batch, batch_idx: int) -> torch.Tensor:
         """Execute a single validation step.
 
         Performs forward pass and computes validation metrics without gradient
@@ -158,25 +161,27 @@ def on_validation_epoch_end(self) -> None:
                         private=True,
                     )
                 except Exception as e:
-                    print(f"Failed to push to hub: {e}")
+                    logger.warning(f"Failed to push to hub: {e}")
 
         self.val_metrics.reset()
         self.val_loss.reset()
         garbage_collection_cuda()
 
-    def configure_optimizers(self) -> tuple[list[torch.optim.Optimizer], list]:
+    def configure_optimizers(
+        self,
+    ) -> tuple[list[torch.optim.Optimizer], list[torch.optim.lr_scheduler.LRScheduler]]:
         """Configure optimizers and learning rate schedulers.
 
-        Sets up AdamW optimizer with weight decay and OneCycleLR scheduler
-        for cosine annealing learning rate schedule with warmup.
+        Sets up AdamW optimizer with weight decay and a configurable scheduler
+        (constant or OneCycleLR with cosine annealing).
 
         Returns:
             Tuple containing:
                 - List with single AdamW optimizer
-                - List with single OneCycleLR scheduler
+                - List with single LR scheduler
 
         Note:
-            The scheduler uses:
+            When using OneCycleLR scheduler:
             - 10% of training for warmup (pct_start=0.1)
             - Cosine annealing strategy
             - Final learning rate of config.min_learning_rate