""" Train the stress classifier (7->16->8->1 MLP) on a synthetic physiologically-motivated dataset, then re-export the ONNX. The synthetic target is a hand-built function of the 7 input features that encodes the same heuristics a clinician would use: - Low average eye aspect ratio -> drowsy - Low brow distance -> furrow / stress - High mouth tension -> clenched jaw - High eye asymmetry -> fatigue / neurological - Low mouth opening -> slack jaw - Late-evening / 3am time-of-day -> circadian low We add small Gaussian noise to inputs and target so the network has something to learn (not just a lookup table) and so the exported ONNX has interesting, well-distributed weights. Run from the hf-space/ directory: python train_stress_model.py Outputs: models/stress_model.onnx (re-exported, trained) models/stress_model_weights.npz (raw numpy weights) models/stress_training_data.npz (synthetic dataset) models/stress_metrics.json (train/val MAE, MSE) """ from __future__ import annotations import json import os import time from pathlib import Path import numpy as np HERE = Path(__file__).parent MODEL_DIR = HERE / "models" ONNX_PATH = MODEL_DIR / "stress_model.onnx" WEIGHTS_PATH = MODEL_DIR / "stress_model_weights.npz" DATA_PATH = MODEL_DIR / "stress_training_data.npz" METRICS_PATH = MODEL_DIR / "stress_metrics.json" RNG_SEED = 42 N_SAMPLES = 2000 VAL_FRACTION = 0.2 EPOCHS = 120 BATCH = 64 LR = 0.01 # Adam learning rate # ─── Synthetic target function ──────────────────────────────────────────────── # All ranges match the published model card. def sample_features(rng: np.random.Generator, n: int) -> np.ndarray: """Sample (n, 7) feature matrix in the documented physiological ranges.""" left_ear = rng.uniform(0.15, 0.45, n) right_ear = rng.uniform(0.15, 0.45, n) brow = rng.uniform(0.02, 0.06, n) mouth_t = rng.uniform(2.0, 12.0, n) eye_sym = np.abs(left_ear - right_ear) / ((left_ear + right_ear) / 2 + 0.001) mouth_w = rng.uniform(0.30, 0.60, n) # mouth width in normalized image coords mouth_h = rng.uniform(0.0, 0.20, n) # mouth height (opening) mouth_o = mouth_h / (mouth_w + 0.001) tod = rng.uniform(0.0, 1.0, n) return np.stack( [left_ear, right_ear, brow, mouth_t, eye_sym, mouth_o, tod], axis=1 ).astype(np.float32) def target_score(x: np.ndarray) -> np.ndarray: """Physiologically-motivated target in [0, 1]. Components are scaled so the sum roughly lives in [0, 1.5] before the final sigmoid, which produces a well-distributed target across "healthy" and "stressed" populations. """ left_ear, right_ear, brow, mouth_t, eye_sym, mouth_o, tod = x.T # Drowsiness: low EAR (eyes closing) avg_ear = (left_ear + right_ear) / 2 drowsy = np.clip((0.32 - avg_ear) / 0.10, 0, 1) * 0.30 # Brow furrow: low brow-to-eye distance furrow = np.clip((0.035 - brow) / 0.015, 0, 1) * 0.22 # Clenched jaw: high mouth_tension clench = np.clip((mouth_t - 6.0) / 4.0, 0, 1) * 0.12 # Eye asymmetry asym = np.clip((eye_sym - 0.10) / 0.10, 0, 1) * 0.12 # Slack jaw: low mouth_opening slack = np.clip((0.10 - mouth_o) / 0.10, 0, 1) * 0.10 # Circadian dip: 3am (tod=0.125) and 3pm (tod=0.625) bumps tod_h = tod * 24 tod_bump = 0.08 * ( np.exp(-((tod_h - 3.0) ** 2) / 4.0) + 0.6 * np.exp(-((tod_h - 15.0) ** 2) / 6.0) ) raw = drowsy + furrow + clench + asym + slack + tod_bump # Squash into [0, 1] with a soft logistic, but allow extremes return 1.0 / (1.0 + np.exp(-(raw * 4.0 - 1.4))) # ─── NumPy MLP with the same shape as the ZK circuit ───────────────────────── # Linear(7,16) -> ReLU -> Linear(16,8) -> ReLU -> Linear(8,1) -> Sigmoid def init_params(rng: np.random.Generator): def he(shape): fan_in = shape[1] return rng.normal(0, np.sqrt(2.0 / fan_in), shape).astype(np.float32) # Init b3 to the logit of the target mean (~0.5) so the network starts # at a sensible constant prediction rather than 0.5 with dead ReLUs. return { "W1": he((16, 7)) * 0.5, "b1": np.zeros(16, dtype=np.float32), "W2": he((8, 16)) * 0.5, "b2": np.zeros(8, dtype=np.float32), "W3": he((1, 8)) * 0.5, "b3": np.array([0.0], dtype=np.float32), } def forward(p, x): z1 = x @ p["W1"].T + p["b1"] a1 = np.maximum(0, z1) z2 = a1 @ p["W2"].T + p["b2"] a2 = np.maximum(0, z2) z3 = a2 @ p["W3"].T + p["b3"] return z3, (x, z1, a1, z2, a2, z3) # return logits; sigmoid applied at export time def sigmoid(z): return 1.0 / (1.0 + np.exp(-np.clip(z, -50, 50))) def bce_loss(y, t, eps=1e-7): y = np.clip(y, eps, 1 - eps) return float(-(t * np.log(y) + (1 - t) * np.log(1 - y)).mean()) def mse_loss(y, t): return float(((y - t) ** 2).mean()) def backward(p, cache, z3, t): """Backward through MSE on the raw logit output (no sigmoid in graph). We use MSE on the pre-sigmoid logit (with target also in logit space). The ONNX graph applies sigmoid at the end, so the probability output is bounded in [0, 1]. Training in logit space avoids the saturation problem of sigmoid + small gradients. """ x, z1, a1, z2, a2, _ = cache n = z3.shape[0] if t.ndim == 1: t = t.reshape(-1, 1) # Convert target probability to logit, clamp to avoid inf t_p = np.clip(t, 1e-5, 1 - 1e-5) t_logit = np.log(t_p / (1.0 - t_p)) # MSE on logits: dL/dz3 = 2(z3 - t_logit) / n dL_dz3 = 2.0 * (z3 - t_logit) / n dW3 = dL_dz3.T @ a2 db3 = dL_dz3.sum(axis=0) dL_da2 = dL_dz3 @ p["W3"] dL_dz2 = dL_da2 * (z2 > 0) dW2 = dL_dz2.T @ a1 db2 = dL_dz2.sum(axis=0) dL_da1 = dL_dz2 @ p["W2"] dL_dz1 = dL_da1 * (z1 > 0) dW1 = dL_dz1.T @ x db1 = dL_dz1.sum(axis=0) return {"W1": dW1, "b1": db1, "W2": dW2, "b2": db2, "W3": dW3, "b3": db3} def train(X, T, *, epochs=EPOCHS, batch=BATCH, lr=LR, seed=RNG_SEED): rng = np.random.default_rng(seed) p = init_params(rng) n = X.shape[0] # Adam state m = {k: np.zeros_like(v) for k, v in p.items()} v = {k: np.zeros_like(v) for k, v in p.items()} b1, b2, eps = 0.9, 0.999, 1e-8 history = [] for epoch in range(epochs): idx = rng.permutation(n) Xs, Ts = X[idx], T[idx] for i in range(0, n, batch): xb, tb = Xs[i:i + batch], Ts[i:i + batch] z3, cache = forward(p, xb) grads = backward(p, cache, z3, tb) for k in p: m[k] = b1 * m[k] + (1 - b1) * grads[k] v[k] = b2 * v[k] + (1 - b2) * (grads[k] ** 2) m_hat = m[k] / (1 - b1 ** (epoch + 1)) v_hat = v[k] / (1 - b2 ** (epoch + 1)) p[k] = p[k] - lr * m_hat / (np.sqrt(v_hat) + eps) if epoch % max(1, epochs // 20) == 0 or epoch == epochs - 1: z3_full, _ = forward(p, X) y_full = sigmoid(z3_full) T_col = T.reshape(-1, 1) if T.ndim == 1 else T loss = mse_loss(y_full, T_col) mae = float(np.mean(np.abs(y_full - T_col))) history.append({"epoch": epoch, "loss": loss, "mae": mae}) print(f" epoch {epoch:4d} loss={loss:.5f} mae={mae:.4f}") return p, history # ─── ONNX export with the trained weights ───────────────────────────────────── def export_onnx(p): import onnx from onnx import helper, TensorProto, numpy_helper initializers = [] nodes = [] def add_linear(name, in_f, out_f, W, b): W_init = numpy_helper.from_array(W.astype(np.float32), name=f"{name}_W") b_init = numpy_helper.from_array(b.astype(np.float32), name=f"{name}_b") matmul = helper.make_node( "Gemm", [f"{name}_in", f"{name}_W", f"{name}_b"], [f"{name}_out"], transB=1, ) return matmul, [W_init, b_init] nodes.append(helper.make_node("Identity", ["input"], ["l1_in"])) n, inits = add_linear("l1", 7, 16, p["W1"], p["b1"]) nodes.append(n) initializers.extend(inits) nodes.append(helper.make_node("Relu", ["l1_out"], ["r1_out"])) nodes.append(helper.make_node("Identity", ["r1_out"], ["l2_in"])) n, inits = add_linear("l2", 16, 8, p["W2"], p["b2"]) nodes.append(n) initializers.extend(inits) nodes.append(helper.make_node("Relu", ["l2_out"], ["r2_out"])) nodes.append(helper.make_node("Identity", ["r2_out"], ["l3_in"])) n, inits = add_linear("l3", 8, 1, p["W3"], p["b3"]) nodes.append(n) initializers.extend(inits) nodes.append(helper.make_node("Sigmoid", ["l3_out"], ["output"])) graph = helper.make_graph( nodes, "stress_mlp", [helper.make_tensor_value_info("input", TensorProto.FLOAT, ["batch", 7])], [helper.make_tensor_value_info("output", TensorProto.FLOAT, ["batch", 1])], initializer=initializers, ) model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 10)]) model.ir_version = 7 onnx.save(model, str(ONNX_PATH)) # ─── Main ───────────────────────────────────────────────────────────────────── def main() -> None: MODEL_DIR.mkdir(parents=True, exist_ok=True) rng = np.random.default_rng(RNG_SEED) X = sample_features(rng, N_SAMPLES) # Add small input noise so the network cannot memorize X = X + rng.normal(0, 0.005, X.shape).astype(np.float32) T = target_score(X).astype(np.float32) # Add small target noise T = np.clip(T + rng.normal(0, 0.03, T.shape).astype(np.float32), 0, 1) # Train/val split n_val = int(N_SAMPLES * VAL_FRACTION) perm = rng.permutation(N_SAMPLES) val_idx, tr_idx = perm[:n_val], perm[n_val:] X_tr, T_tr = X[tr_idx], T[tr_idx] X_val, T_val = X[val_idx], T[val_idx] print(f"Training stress MLP on {N_SAMPLES} synthetic samples...") t0 = time.time() p, history = train(X_tr, T_tr) dt = time.time() - t0 z3_val, _ = forward(p, X_val) y_val_p = sigmoid(z3_val) T_val_col = T_val.reshape(-1, 1) if T_val.ndim == 1 else T_val val_mae = float(np.mean(np.abs(y_val_p - T_val_col))) val_mse = float(np.mean((y_val_p - T_val_col) ** 2)) print(f"\nTrained in {dt:.1f}s. Val MAE={val_mae:.4f} Val MSE={val_mse:.4f}") # Save weights, data, metrics np.savez(WEIGHTS_PATH, **p) np.savez(DATA_PATH, X=X, T=T, val_idx=val_idx, tr_idx=tr_idx) metrics = { "n_samples": int(N_SAMPLES), "n_train": int(len(tr_idx)), "n_val": int(len(val_idx)), "epochs": int(EPOCHS), "batch": int(BATCH), "lr": float(LR), "val_mae": val_mae, "val_mse": val_mse, "train_history_tail": history[-5:], "seed": int(RNG_SEED), "train_seconds": round(dt, 2), } METRICS_PATH.write_text(json.dumps(metrics, indent=2)) # Re-export ONNX print(f"Exporting trained ONNX to {ONNX_PATH}...") export_onnx(p) print(f"ONNX size: {ONNX_PATH.stat().st_size} bytes") # Round-trip check via onnxruntime try: import onnxruntime as ort sess = ort.InferenceSession(str(ONNX_PATH)) test = X_val[:5] y_onnx = sess.run(None, {sess.get_inputs()[0].name: test})[0] z3_np, _ = forward(p, test) y_np = sigmoid(z3_np) max_diff = float(np.max(np.abs(y_onnx - y_np))) print(f"ONNX vs numpy max diff: {max_diff:.2e}") except Exception as e: print(f"Round-trip check skipped: {e}") if __name__ == "__main__": main()