""" NextActionPredictor training script. Data source: ~6 000 synthetic (domain, entity_flags, prior_contact → action) examples generated in-memory from DOMAIN_ACTION_PRIORS. No download required. Trains in < 30 seconds on CPU. Label assignment rules (derived from DOMAIN_ACTION_PRIORS): prior_contact = 0 → company_support (always try company first) prior_contact = 1 → priors[domain][1] (domain regulatory authority) prior_contact = 1 + has_AMOUNT + has_REF_ID + domain has a tertiary priors[2] + 20% chance → priors[domain][2] (escalate to legal / ombudsman) This creates a distribution the MLP learns naturally: company_support is the dominant class (~50%) while the 5 escalation authorities split the rest. CLI usage: python -m src.next_action.train --output_path models/next_action/model.pt """ from __future__ import annotations import argparse import logging import os import random import torch import torch.nn as nn from torch.utils.data import DataLoader, TensorDataset from src.next_action.model import ( ACTION2ID, DOMAIN_LABELS_ORDERED, ENTITY_NAMES, FEATURE_DIM, GUIDE_MLP, NUM_ACTIONS, ACTION_LABELS, build_feature_vector, ) from src.next_action.priors import DOMAIN_ACTION_PRIORS logger = logging.getLogger(__name__) # --------------------------------------------------------------------------- # Label assignment # --------------------------------------------------------------------------- def _label_for( domain: str, entity_flags: list[float], prior_contact: float, rng: random.Random, ) -> int: """Assign a training label from the priors and feature context.""" priors = DOMAIN_ACTION_PRIORS[domain] # ordered list of action strings if prior_contact < 0.5: action = priors[0] # always "company_support" else: has_amount = entity_flags[1] > 0.5 # ENTITY_NAMES[1] = "AMOUNT" has_ref_id = entity_flags[3] > 0.5 # ENTITY_NAMES[3] = "REF_ID" strong_case = has_amount and has_ref_id if strong_case and len(priors) >= 3 and rng.random() < 0.20: action = priors[2] # tertiary (legal / secondary ombudsman) else: action = priors[1] # primary regulatory authority return ACTION2ID[action] # --------------------------------------------------------------------------- # Dataset builder # --------------------------------------------------------------------------- def build_synthetic_dataset( n_samples: int = 6000, seed: int = 42, ) -> tuple[list[list[float]], list[int]]: """ Generate *n_samples* (feature_vector, label_id) pairs in memory. Distribution: ~50% prior_contact = 0 → company_support label ~50% prior_contact = 1 → domain authority or legal label Entity flags are sampled independently with p=0.55 of being set (biased toward having evidence, which is realistic for complaint scenarios). """ rng = random.Random(seed) X: list[list[float]] = [] y: list[int] = [] for _ in range(n_samples): domain = rng.choice(DOMAIN_LABELS_ORDERED) entity_flags = [float(rng.random() < 0.55) for _ in ENTITY_NAMES] prior_contact = float(rng.random() < 0.50) fv = build_feature_vector(domain, entity_flags, prior_contact) label = _label_for(domain, entity_flags, prior_contact, rng) X.append(fv) y.append(label) # Log class distribution for transparency counts: dict[int, int] = {} for lbl in y: counts[lbl] = counts.get(lbl, 0) + 1 dist = ", ".join( f"{ACTION_LABELS[k]}={v}" for k, v in sorted(counts.items()) ) logger.info("Synthetic dataset: %d examples — %s", n_samples, dist) return X, y # --------------------------------------------------------------------------- # Training entry point # --------------------------------------------------------------------------- def train(args: argparse.Namespace) -> None: """Train GUIDE_MLP and save checkpoint to args.output_path.""" logging.basicConfig(level=logging.INFO) # 1. Data X, y = build_synthetic_dataset(n_samples=args.n_samples) X_t = torch.tensor(X, dtype=torch.float32) y_t = torch.tensor(y, dtype=torch.long) dataset = TensorDataset(X_t, y_t) loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True) # 2. Model, loss, optimiser model = GUIDE_MLP() criterion = nn.CrossEntropyLoss() optimiser = torch.optim.Adam( model.parameters(), lr=args.lr, weight_decay=1e-4 ) # 3. Training loop (pure PyTorch — no Trainer — keeps startup time minimal) model.train() for epoch in range(1, args.epochs + 1): epoch_loss = 0.0 correct = 0 for xb, yb in loader: optimiser.zero_grad() logits = model(xb) loss = criterion(logits, yb) loss.backward() optimiser.step() epoch_loss += loss.item() * len(yb) correct += (logits.argmax(1) == yb).sum().item() if epoch % 10 == 0 or epoch == args.epochs: acc = correct / len(dataset) logger.info( "Epoch %3d/%d — loss: %.4f acc: %.3f", epoch, args.epochs, epoch_loss / len(dataset), acc, ) # 4. Evaluate final accuracy model.eval() with torch.no_grad(): all_logits = model(X_t) final_acc = (all_logits.argmax(1) == y_t).float().mean().item() logger.info("Training complete — final train acc: %.3f", final_acc) # 5. Save checkpoint os.makedirs(os.path.dirname(os.path.abspath(args.output_path)), exist_ok=True) torch.save( { "state_dict": model.state_dict(), "feature_dim": FEATURE_DIM, "num_actions": NUM_ACTIONS, "action_labels": ACTION_LABELS, "domain_labels": DOMAIN_LABELS_ORDERED, "entity_names": ENTITY_NAMES, }, args.output_path, ) logger.info("NextActionPredictor checkpoint saved to %s", args.output_path) # --------------------------------------------------------------------------- # CLI # --------------------------------------------------------------------------- def parse_args() -> argparse.Namespace: p = argparse.ArgumentParser(description="Train NextActionPredictor MLP") p.add_argument( "--output_path", default="models/next_action/model.pt", help="Path to save the .pt checkpoint", ) p.add_argument( "--n_samples", type=int, default=6000, help="Number of synthetic training examples", ) p.add_argument( "--epochs", type=int, default=60, help="Training epochs (default 60 → well under 30 s on CPU)", ) p.add_argument( "--batch_size", type=int, default=128, help="Mini-batch size", ) p.add_argument( "--lr", type=float, default=3e-3, help="Adam learning rate", ) return p.parse_args() if __name__ == "__main__": train(parse_args())