Spaces:
Running on Zero
Running on Zero
File size: 2,639 Bytes
81a4f72 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 | """
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()
|