feat: add TorchSklearnWrapper for PyTorch model compatibility with sklearn and enhance final model training process

app/training/train_deep_classifiers.py

@@ -273,8 +273,129 @@ def evaluate_cv(
     }
 
 
+class TorchSklearnWrapper:
+    """
+    Sklearn-compatible wrapper for trained PyTorch classifiers.
+    Saves model class name + state dict so it can be pickled and reloaded.
+    """
+
+    def __init__(
+        self,
+        model_class: type,
+        n_features: int,
+        state_dict: dict,
+        scaler: StandardScaler,
+    ) -> None:
+        self.model_class_name = model_class.__name__
+        self._model_class = model_class
+        self.n_features = n_features
+        self.state_dict = state_dict
+        self.scaler = scaler
+        self.n_features_in_ = n_features
+
+    def _build_model(self) -> nn.Module:
+        model = self._model_class(self.n_features)
+        model.load_state_dict(self.state_dict)
+        model.eval()
+        return model
+
+    def predict_proba(self, X: np.ndarray) -> np.ndarray:
+        model = self._build_model().to("cpu")
+        X_scaled = self.scaler.transform(X)
+        x_t = torch.tensor(X_scaled, dtype=torch.float32)
+        with torch.no_grad():
+            logits = model(x_t)
+            probs = torch.sigmoid(logits).numpy().flatten()
+        return np.column_stack([1.0 - probs, probs])
+
+    def __getstate__(self) -> dict:
+        state = self.__dict__.copy()
+        state.pop("_model_class", None)
+        return state
+
+    def __setstate__(self, state: dict) -> None:
+        self.__dict__.update(state)
+        # Re-attach class from global lookup
+        _CLASS_MAP = {
+            "DeepMLP": DeepMLP,
+            "Conv1DClassifier": Conv1DClassifier,
+            "ResidualMLP": ResidualMLP,
+            "AttentionMLP": AttentionMLP,
+        }
+        self._model_class = _CLASS_MAP.get(self.model_class_name, DeepMLP)
+
+
+def train_final_model(
+    model_class: type,
+    X: np.ndarray,
+    y: np.ndarray,
+    epochs: int = EPOCHS,
+    patience: int = PATIENCE,
+) -> TorchSklearnWrapper:
+    """Train model on full dataset and return sklearn-compatible wrapper."""
+    from sklearn.model_selection import train_test_split
+
+    scaler = StandardScaler()
+    X_tr_raw, X_val_raw, y_tr, y_val = train_test_split(
+        X, y, test_size=0.1, stratify=y, random_state=SEED
+    )
+    X_tr = scaler.fit_transform(X_tr_raw)
+    X_v = scaler.transform(X_val_raw)
+
+    n_features = X.shape[1]
+    model = model_class(n_features).to(DEVICE)
+    optimizer = torch.optim.AdamW(model.parameters(), lr=LR, weight_decay=1e-4)
+    criterion = nn.BCEWithLogitsLoss()
+    loader = DataLoader(
+        TensorDataset(
+            torch.tensor(X_tr, dtype=torch.float32),
+            torch.tensor(y_tr, dtype=torch.float32),
+        ),
+        batch_size=BATCH_SIZE,
+        shuffle=True,
+    )
+    val_X = torch.tensor(X_v, dtype=torch.float32).to(DEVICE)
+    val_y = torch.tensor(y_val, dtype=torch.float32)
+
+    best_auc = 0.0
+    best_state = None
+    patience_ctr = 0
+
+    for epoch in range(epochs):
+        model.train()
+        for bx, by in loader:
+            bx, by = bx.to(DEVICE), by.to(DEVICE)
+            optimizer.zero_grad()
+            criterion(model(bx), by).backward()
+            optimizer.step()
+
+        model.eval()
+        with torch.no_grad():
+            probs = torch.sigmoid(model(val_X)).cpu().numpy()
+        auc = roc_auc_score(val_y.numpy(), probs)
+        if auc > best_auc:
+            best_auc = auc
+            best_state = {k: v.cpu().clone() for k, v in model.state_dict().items()}
+            patience_ctr = 0
+        else:
+            patience_ctr += 1
+            if patience_ctr >= patience:
+                break
+
+    return TorchSklearnWrapper(model_class, n_features, best_state or model.state_dict(), scaler)
+
+
+def _safe_name(name: str) -> str:
+    return name.lower().replace(" ", "_").replace("(", "").replace(")", "").replace("-", "_")
+
+
 def main() -> None:
+    import pickle
+
     csv_path = sys.argv[1] if len(sys.argv) > 1 else "../DataSet/features.csv"
+    out_dir = Path(sys.argv[2]) if len(sys.argv) > 2 else Path("models")
+    out_dir.mkdir(parents=True, exist_ok=True)
+
     print(f"Device: {DEVICE}")
     print(f"Loading: {csv_path}")
 
@@ -283,7 +404,7 @@ def main() -> None:
     print(f"Samples: {len(y)}, Features: {n_features}")
     print(f"AI: {np.sum(y == 1)}, Human: {np.sum(y == 0)}")
 
-    models = {
+    model_classes = {
         "Deep MLP (512-256-128-64)": DeepMLP,
         "1D-CNN": Conv1DClassifier,
         "Residual MLP (3 blocks)": ResidualMLP,
@@ -291,16 +412,24 @@ def main() -> None:
     }
 
     all_results = {}
-    for name, cls in models.items():
+    for name, cls in model_classes.items():
         print(f"\n{'='*60}")
         print(f"  {name}")
         print(f"{'='*60}")
         result = evaluate_cv(cls, X, y, n_features)
-        all_results[name] = result
+        all_results[name] = {**result, "type": "deep_learning"}
         print(f"  => Acc={result['accuracy']:.4f}  AUC={result['roc_auc']:.4f}  "
               f"F1={result['f1']:.4f}  Time={result['train_time_sec']:.0f}s")
 
-    out_path = Path("models/deep_learning_results.json")
+        print(f"  Training final model for {name}...")
+        wrapper = train_final_model(cls, X, y)
+        pkl_path = out_dir / f"model_dl_{_safe_name(name)}.pkl"
+        with open(pkl_path, "wb") as f:
+            pickle.dump(wrapper, f)
+        all_results[name]["model_path"] = str(pkl_path)
+        print(f"  Saved: {pkl_path}")
+
+    out_path = out_dir / "deep_learning_results.json"
     with open(out_path, "w") as f:
         json.dump(all_results, f, indent=2)
     print(f"\nResults saved: {out_path}")