Deep Learning and Neural Networks

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
import mlflow, mlflow.pytorch

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Training on: {device}")
if torch.cuda.is_available():
    print(f"GPU: {torch.cuda.get_device_name(0)}")
    print(f"VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB")


class ContractRiskNet(nn.Module):
    """
    Feedforward network with embeddings for tabular government data.
    Uses embeddings for high-cardinality categoricals (hull_class, naics_code)
    and standard linear layers for numeric features.
    """
    def __init__(self, n_numeric: int, naics_vocab: int, hull_vocab: int,
                 embed_dim: int = 8, hidden_dim: int = 128):
        super().__init__()
        self.naics_embed = nn.Embedding(naics_vocab, embed_dim)
        self.hull_embed  = nn.Embedding(hull_vocab, embed_dim)

        input_dim = n_numeric + 2 * embed_dim
        self.network = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.BatchNorm1d(hidden_dim),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.BatchNorm1d(hidden_dim // 2),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(hidden_dim // 2, 1),
        )

    def forward(self, numeric, naics_idx, hull_idx):
        naics_emb = self.naics_embed(naics_idx)
        hull_emb  = self.hull_embed(hull_idx)
        x = torch.cat([numeric, naics_emb, hull_emb], dim=1)
        return self.network(x).squeeze(1)


model = ContractRiskNet(
    n_numeric=6, naics_vocab=500, hull_vocab=30
).to(device)

optimizer  = optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-4)
criterion  = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([5.0]).to(device))
scheduler  = optim.lr_scheduler.OneCycleLR(
    optimizer, max_lr=1e-2, steps_per_epoch=100, epochs=20
)

mlflow.set_experiment("/models/contract_risk_neural")

with mlflow.start_run(run_name="feedforward_baseline"):
    mlflow.log_params({
        "architecture": "feedforward_with_embeddings",
        "hidden_dim": 128, "embed_dim": 8, "dropout": 0.3,
        "optimizer": "AdamW", "lr": 1e-3, "epochs": 20,
    })

    for epoch in range(20):
        model.train()
        epoch_loss = 0.0
        for batch in train_loader:
            numeric, naics_idx, hull_idx, labels = [b.to(device) for b in batch]
            optimizer.zero_grad()
            logits = model(numeric, naics_idx, hull_idx)
            loss = criterion(logits, labels.float())
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            scheduler.step()
            epoch_loss += loss.item()

        mlflow.log_metric("train_loss", epoch_loss, step=epoch)

    mlflow.pytorch.log_model(model, artifact_path="model")

from pyspark.ml.torch.distributor import TorchDistributor
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler

def train_fn(n_gpus: int):
    """Training function executed on each GPU worker."""
    local_rank = int(os.environ.get("LOCAL_RANK", 0))
    dist.init_process_group(backend="nccl")
    torch.cuda.set_device(local_rank)

    model = build_facility_classifier(freeze_backbone=False)
    model = model.to(local_rank)
    model = DDP(model, device_ids=[local_rank])

    train_ds = SatelliteImageDataset(train_dir, transform=train_transforms)
    sampler  = DistributedSampler(train_ds)
    loader   = DataLoader(train_ds, batch_size=32, sampler=sampler, num_workers=4)

    # ... training loop ...
    dist.destroy_process_group()

n_gpus = torch.cuda.device_count()
distributor = TorchDistributor(
    num_processes=n_gpus,
    local_mode=False,
    use_gpu=True,
)
distributor.run(train_fn, n_gpus)

class OperationalInferencePipeline:
    """
    Wrapper enforcing DoD 3000.09 requirements:
    - Confidence threshold gate (below threshold → human review)
    - Full probability distribution logged with every prediction
    - Audit trail for every inference
    """
    def __init__(self, model, confidence_threshold: float = 0.80):
        self.model               = model
        self.confidence_threshold = confidence_threshold

    def predict(self, inputs, record_id: str) -> dict:
        probs      = self.model.predict_proba([inputs])[0]
        pred_class = int(np.argmax(probs))
        confidence = float(probs.max())

        result = {
            "record_id":            record_id,
            "predicted_class":      pred_class,
            "confidence":           confidence,
            "probability_dist":     probs.tolist(),   # ALWAYS log full distribution
            "requires_human_review": confidence < self.confidence_threshold,
            "timestamp":            pd.Timestamp.now().isoformat(),
        }

        # Log to audit trail — every inference, no exceptions
        self._log_to_audit_table(result)
        return result

    def _log_to_audit_table(self, result: dict) -> None:
        # On Foundry: write to an Ontology Action
        # On Databricks: append to a Delta audit table
        pass