Reliability & Safety

# pip install torch torchvision adversarial-robustness-toolbox evidently

# ─── 1. FGSM Adversarial Attack ──────────────────────────────
import torch
import torch.nn as nn
import numpy as np
from torchvision import models, transforms
from PIL import Image

def fgsm_attack(model, loss_fn, image, label, epsilon=0.03):
    """Fast Gradient Sign Method adversarial attack."""
    image.requires_grad = True
    output = model(image)
    loss = loss_fn(output, label)
    model.zero_grad()
    loss.backward()
    perturbation = epsilon * image.grad.data.sign()
    adversarial = torch.clamp(image + perturbation, 0, 1)
    return adversarial

def pgd_attack(model, loss_fn, image, label, epsilon=0.03, alpha=0.01, num_steps=40):
    """Projected Gradient Descent — strongest first-order attack."""
    adv = image.clone().detach().requires_grad_(True)
    for _ in range(num_steps):
        output = model(adv)
        loss = loss_fn(output, label)
        loss.backward()
        with torch.no_grad():
            adv_update = adv + alpha * adv.grad.sign()
            # Project back onto epsilon-ball
            adv_update = image + torch.clamp(adv_update - image, -epsilon, epsilon)
            adv_update = torch.clamp(adv_update, 0, 1)
        adv = adv_update.detach().requires_grad_(True)
    return adv.detach()

# Simple CNN for demonstration
class SmallCNN(nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
        )
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(64 * 7 * 7, 256), nn.ReLU(), nn.Dropout(0.5),
            nn.Linear(256, num_classes),
        )
    def forward(self, x):
        return self.classifier(self.features(x))

model = SmallCNN()
loss_fn = nn.CrossEntropyLoss()

# Demo: measure accuracy under attack on a batch
batch_size = 32
x = torch.rand(batch_size, 1, 28, 28)  # synthetic MNIST-like
y = torch.randint(0, 10, (batch_size,))

with torch.no_grad():
    clean_preds = model(x).argmax(dim=1)
    clean_acc = (clean_preds == y).float().mean()
    print(f"Clean accuracy: {clean_acc:.3f}")

adv_x = fgsm_attack(model, loss_fn, x.clone().requires_grad_(True), y, epsilon=0.1)
with torch.no_grad():
    adv_preds = model(adv_x).argmax(dim=1)
    adv_acc = (adv_preds == y).float().mean()
    print(f"Accuracy under FGSM (ε=0.1): {adv_acc:.3f}")
    print(f"  Attack success rate: {1 - adv_acc:.3f}")

# ─── 2. Monte Carlo Dropout Uncertainty ──────────────────────
def mc_dropout_predict(model, x, n_passes=50):
    """Enable dropout during inference for uncertainty estimation."""
    model.train()  # Keep dropout active
    preds = torch.stack([
        torch.softmax(model(x), dim=1) for _ in range(n_passes)
    ])  # (n_passes, batch, classes)
    mean_pred = preds.mean(dim=0)          # Mean prediction
    uncertainty = preds.var(dim=0).sum(dim=1)  # Epistemic uncertainty
    return mean_pred, uncertainty

with torch.no_grad():
    mean_preds, uncertainties = mc_dropout_predict(model, x[:10])

print("\nMC Dropout Uncertainty Estimation:")
for i in range(5):
    pred_class = mean_preds[i].argmax().item()
    confidence = mean_preds[i].max().item()
    unc = uncertainties[i].item()
    flag = " ⚠ HIGH UNCERTAINTY" if unc > 0.3 else ""
    print(f"  Sample {i}: pred={pred_class}, conf={confidence:.3f}, uncertainty={unc:.4f}{flag}")

# ─── 3. Conformal Prediction ──────────────────────────────────
from sklearn.datasets import make_classification
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
import numpy as np

X_all, y_all = make_classification(n_samples=2000, n_features=20, n_classes=3,
                                    n_informative=15, random_state=42)
X_train, X_temp, y_train, y_temp = train_test_split(X_all, y_all, test_size=0.4)
X_cal, X_test, y_cal, y_test = train_test_split(X_temp, y_temp, test_size=0.5)

clf = RandomForestClassifier(n_estimators=100, random_state=42)
clf.fit(X_train, y_train)

# Conformal calibration: compute non-conformity scores on calibration set
cal_probs = clf.predict_proba(X_cal)
cal_scores = 1 - cal_probs[np.arange(len(y_cal)), y_cal]  # how "wrong" the model is on true class

# At alpha=0.1 → 90% coverage guarantee
alpha = 0.1
q_hat = np.quantile(cal_scores, 1 - alpha)
print(f"\nConformal Prediction (α={alpha}, target coverage={100*(1-alpha):.0f}%):")
print(f"  Calibrated quantile q̂ = {q_hat:.4f}")

# Generate prediction sets for test instances
test_probs = clf.predict_proba(X_test)
prediction_sets = (test_probs >= 1 - q_hat)  # include class if score ≥ threshold

# Evaluate coverage and set size
covered = prediction_sets[np.arange(len(y_test)), y_test].mean()
avg_set_size = prediction_sets.sum(axis=1).mean()
print(f"  Empirical coverage: {covered:.4f} (should be ≥ {1-alpha:.2f})")
print(f"  Average prediction set size: {avg_set_size:.2f} (smaller = more informative)")

# Show examples
print("\n  Example prediction sets:")
for i in range(3):
    pset = [j for j in range(3) if prediction_sets[i, j]]
    correct = y_test[i] in pset
    print(f"    Sample {i}: set={pset}, true={y_test[i]}, covered={correct}")

# ─── 4. Fail-Safe Circuit Breaker Pattern ────────────────────
import time
from dataclasses import dataclass
from typing import Callable, Optional

class CircuitState:
    CLOSED = "CLOSED"       # Normal operation
    OPEN = "OPEN"           # Failing — reject requests, use fallback
    HALF_OPEN = "HALF_OPEN" # Testing recovery

@dataclass
class CircuitBreaker:
    """Circuit breaker for AI model inference — fail fast and safely."""
    name: str
    failure_threshold: int = 5     # failures before opening
    recovery_timeout: float = 60.0 # seconds before trying half-open
    success_threshold: int = 2     # successes in half-open before closing

    failure_count: int = 0
    success_count: int = 0
    state: str = CircuitState.CLOSED
    last_failure_time: float = 0.0

    def call(self, fn: Callable, fallback: Callable, *args, **kwargs):
        """Call fn; use fallback if circuit is open."""
        if self.state == CircuitState.OPEN:
            if time.time() - self.last_failure_time > self.recovery_timeout:
                self.state = CircuitState.HALF_OPEN
                print(f"  [Circuit {self.name}] HALF-OPEN — testing recovery")
            else:
                print(f"  [Circuit {self.name}] OPEN — using fallback")
                return fallback(*args, **kwargs)

        try:
            result = fn(*args, **kwargs)
            self._on_success()
            return result
        except Exception as e:
            self._on_failure(e)
            print(f"  [Circuit {self.name}] Error: {e} — using fallback")
            return fallback(*args, **kwargs)

    def _on_success(self):
        if self.state == CircuitState.HALF_OPEN:
            self.success_count += 1
            if self.success_count >= self.success_threshold:
                self.state = CircuitState.CLOSED
                self.failure_count = 0
                self.success_count = 0
                print(f"  [Circuit {self.name}] CLOSED — recovered")
        elif self.state == CircuitState.CLOSED:
            self.failure_count = 0  # reset on success

    def _on_failure(self, error):
        self.failure_count += 1
        self.last_failure_time = time.time()
        if self.failure_count >= self.failure_threshold:
            self.state = CircuitState.OPEN
            print(f"  [Circuit {self.name}] OPEN after {self.failure_count} failures")

# AI model with circuit breaker + uncertainty-based escalation
call_count = 0
def ai_model_predict(input_data: dict) -> dict:
    global call_count
    call_count += 1
    if call_count % 3 == 0:  # simulate occasional failures
        raise RuntimeError("Model inference timeout")
    uncertainty = np.random.uniform(0, 1)
    pred = np.random.randint(0, 2)
    return {"prediction": pred, "confidence": 1 - uncertainty, "uncertainty": uncertainty}

def rule_based_fallback(input_data: dict) -> dict:
    """Simple rule-based fallback when AI model fails."""
    return {"prediction": 0, "confidence": 0.5, "source": "FALLBACK_RULES", "escalate": True}

circuit = CircuitBreaker("LoanModel", failure_threshold=3, recovery_timeout=5.0)

UNCERTAINTY_THRESHOLD = 0.7
HUMAN_REVIEW_THRESHOLD = 0.5

print("\nFail-Safe AI System (circuit breaker + uncertainty escalation):")
for i in range(8):
    result = circuit.call(ai_model_predict, rule_based_fallback, {"applicant_id": f"APP-{i}"})
    if result.get("source") == "FALLBACK_RULES" or result.get("escalate"):
        action = "→ HUMAN REVIEW (fallback triggered)"
    elif result["uncertainty"] > UNCERTAINTY_THRESHOLD:
        action = f"→ HUMAN REVIEW (high uncertainty={result['uncertainty']:.2f})"
    elif result["confidence"] < HUMAN_REVIEW_THRESHOLD:
        action = f"→ CONFIDENCE CHECK (conf={result['confidence']:.2f})"
    else:
        action = f"→ AUTO-APPROVE (conf={result['confidence']:.2f})"
    print(f"  Request {i+1}: {action}")