Building Your First Neural Network with Keras

# io.thecodeforge: Neural Network with Keras Sequential API
# Keras 3.0+ — import from keras directly, not tensorflow.keras
# This preserves backend flexibility (TensorFlow, PyTorch, JAX)
import os
os.environ['KERAS_BACKEND'] = 'tensorflow'  # Switch to 'jax' or 'torch' to change backend

import keras
from keras import layers, models
import numpy as np


def create_forge_classifier(input_dim: int, num_classes: int) -> keras.Model:
    """
    Build a regularized dense classifier for tabular or flattened image data.

    Args:
        input_dim:   Number of features per sample (e.g., 784 for 28x28 MNIST images).
        num_classes: Number of output classes (e.g., 10 for digits 0-9).

    Returns:
        Compiled Keras model ready for training.

    Architecture decisions documented:
    - He Normal init: correct for ReLU activations (Glorot is for Tanh/Sigmoid)
    - L2 regularization: weight penalty prevents memorization on small datasets
    - Dropout(0.3): applied after the larger hidden layer, not before the output
    - BatchNormalization: stabilizes training by normalizing layer inputs per batch
    """
    model = models.Sequential(
        [
            # Input shape declaration — batch dimension (None) is implicit.
            # Always declare Input explicitly rather than relying on first layer
            # inference — it makes model.input_shape reliable before the first fit.
            layers.Input(shape=(input_dim,), name='input'),

            # First hidden layer.
            # kernel_initializer='he_normal' pairs with ReLU — He accounts for
            # the fact that ReLU zeroes half of all inputs, so variance needs
            # to be higher at init to keep gradient magnitudes stable.
            layers.Dense(
                256,
                activation='relu',
                kernel_initializer='he_normal',
                kernel_regularizer=keras.regularizers.l2(1e-4),
                name='hidden_1'
            ),
            layers.BatchNormalization(name='bn_1'),
            layers.Dropout(0.3, name='dropout_1'),

            # Second hidden layer — narrower, forcing compression.
            layers.Dense(
                128,
                activation='relu',
                kernel_initializer='he_normal',
                kernel_regularizer=keras.regularizers.l2(1e-4),
                name='hidden_2'
            ),
            layers.BatchNormalization(name='bn_2'),
            layers.Dropout(0.2, name='dropout_2'),

            # Output layer.
            # softmax: outputs sum to 1.0 across all classes — correct for
            # multi-class classification. Pair with sparse_categorical_crossentropy
            # if labels are integers, categorical_crossentropy if one-hot.
            layers.Dense(
                num_classes,
                activation='softmax',
                name='output'
            ),
        ],
        name='forge_classifier'
    )

    model.compile(
        # Adam with a slightly reduced learning rate from the default 0.001.
        # Default often causes loss instability in the first few epochs on
        # small tabular datasets — 0.0005 is a safer starting point.
        optimizer=keras.optimizers.Adam(learning_rate=5e-4),
        loss='sparse_categorical_crossentropy',
        metrics=['accuracy', keras.metrics.TopKCategoricalAccuracy(k=3, name='top3_acc')]
    )
    return model


if __name__ == '__main__':
    # Verify the architecture before touching real data.
    model = create_forge_classifier(input_dim=784, num_classes=10)
    model.summary()

    # Smoke test with synthetic data — catches shape mismatches before
    # the full training job is submitted to a GPU cluster.
    X_dummy = np.random.rand(32, 784).astype('float32')  # Already in [0, 1]
    y_dummy = np.random.randint(0, 10, size=(32,))
    loss_val = model.evaluate(X_dummy, y_dummy, verbose=0)
    print(f'Smoke test loss: {loss_val[0]:.4f} — model is wired correctly if this is a reasonable number')

-- io.thecodeforge: Production training data schema
-- Design decisions:
--   feature_id BIGSERIAL: monotonically increasing, enables deterministic ORDER BY
--   dataset_version TEXT: ties each row to a specific labeled dataset snapshot
--   split TEXT: train/val/test assignment baked into the table, not computed at query time
--   normalization_version TEXT: tracks which preprocessing run produced the scaled values
--     Normalization parameters (mean, std per feature) are stored in a separate table
--     so they can be retrieved and applied identically at inference time.

CREATE TABLE IF NOT EXISTS io_thecodeforge_training_features (
    feature_id      BIGSERIAL PRIMARY KEY,
    dataset_version TEXT      NOT NULL,   -- e.g., 'fraud_v3_2026q1'
    split           TEXT      NOT NULL CHECK (split IN ('train', 'val', 'test')),
    raw_vector      FLOAT8[]  NOT NULL,   -- original unscaled values
    scaled_vector   FLOAT8[]  NOT NULL,   -- normalized to [0,1] or z-score
    label_index     SMALLINT  NOT NULL,
    normalization_version TEXT NOT NULL,  -- links to normalization_params table
    created_at      TIMESTAMPTZ DEFAULT NOW()
);

-- Normalization parameters: stored at the pipeline level, not computed on the fly.
-- At inference time, retrieve these and apply the same transformation to incoming data.
-- This is the contract between training and serving.
CREATE TABLE IF NOT EXISTS io_thecodeforge_normalization_params (
    normalization_version TEXT PRIMARY KEY,
    feature_means   FLOAT8[] NOT NULL,
    feature_stds    FLOAT8[] NOT NULL,
    feature_mins    FLOAT8[] NOT NULL,
    feature_maxes   FLOAT8[] NOT NULL,
    computed_at     TIMESTAMPTZ DEFAULT NOW(),
    row_count       BIGINT   NOT NULL  -- how many rows this was computed over
);

-- Training query — deterministic ordering by feature_id is mandatory.
-- Without ORDER BY, batch composition differs across runs and experiments
-- are not reproducible.
SELECT
    scaled_vector,
    label_index
FROM io_thecodeforge_training_features
WHERE
    dataset_version = 'fraud_v3_2026q1'
    AND split = 'train'
    AND label_index IS NOT NULL
ORDER BY feature_id  -- deterministic ordering: same batches every run
LIMIT 100000;

-- Validation query — same dataset_version, different split.
-- Never use the test split during training or hyperparameter tuning.
SELECT scaled_vector, label_index
FROM io_thecodeforge_training_features
WHERE dataset_version = 'fraud_v3_2026q1'
  AND split = 'val'
ORDER BY feature_id;

# io.thecodeforge: Production ML Training Environment
# Pin exact TensorFlow + CUDA + cuDNN versions.
# tensorflow/tensorflow:2.16.1-gpu includes:
#   CUDA 12.2, cuDNN 8.9, Python 3.11
# Verify compatibility at: https://www.tensorflow.org/install/source#gpu
#
# NEVER use :latest — it resolves to different images on different days
# and breaks the reproducibility guarantee that makes Docker valuable.
FROM tensorflow/tensorflow:2.16.1-gpu

# Set working directory.
WORKDIR /app

# Install Python dependencies in a separate layer before copying application code.
# This layer is cached by Docker — if requirements.txt doesn't change,
# pip install doesn't re-run on every build, saving 3-5 minutes per iteration.
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

# Copy application code last — changes here only invalidate this layer,
# not the dependency installation layer above.
COPY . .

# Set Keras backend explicitly — don't rely on default detection.
ENV KERAS_BACKEND=tensorflow

# Disable TensorFlow's aggressive GPU memory pre-allocation.
# Without this, TF allocates all available GPU memory at startup, blocking
# other processes on shared GPU machines (common in dev and staging).
ENV TF_FORCE_GPU_ALLOW_GROWTH=true

# Mixed precision: halves memory usage with minimal accuracy impact on
# supported GPUs (NVIDIA Ampere and newer — A100, RTX 30xx, RTX 40xx).
# Set to '' to disable on older hardware.
ENV TF_ENABLE_AUTO_MIXED_PRECISION=1

# Default command — override with docker run ... python train_job.py
CMD ["python", "forge_nn_basic.py"]

# io.thecodeforge: Complete training pipeline with normalization,
# callbacks, and the benchmark-first pattern.

import os
os.environ['KERAS_BACKEND'] = 'tensorflow'

import numpy as np
import keras
from keras import layers, models, callbacks
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import roc_auc_score
import time


def benchmark_first(X_train: np.ndarray, y_train: np.ndarray,
                   X_val: np.ndarray, y_val: np.ndarray) -> float:
    """
    Run a baseline sklearn model before training the neural network.
    If the baseline AUC exceeds 0.90, evaluate whether a neural network
    adds enough value to justify the operational complexity.

    Returns:
        Baseline AUC on validation set.
    """
    print('\n=== Benchmark: RandomForestClassifier (no tuning) ===')
    t0 = time.time()
    rf = RandomForestClassifier(n_estimators=100, random_state=42, n_jobs=-1)
    rf.fit(X_train, y_train)
    rf_preds = rf.predict_proba(X_val)[:, 1]
    rf_auc = roc_auc_score(y_val, rf_preds)
    print(f'Random Forest AUC: {rf_auc:.4f} | Training time: {time.time() - t0:.1f}s')
    print('If RF AUC > 0.90, question whether the neural network is necessary.\n')
    return rf_auc


def prepare_data(X_raw: np.ndarray, y_raw: np.ndarray):
    """
    Standard preprocessing pipeline:
    1. Train/val/test split (70/15/15)
    2. StandardScaler fit on training data only
    3. Transform val and test with training scaler (no data leakage)

    Returns:
        Tuple of (X_train, X_val, X_test, y_train, y_val, y_test, scaler)
        The scaler must be saved alongside the model for inference.
    """
    X_train, X_temp, y_train, y_temp = train_test_split(
        X_raw, y_raw, test_size=0.30, random_state=42, stratify=y_raw
    )
    X_val, X_test, y_val, y_test = train_test_split(
        X_temp, y_temp, test_size=0.50, random_state=42, stratify=y_temp
    )

    # Fit scaler on training data only.
    # Fitting on full dataset is data leakage — val/test statistics
    # should never influence the preprocessing parameters.
    scaler = StandardScaler()
    X_train = scaler.fit_transform(X_train).astype('float32')
    X_val   = scaler.transform(X_val).astype('float32')
    X_test  = scaler.transform(X_test).astype('float32')

    # Audit the normalized data before training.
    print(f'Training data after normalization:')
    print(f'  mean={X_train.mean():.4f} (target: ~0.0)')
    print(f'  std={X_train.std():.4f}  (target: ~1.0)')
    print(f'  min={X_train.min():.4f}, max={X_train.max():.4f}')

    return X_train, X_val, X_test, y_train, y_val, y_test, scaler


def build_and_train(X_train, X_val, y_train, y_val, num_classes: int):
    """
    Build, compile, and train the classifier with production-grade callbacks.
    """
    input_dim = X_train.shape[1]

    model = models.Sequential(
        [
            layers.Input(shape=(input_dim,), name='input'),
            layers.Dense(256, activation='relu', kernel_initializer='he_normal',
                         kernel_regularizer=keras.regularizers.l2(1e-4), name='dense_1'),
            layers.BatchNormalization(),
            layers.Dropout(0.3),
            layers.Dense(128, activation='relu', kernel_initializer='he_normal',
                         kernel_regularizer=keras.regularizers.l2(1e-4), name='dense_2'),
            layers.BatchNormalization(),
            layers.Dropout(0.2),
            # Binary classification: sigmoid output + binary_crossentropy.
            # For multi-class: softmax + sparse_categorical_crossentropy.
            layers.Dense(1, activation='sigmoid', name='output'),
        ],
        name='forge_binary_classifier'
    )

    model.compile(
        optimizer=keras.optimizers.Adam(learning_rate=5e-4),
        loss='binary_crossentropy',
        metrics=[
            'accuracy',
            keras.metrics.AUC(name='auc'),
            keras.metrics.Precision(name='precision'),
            keras.metrics.Recall(name='recall')
        ]
    )

    training_callbacks = [
        # Stop training when val_auc stops improving.
        # restore_best_weights=True rolls back to the best checkpoint automatically.
        callbacks.EarlyStopping(
            monitor='val_auc',
            patience=10,
            mode='max',
            restore_best_weights=True,
            verbose=1
        ),
        # Save best model checkpoint to disk during training.
        # If training is interrupted (OOM, preemption), the best weights are preserved.
        callbacks.ModelCheckpoint(
            filepath='checkpoints/forge_model_best.keras',
            monitor='val_auc',
            mode='max',
            save_best_only=True,
            verbose=1
        ),
        # Reduce learning rate when val_loss plateaus.
        # factor=0.5 halves the learning rate; min_lr prevents it from going to zero.
        callbacks.ReduceLROnPlateau(
            monitor='val_loss',
            factor=0.5,
            patience=5,
            min_lr=1e-6,
            verbose=1
        ),
    ]

    history = model.fit(
        X_train, y_train,
        validation_data=(X_val, y_val),
        epochs=100,          # EarlyStopping will stop before this in practice
        batch_size=256,
        callbacks=training_callbacks,
        verbose=1
    )

    return model, history


if __name__ == '__main__':
    # Synthetic dataset — replace with your actual data loading.
    np.random.seed(42)
    X_raw = np.random.randn(10_000, 40) * 1000  # Raw unscaled tabular features
    y_raw = (np.random.rand(10_000) > 0.5).astype(int)  # Binary labels

    X_train, X_val, X_test, y_train, y_val, y_test, scaler = prepare_data(X_raw, y_raw)

    # Benchmark first — always.
    rf_auc = benchmark_first(X_train, y_train, X_val, y_val)

    # Proceed with neural network only if we expect to improve on the baseline
    # or if interpretability requirements favor a NN over tree ensembles.
    nn_model, history = build_and_train(X_train, X_val, y_train, y_val, num_classes=2)

    # Save scaler alongside the model — inference requires both.
    import joblib
    joblib.dump(scaler, 'checkpoints/forge_scaler.pkl')
    print('Model and scaler saved to checkpoints/')