src/train.py

"""
train.py
Trains a prompt-injection classifier (TF-IDF + heuristics -> Logistic Regression)
and evaluates it with security-appropriate metrics (recall on the injection
class matters most -- a missed attack is worse than a false alarm).
"""

import pandas as pd
import numpy as np
import joblib
from sklearn.model_selection import train_test_split
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import (
    classification_report, confusion_matrix, precision_recall_curve, f1_score
)

from features import build_features


def load_data(path="data/prompts.csv"):
    df = pd.read_csv(path)
    return df["text"].tolist(), df["label"].tolist()


def train_and_evaluate():
    texts, labels = load_data()
    X_train_text, X_test_text, y_train, y_test = train_test_split(
        texts, labels, test_size=0.25, random_state=42, stratify=labels
    )

    vectorizer = TfidfVectorizer(
        ngram_range=(1, 2), max_features=3000, lowercase=True, stop_words="english"
    )

    X_train = build_features(X_train_text, vectorizer, fit=True)
    X_test = build_features(X_test_text, vectorizer, fit=False)

    models = {
        "logistic_regression": LogisticRegression(max_iter=1000, class_weight="balanced"),
        "random_forest": RandomForestClassifier(n_estimators=200, random_state=42, class_weight="balanced"),
    }

    results = {}
    for name, model in models.items():
        model.fit(X_train, y_train)
        y_pred = model.predict(X_test)
        report = classification_report(y_test, y_pred, target_names=["benign", "injection"], output_dict=True)
        cm = confusion_matrix(y_test, y_pred)
        results[name] = {"model": model, "report": report, "cm": cm}

        print(f"\n=== {name} ===")
        print(classification_report(y_test, y_pred, target_names=["benign", "injection"]))
        print("Confusion matrix (rows=true, cols=pred) [benign, injection]:")
        print(cm)

    # Pick the model with the best recall on the injection class
    # (in security, missing an attack is worse than a false alarm)
    best_name = max(results, key=lambda n: results[n]["report"]["injection"]["recall"])
    best_model = results[best_name]["model"]
    print(f"\nSelected model: {best_name} (highest recall on injection class)")

    joblib.dump(best_model, "models/classifier.joblib")
    joblib.dump(vectorizer, "models/vectorizer.joblib")
    print("Saved model + vectorizer to models/")

    return best_name, results


if __name__ == "__main__":
    train_and_evaluate()