AI and Machine Learning Background Questions

import os
import json
import tempfile
import shutil
from pathlib import Path

import numpy as np
import tensorflow as tf
from sklearn.base import BaseEstimator, ClassifierMixin

class KerasSKClassifier(BaseEstimator, ClassifierMixin):
    """
    Scikit-learn wrapper for a TensorFlow Keras classification model.

    Parameters (examples):
    - build_fn: callable(**params) -> compiled tf.keras.Model
    - epochs: int
    - batch_size: int
    - verbose: int
    - seed: int or None
    - model_dir: str or None (if None, uses temporary dir)
    - **fit_kwargs forwarded to model.fit (e.g., validation_data)
    """
    def __init__(self, build_fn,
                 epochs=10,
                 batch_size=32,
                 verbose=0,
                 seed=None,
                 model_dir=None,
                 **kwargs):
        self.build_fn = build_fn
        self.epochs = int(epochs)
        self.batch_size = int(batch_size)
        self.verbose = int(verbose)
        self.seed = seed
        self.model_dir = model_dir
        # any extra hyperparameters forwarded to build_fn
        self.build_params = kwargs

        # internal
        self._model = None
        self._model_path = None

    def _set_seed(self):
        if self.seed is None:
            return
        os.environ['PYTHONHASHSEED'] = str(self.seed)
        np.random.seed(self.seed)
        tf.random.set_seed(self.seed)

    def _ensure_model_dir(self):
        if self.model_dir:
            os.makedirs(self.model_dir, exist_ok=True)
            return self.model_dir
        td = tempfile.mkdtemp(prefix="keras_sklearn_")
        return td

    def fit(self, X, y, sample_weight=None, class_weight=None, **fit_kwargs):
        """
        Fit the Keras model.
        Accepts sample_weight and class_weight.
        """
        self._set_seed()
        # build model
        params = dict(self.build_params)
        model = self.build_fn(**params)
        if not isinstance(model, tf.keras.Model):
            raise ValueError("build_fn must return a compiled tf.keras.Model")

        # convert X,y to numpy
        X_np = np.asarray(X)
        y_np = np.asarray(y)

        # fit
        history = model.fit(
            X_np, y_np,
            epochs=self.epochs,
            batch_size=self.batch_size,
            verbose=self.verbose,
            sample_weight=sample_weight,
            class_weight=class_weight,
            **fit_kwargs
        )

        # save model architecture + weights to allow GridSearchCV cloning/serialization
        model_dir = self._ensure_model_dir()
        arch_path = Path(model_dir) / "model_arch.json"
        weights_path = Path(model_dir) / "model_weights.h5"

        # architecture
        with open(arch_path, "w") as f:
            f.write(model.to_json())
        # weights
        model.save_weights(str(weights_path))

        # keep paths and a loaded model for predictions
        self._model_path = model_dir
        # to ensure a fresh loaded model (simulate serialization)
        loaded = tf.keras.models.model_from_json(json.load(open(arch_path)))
        loaded.load_weights(str(weights_path))
        # compiled state is not preserved by model_from_json; recompile using build_fn or compile args:
        # if build_fn returns compiled model, prefer re-building compile config:
        try:
            # attempt to recompile using the original compiled optimizer/loss if available
            compiled = model.optimizer.get_config(), model.loss
            # best-effort: recompile loaded with same optimizer config
            opt_config = model.optimizer.get_config()
            opt_class = model.optimizer.__class__
            loaded.compile(optimizer=opt_class.from_config(opt_config), loss=model.loss, metrics=getattr(model, 'metrics', None))
        except Exception:
            # fallback: skip compile; predict/predict_proba don't need compilation
            pass

        self._model = loaded
        return self

    def predict_proba(self, X):
        if self._model is None:
            raise ValueError("Model is not fitted yet.")
        X_np = np.asarray(X)
        probs = self._model.predict(X_np, batch_size=self.batch_size, verbose=self.verbose)
        # ensure 2D probability output
        return probs

    def predict(self, X):
        probs = self.predict_proba(X)
        # for binary and multiclass
        if probs.ndim == 1 or probs.shape[1] == 1:
            # binary: threshold 0.5
            return (probs.ravel() > 0.5).astype(int)
        else:
            return np.argmax(probs, axis=1)

    def save(self, path):
        """Persist wrapper (model artifacts already saved during fit)."""
        if self._model_path is None:
            raise ValueError("No model artifacts to save. Fit first.")
        shutil.copytree(self._model_path, path, dirs_exist_ok=True)

    def load(self, path):
        """Load persisted artifacts from path."""
        arch_path = Path(path) / "model_arch.json"
        weights_path = Path(path) / "model_weights.h5"
        if not arch_path.exists() or not weights_path.exists():
            raise ValueError("Missing model artifacts in path.")
        model = tf.keras.models.model_from_json(json.load(open(arch_path)))
        model.load_weights(str(weights_path))
        self._model = model
        self._model_path = path
        return self

    # scikit-learn compatibility
    def get_params(self, deep=True):
        # include build_fn and build_params
        params = {
            "build_fn": self.build_fn,
            "epochs": self.epochs,
            "batch_size": self.batch_size,
            "verbose": self.verbose,
            "seed": self.seed,
            "model_dir": self.model_dir,
        }
        params.update(self.build_params)
        return params

    def set_params(self, **params):
        # set known hyperparameters; unknown pass into build_params
        for key, val in params.items():
            if hasattr(self, key):
                setattr(self, key, val)
            else:
                self.build_params[key] = val
        return self

from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler

pipe = Pipeline([
  ("scaler", StandardScaler()),
  ("keras", KerasSKClassifier(build_fn=build_fn, epochs=20, batch_size=64, seed=42))
])

from sklearn.model_selection import GridSearchCV
param_grid = {"keras__epochs": [10,20], "keras__learning_rate": [1e-3, 1e-4]}
gs = GridSearchCV(pipe, param_grid, cv=3)
gs.fit(X_train, y_train)