scripts/train_model.py

"""Train and evaluate the N.White AI operations intent classifier.

The script fits exactly one pre-declared CPU model on the training split only.
Validation and test records are used solely for post-fit evaluation.
"""

from __future__ import annotations

import csv
import hashlib
import json
from pathlib import Path

import joblib
from export_web_model import export_web_model
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import (
    accuracy_score,
    classification_report,
    confusion_matrix,
    f1_score,
    precision_recall_fscore_support,
)
from sklearn.pipeline import Pipeline


ROOT = Path(__file__).resolve().parents[1]
PACKAGE_ROOT = ROOT.parent
DATASET_ROOT = PACKAGE_ROOT / "nwhite-ai-operations-intent-dataset"
DATA_DIR = DATASET_ROOT / "data"
REPORT_DIR = ROOT / "reports"
RANDOM_STATE = 20260801

SMOKE_EXAMPLES = [
    {
        "expected_intent": "workflow_automation",
        "text": "Design a controlled sequence that assigns a new support request, records each hand-off and pauses for a supervisor before any consequential change.",
    },
    {
        "expected_intent": "document_processing",
        "text": "Extract the reference, date and category from a batch of synthetic forms, preserve page citations and queue uncertain fields for review.",
    },
    {
        "expected_intent": "analytics",
        "text": "Analyse the demonstration queue by service type and ageing band, showing median turnaround, missing values and reproducible aggregates.",
    },
    {
        "expected_intent": "knowledge_retrieval",
        "text": "Find the approved guidance for an internal procedure, cite the exact section and say clearly when the knowledge base does not contain an answer.",
    },
    {
        "expected_intent": "exception_handling",
        "text": "Quarantine the affected record when validation fails, preserve the evidence, alert its owner and define safe retry conditions without stopping other work.",
    },
    {
        "expected_intent": "human_escalation",
        "text": "Prepare a privacy-minimised review pack and route the uncertain high-impact request to a named manager for an explicit recorded decision.",
    },
    {
        "expected_intent": "api_integration",
        "text": "Synchronise approved synthetic records between two sandbox systems with schema validation, idempotency keys, bounded retries and reconciliation totals.",
    },
    {
        "expected_intent": "reporting",
        "text": "Generate a monthly operational summary with reporting period, volumes, overdue work, definitions, source coverage and data-quality caveats.",
    },
]


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for block in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(block)
    return digest.hexdigest()


def read_split(split: str) -> tuple[list[str], list[str], list[str]]:
    path = DATA_DIR / f"{split}.csv"
    with path.open("r", encoding="utf-8", newline="") as handle:
        rows = list(csv.DictReader(handle))
    return (
        [row["user_request"] for row in rows],
        [row["intent"] for row in rows],
        [row["id"] for row in rows],
    )


def metrics_for(y_true: list[str], y_pred: list[str], labels: list[str]) -> dict[str, object]:
    precision, recall, f1, support = precision_recall_fscore_support(
        y_true,
        y_pred,
        labels=labels,
        zero_division=0,
    )
    per_class = {
        label: {
            "precision": float(precision[index]),
            "recall": float(recall[index]),
            "f1": float(f1[index]),
            "support": int(support[index]),
        }
        for index, label in enumerate(labels)
    }
    return {
        "accuracy": float(accuracy_score(y_true, y_pred)),
        "macro_f1": float(f1_score(y_true, y_pred, labels=labels, average="macro", zero_division=0)),
        "weighted_f1": float(f1_score(y_true, y_pred, labels=labels, average="weighted", zero_division=0)),
        "per_class": per_class,
        "confusion_matrix": confusion_matrix(y_true, y_pred, labels=labels).tolist(),
        "classification_report": classification_report(
            y_true,
            y_pred,
            labels=labels,
            output_dict=True,
            zero_division=0,
        ),
    }


def write_predictions(
    path: Path,
    ids: list[str],
    texts: list[str],
    y_true: list[str],
    y_pred: list[str],
    probabilities: list[list[float]],
    labels: list[str],
) -> None:
    fields = ["id", "user_request", "true_intent", "predicted_intent", "correct", "predicted_probability"]
    with path.open("w", encoding="utf-8", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=fields, lineterminator="\n")
        writer.writeheader()
        for record_id, text, truth, prediction, probability_row in zip(ids, texts, y_true, y_pred, probabilities):
            probability = probability_row[labels.index(prediction)]
            writer.writerow(
                {
                    "id": record_id,
                    "user_request": text,
                    "true_intent": truth,
                    "predicted_intent": prediction,
                    "correct": truth == prediction,
                    "predicted_probability": f"{probability:.12f}",
                }
            )


def markdown_table(labels: list[str], matrix: list[list[int]]) -> str:
    short = [label.replace("_", " ") for label in labels]
    lines = [
        "| True \\ predicted | " + " | ".join(short) + " |",
        "|---|" + "---:|" * len(labels),
    ]
    for label, row in zip(short, matrix):
        lines.append(f"| {label} | " + " | ".join(str(value) for value in row) + " |")
    return "\n".join(lines)


def render_report(payload: dict[str, object]) -> str:
    labels = payload["labels"]
    validation = payload["validation"]
    test = payload["test"]
    test_rows = [
        "| Intent | Precision | Recall | F1 | Support |",
        "|---|---:|---:|---:|---:|",
    ]
    for label in labels:
        row = test["per_class"][label]
        test_rows.append(
            f"| `{label}` | {row['precision']:.4f} | {row['recall']:.4f} | {row['f1']:.4f} | {row['support']} |"
        )

    return f"""# Evaluation report

## Evaluation boundary

The TF-IDF plus logistic-regression pipeline was fitted once on **128 training records only**. No validation or test text was passed to `fit`. The fixed 32-record validation and 32-record test splits were used after fitting; no hyperparameter search or test-driven tuning was performed.

All records are synthetic and share an authoring framework. These scores measure performance only on the published synthetic split and do not establish production performance.

## Fixed model configuration

- Python: 3.11
- Vectoriser: word-level TF-IDF, unigrams and bigrams, Unicode accent stripping, sublinear term frequency
- Classifier: scikit-learn multinomial logistic regression, `solver=lbfgs`, `C=3.0`, `max_iter=2000`, `random_state=20260801`
- Training data SHA-256: `{payload['data']['train_sha256']}`
- Dataset version: 1.0.0

## Aggregate results

| Split | Records | Accuracy | Macro F1 | Weighted F1 |
|---|---:|---:|---:|---:|
| Validation | {payload['data']['validation_records']} | {validation['accuracy']:.4f} | {validation['macro_f1']:.4f} | {validation['weighted_f1']:.4f} |
| Test | {payload['data']['test_records']} | {test['accuracy']:.4f} | {test['macro_f1']:.4f} | {test['weighted_f1']:.4f} |

## Held-out test results by class

{chr(10).join(test_rows)}

## Test confusion matrix

Rows are true intents and columns are predicted intents, in the label order shown.

{markdown_table(labels, test['confusion_matrix'])}

## Additional inference smoke set

Eight independently worded, non-sensitive requests—one per intended class—were run after the model artefact was reloaded. Their predictions are stored in `sample_predictions.json`. This is a functional smoke test, not a benchmark.

## Interpretation and limitations

The corpus is deliberately balanced, compact and lexically explicit. Intent-specific control language appears across splits, and all examples were produced by one deterministic authoring system. Strong held-out scores can therefore reflect generator regularities and should not be generalised to natural request traffic, code-switching, African languages, speech transcripts, ambiguous multi-intent requests or organisation-specific terminology.

Before any real operational use, obtain lawful and representative local evaluation data, test misrouting costs, define abstention thresholds, monitor drift and keep an accountable human responsible for consequential actions. The model must not make insurance, financial, legal, medical, safety, employment or education decisions.
"""


def main() -> None:
    REPORT_DIR.mkdir(parents=True, exist_ok=True)
    x_train, y_train, train_ids = read_split("train")
    x_validation, y_validation, validation_ids = read_split("validation")
    x_test, y_test, test_ids = read_split("test")

    if len(set(train_ids) | set(validation_ids) | set(test_ids)) != 192:
        raise ValueError("Dataset splits overlap or do not cover 192 unique records")
    if set(train_ids) & (set(validation_ids) | set(test_ids)):
        raise ValueError("Evaluation data overlaps the training split")

    pipeline = Pipeline(
        steps=[
            (
                "tfidf",
                TfidfVectorizer(
                    lowercase=True,
                    strip_accents="unicode",
                    ngram_range=(1, 2),
                    min_df=1,
                    max_df=0.98,
                    sublinear_tf=True,
                    norm="l2",
                ),
            ),
            (
                "classifier",
                LogisticRegression(
                    C=3.0,
                    solver="lbfgs",
                    max_iter=2000,
                    random_state=RANDOM_STATE,
                ),
            ),
        ]
    )
    pipeline.fit(x_train, y_train)
    labels = [str(label) for label in pipeline.classes_]

    validation_predictions = pipeline.predict(x_validation).tolist()
    test_predictions = pipeline.predict(x_test).tolist()
    validation_probabilities = pipeline.predict_proba(x_validation).tolist()
    test_probabilities = pipeline.predict_proba(x_test).tolist()

    payload: dict[str, object] = {
        "model_name": "nwhite-ai-operations-intent-classifier",
        "model_version": "1.0.0",
        "evaluation_date": "2026-08-01",
        "algorithm": "TF-IDF (word unigrams and bigrams) with logistic regression",
        "labels": labels,
        "data": {
            "dataset": "nwhite-systems/nwhite-ai-operations-intent-dataset",
            "dataset_version": "1.0.0",
            "training_records": len(x_train),
            "validation_records": len(x_validation),
            "test_records": len(x_test),
            "train_sha256": sha256(DATA_DIR / "train.csv"),
            "validation_sha256": sha256(DATA_DIR / "validation.csv"),
            "test_sha256": sha256(DATA_DIR / "test.csv"),
        },
        "validation": metrics_for(y_validation, validation_predictions, labels),
        "test": metrics_for(y_test, test_predictions, labels),
        "limitations_note": "Synthetic held-out performance does not establish production performance.",
    }

    joblib.dump(pipeline, ROOT / "model.joblib", compress=3, protocol=5)
    export_web_model(ROOT / "model.joblib", ROOT / "web_model.json")
    with (ROOT / "metrics.json").open("w", encoding="utf-8", newline="\n") as handle:
        json.dump(payload, handle, ensure_ascii=False, indent=2)
        handle.write("\n")

    label_mapping = {
        "model_version": "1.0.0",
        "id_to_label": {str(index): label for index, label in enumerate(labels)},
        "label_to_id": {label: index for index, label in enumerate(labels)},
    }
    with (ROOT / "label_mapping.json").open("w", encoding="utf-8", newline="\n") as handle:
        json.dump(label_mapping, handle, ensure_ascii=False, indent=2)
        handle.write("\n")

    model_config = {
        "model_version": "1.0.0",
        "framework": "scikit-learn",
        "pipeline": ["TfidfVectorizer", "LogisticRegression"],
        "random_state": RANDOM_STATE,
        "training_split_only": True,
        "fit_record_count": len(x_train),
        "text_field": "user_request",
        "target_field": "intent",
        "probability_output": True,
        "abstention_threshold": None,
    }
    with (ROOT / "model_config.json").open("w", encoding="utf-8", newline="\n") as handle:
        json.dump(model_config, handle, ensure_ascii=False, indent=2)
        handle.write("\n")

    write_predictions(
        REPORT_DIR / "validation_predictions.csv",
        validation_ids,
        x_validation,
        y_validation,
        validation_predictions,
        validation_probabilities,
        labels,
    )
    write_predictions(
        REPORT_DIR / "test_predictions.csv",
        test_ids,
        x_test,
        y_test,
        test_predictions,
        test_probabilities,
        labels,
    )
    with (REPORT_DIR / "evaluation_report.md").open("w", encoding="utf-8", newline="\n") as handle:
        handle.write(render_report(payload))

    reloaded = joblib.load(ROOT / "model.joblib")
    sample_texts = [item["text"] for item in SMOKE_EXAMPLES]
    sample_predictions = reloaded.predict(sample_texts).tolist()
    sample_probabilities = reloaded.predict_proba(sample_texts).tolist()
    samples: list[dict[str, object]] = []
    for example, prediction, probability_row in zip(SMOKE_EXAMPLES, sample_predictions, sample_probabilities):
        ranked = sorted(zip(labels, probability_row), key=lambda item: item[1], reverse=True)
        samples.append(
            {
                "text": example["text"],
                "expected_intent": example["expected_intent"],
                "predicted_intent": prediction,
                "matches_expected": prediction == example["expected_intent"],
                "top_probabilities": [
                    {"intent": label, "probability": float(probability)} for label, probability in ranked[:3]
                ],
            }
        )
    with (ROOT / "sample_predictions.json").open("w", encoding="utf-8", newline="\n") as handle:
        json.dump({"description": "Functional smoke examples, not a benchmark", "examples": samples}, handle, ensure_ascii=False, indent=2)
        handle.write("\n")

    output = {
        "status": "trained_and_evaluated",
        "training_records": len(x_train),
        "validation_records": len(x_validation),
        "test_records": len(x_test),
        "validation_accuracy": payload["validation"]["accuracy"],
        "validation_macro_f1": payload["validation"]["macro_f1"],
        "test_accuracy": payload["test"]["accuracy"],
        "test_macro_f1": payload["test"]["macro_f1"],
        "reload_prediction_count": len(sample_predictions),
        "reload_smoke_matches": sum(item["matches_expected"] for item in samples),
    }
    print(json.dumps(output, indent=2))


if __name__ == "__main__":
    main()