"""Shared training loop with validation, TensorBoard logging, and BestModelCallback. Reusable by both `scripts/train.py` and `/retrain` endpoint in InferenceEngine. """ import os from datetime import datetime, timezone from typing import Optional import numpy as np import tensorflow as tf from src.trainers.trainer import BestModelCallback, sampled_softmax_loss_weighted def _batch_iterator( stu_feat: np.ndarray, sch_feat: np.ndarray, weights: np.ndarray, batch_size: int = 256, shuffle: bool = False, seed: Optional[int] = None, ) -> list[tuple[np.ndarray, np.ndarray, np.ndarray]]: """Yield mini-batches from numpy arrays. Args: stu_feat: (N, D_stu) student features. sch_feat: (N, D_sch) scholarship features. weights: (N,) per-sample feedback weights. batch_size: Mini-batch size. shuffle: If True, randomly permute indices before batching. seed: Optional RNG seed for reproducible shuffling. Yields: (stu_batch, sch_batch, weight_batch) triples. """ n = len(stu_feat) indices = np.arange(n) if shuffle: np.random.default_rng(seed).shuffle(indices) batches = [] for start in range(0, n, batch_size): end = min(start + batch_size, n) idx = indices[start:end] batches.append((stu_feat[idx], sch_feat[idx], weights[idx])) return batches def run_training( student_tower: tf.keras.Model, scholarship_tower: tf.keras.Model, train_stu_feat: np.ndarray, train_sch_feat: np.ndarray, train_weights: np.ndarray, val_stu_feat: np.ndarray, val_sch_feat: np.ndarray, val_weights: np.ndarray, optimizer: tf.keras.optimizers.Optimizer, temperature: float, epochs: int, checkpoint_dir: str, log_dir: str, tb_enabled: bool = False, k: int = 5, batch_size: int = 256, seed: Optional[int] = None, # ── Validation metrics (optional, for full Recall/NDCG/MRR) ─────── val_df: Optional[object] = None, stu_struct: Optional[np.ndarray] = None, sch_struct: Optional[np.ndarray] = None, stu_text_emb: Optional[np.ndarray] = None, sch_text_emb: Optional[np.ndarray] = None, stu_id_to_idx: Optional[dict] = None, sch_ids: Optional[list] = None, ) -> dict: """Run training loop with validation split, TensorBoard logging, and best-model saving. Args: student_tower / scholarship_tower: Tower models to finetune. train_*_feat / val_*_feat: Feature matrices for train/val splits. train_weights / val_weights: Per-sample feedback weights. optimizer: Keras optimizer. temperature: Contrastive temperature for sampled softmax loss. epochs: Number of training epochs. checkpoint_dir: Directory to save best model weights. log_dir: TensorBoard log directory. tb_enabled: Whether to write TensorBoard summaries. k: Recall@K value for validation metrics. batch_size: Mini-batch size for training (default 256). val_df: (optional) Validation DataFrame for Evaluator.compute_metrics(). If provided, full Recall/NDCG/MRR metrics are computed each epoch. stu_struct / sch_struct: Structured feature arrays. stu_text_emb / sch_text_emb: Text embedding arrays. stu_id_to_idx / sch_ids: ID mappings for retrieval evaluation. Returns: dict with keys: train_loss, val_loss, and optionally Recall/NDCG/MRR per epoch. """ # ── Setup callback ─────────────────────────────────────────────────────── os.makedirs(checkpoint_dir, exist_ok=True) cb = BestModelCallback( student_tower=student_tower, scholarship_tower=scholarship_tower, checkpoint_dir=checkpoint_dir, monitor=f"recall@{k}", # matches evaluator key (no "val_" prefix) ) # ── Setup TensorBoard ──────────────────────────────────────────────────── summary_writer = None if tb_enabled: log_dir = os.path.join(log_dir, f"retrain_{datetime.now(timezone.utc).strftime('%Y%m%d_%H%M%S')}") summary_writer = tf.summary.create_file_writer(log_dir) # ── Import evaluator lazily ──────────────────────────────────────────── from src.evaluators import Evaluator evaluator = Evaluator(k=k) # ── Pre-compute validation embeddings (once per epoch) ──────────────── val_stu_emb_all = student_tower(val_stu_feat, training=False).numpy() val_sch_emb_all = scholarship_tower(val_sch_feat, training=False).numpy() # ── Training loop ──────────────────────────────────────────────────────── cb.on_train_begin() metrics_history = {} for epoch in range(1, epochs + 1): # Train step — shuffle indices each epoch to prevent temporal ordering bias train_losses = [] for stu_b, sch_b, w_b in _batch_iterator(train_stu_feat, train_sch_feat, train_weights, batch_size, shuffle=True, seed=None if seed is None else seed + epoch): with tf.GradientTape() as tape: stu_emb = student_tower(stu_b, training=True) sch_emb = scholarship_tower(sch_b, training=True) train_loss = sampled_softmax_loss_weighted(stu_emb, sch_emb, w_b, temperature) all_vars = (student_tower.trainable_variables + scholarship_tower.trainable_variables) grads = tape.gradient(train_loss, all_vars) optimizer.apply_gradients(zip(grads, all_vars)) train_losses.append(float(train_loss)) train_loss_mean = float(np.mean(train_losses)) # Validation loss (batched) val_losses = [] for stu_b, sch_b, w_b in _batch_iterator(val_stu_feat, val_sch_feat, val_weights, batch_size): with tf.GradientTape(watch_accessed_variables=False): val_stu_emb = student_tower(stu_b, training=False) val_sch_emb = scholarship_tower(sch_b, training=False) val_loss = sampled_softmax_loss_weighted( val_stu_emb, val_sch_emb, w_b, temperature ) val_losses.append(float(val_loss)) val_loss_mean = float(np.mean(val_losses)) # Build logs dict — include retrieval metrics if evaluator data provided logs = { "train_loss": train_loss_mean, "val_loss": val_loss_mean, } if val_df is not None and stu_struct is not None: eval_metrics = evaluator.compute_metrics( df=val_df, student_tower=student_tower, scholarship_tower=scholarship_tower, stu_struct=stu_struct, sch_struct=sch_struct, stu_text_emb=stu_text_emb, sch_text_emb=sch_text_emb, stu_id_to_idx=stu_id_to_idx, sch_ids=sch_ids, ) logs.update(eval_metrics) cb.on_epoch_end(epoch - 1, logs=logs) # TensorBoard logging if tb_enabled and summary_writer is not None: with summary_writer.as_default(): tf.summary.scalar("loss/train_loss", train_loss_mean, step=epoch) tf.summary.scalar("loss/val_loss", val_loss_mean, step=epoch) for mk, mv in logs.items(): if mk not in ("train_loss", "val_loss"): tf.summary.scalar(f"{mk}", mv, step=epoch) # Print progress — include retrieval metrics when available if val_df is not None: print( f" Epoch {epoch}/{epochs} | " f"train_loss={train_loss_mean:.4f} val_loss={val_loss_mean:.4f} | " f"Recall@{k}={logs.get(f'recall@{k}', 0):.4f} " f"NDCG@{k}={logs.get(f'ndcg@{k}', 0):.4f} " f"MRR={logs.get('mrr', 0):.4f}" ) else: print( f" Epoch {epoch}/{epochs} — train_loss={train_loss_mean:.4f} val_loss={val_loss_mean:.4f}" ) metrics_history[epoch] = dict(logs) cb.on_train_end() # ── Close TensorBoard writer ───────────────────────────────────────────── if summary_writer is not None: summary_writer.close() return metrics_history