(chore): unify model selection via shared registry (train uses 'choices()'/'build()')

models/registry.py

@@ -0,0 +1,24 @@
+# models/registry.py
+from typing import Callable, Dict
+from models.figure2_cnn import Figure2CNN
+from models.resnet_cnn import ResNet1D
+# from models.resnet18_vision import ResNet18Vision # (Step 2)
+
+# Internal registry of model builders keyed by short name.
+_REGISTRY: Dict[str, Callable[[int], object]] = {
+    "figure2": lambda L: Figure2CNN(input_length=L),
+    "resnet": lambda L: ResNet1D(input_length=L),
+    # "resnet18vision": lambda L: ResNet18Vision(input_length=L)
+}
+
+def choices():
+    """Return the list of available model keys."""
+    return list(_REGISTRY.keys())
+
+def build(name: str, input_length: int):
+    """Instantiate a model by short name with the given input length."""
+    if name not in _REGISTRY:
+        raise ValueError(f"Unknown model '{name}'. Choices: {choices()}")
+    return _REGISTRY[name](input_length)
+
+__all__ = ["choices", "build"]

scripts/train_model.py

@@ -7,10 +7,23 @@ from torch.utils.data import TensorDataset, DataLoader
 from sklearn.model_selection import StratifiedKFold
 from sklearn.metrics import confusion_matrix
 
+import random
+import json
+
+# Reproducibility
+SEED = 42
+random.seed(SEED)
+np.random.seed(SEED)
+torch.manual_seed(SEED)
+if torch.cuda.is_available():
+    torch.cuda.manual_seed_all(SEED)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
 # Add project-specific imports
 from scripts.preprocess_dataset import preprocess_dataset
-from models.figure2_cnn import Figure2CNN
-from models.resnet_cnn import ResNet1D
+from models.registry import choices as model_choices, build as build_model
+
 
 # Argument parser for CLI usage
 parser = argparse.ArgumentParser(