Update core/model_runner.py

core/model_runner.py

@@ -26,39 +26,85 @@ def get_model(
     target,
     model_name="LSTM",
     horizon=1,
-    hidden=64,
-    hidden_units=None,  # alias for backward compatibility
-    layers=1,
+    # Hidden size aliases
+    hidden=None,
+    hidden_units=None,
+    # Layers aliases
+    layers=None,
+    n_layers=None,
+    # Learning rate aliases
+    lr=None,
+    learning_rate=None,
+    # Betas for optimizer
+    beta1=0.9,
+    beta2=0.999,
+    # Other hyperparams
     epochs=50,
-    lr=0.001,
     weight_decay=0.01,
     dropout=0.2,
-    window=30,
+    # Window aliases
+    window=None,
+    window_size=None,
     test_split=0.2,
     selector_method="RandomForest",
     importance_threshold=0.0,
     scheduler_type="None",
+    device=None,
     verbose=True,
 ):
+    """
+    Wrapper that accepts many common argument names used by the UI/calls,
+    normalizes them, and calls train_and_evaluate(...) with the canonical names.
+    """
     try:
-        # prefer hidden_units if passed
+        # --- Normalize aliases & defaults ---
+        # hidden size: prefer explicit hidden_units, then hidden, else default 64
         if hidden_units is not None:
             hidden = hidden_units
+        if hidden is None:
+            hidden = 64
 
+        # layers: prefer explicit n_layers, then layers, else default 1
+        if n_layers is not None:
+            layers = n_layers
+        if layers is None:
+            layers = 1
+
+        # learning rate: prefer learning_rate then lr, else default 0.001
+        if learning_rate is not None:
+            lr = learning_rate
+        if lr is None:
+            lr = 0.001
+
+        # window size: prefer window_size then window, else default 30
+        if window_size is not None:
+            window = window_size
+        if window is None:
+            window = 30
+
+        # device: caller may pass it; otherwise detect automatically
+        if device is None:
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+
+        logging.info(
+            f"get_model called: model={model_name}, device={device}, hidden={hidden}, layers={layers}, lr={lr}, window={window}, epochs={epochs}"
+        )
+
+        # --- Select model class mapping (keys as used in UI) ---
         model_classes = {
             "LSTM": LSTMModel,
             "GRU": GRUModel,
             "CNN": CNNModel,
             "MLP": MLPModel,
             "Hybrid": HybridCNNGRUModel,
+            "HybridCNNGRU": HybridCNNGRUModel,
             "Transformer": TransformerModel,
             "BiLSTM": BiLSTMModel,
         }
 
         model_cls = model_classes.get(model_name, LSTMModel)
-        device = "cuda" if torch.cuda.is_available() else "cpu"
-        logging.info(f"Running {model_name} on device={device}")
 
+        # --- Call the core training function with canonical param names ---
         result = train_and_evaluate(
             df=df,
             features=features,
@@ -69,6 +115,8 @@ def get_model(
             layers=layers,
             epochs=epochs,
             lr=lr,
+            beta1=beta1,
+            beta2=beta2,
             weight_decay=weight_decay,
             dropout=dropout,
             window=window,
@@ -80,14 +128,17 @@ def get_model(
             verbose=verbose,
         )
 
+        # --- Normalize return --- 
         if not result:
             logging.error(f"{model_name} returned empty result.")
-            return {"error": "Empty result"}
+            return {"error": "Empty result from training"}
         if isinstance(result, dict) and result.get("error"):
             logging.error(f"{model_name} training error: {result['error']}")
             return {"error": result["error"]}
+
+        logging.info(f"{model_name} training completed successfully")
         return result
 
     except Exception as e:
-        logging.error(f"Model runner error for {model_name}: {str(e)}")
+        logging.error(f"Model runner error for {model_name}: {str(e)}", exc_info=True)
         return {"error": str(e)}