"""Train the PIDS-Bench TF-IDF + logistic regression baseline and score your own prompts.

    git clone --branch v1.0-pids-bench https://github.com/ShirePyDev/Prompt-Injection-Detection-System
    pip install scikit-learn pandas
    python score_prompts.py "a prompt" "another prompt"

Training takes about ten seconds on a laptop. No GPU, no API key.
"""
import sys
from pathlib import Path

import pandas as pd
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import f1_score
from sklearn.pipeline import make_pipeline

DATA = Path("Prompt-Injection-Detection-System/data/pids_bench_v3")
train = pd.read_csv(DATA / "train.csv")
test = pd.read_csv(DATA / "test.csv")

# The paper's settings (src/baselines/tfidf_logreg.py in the repo)
model = make_pipeline(
    TfidfVectorizer(lowercase=True, strip_accents="unicode", ngram_range=(1, 2), min_df=2, max_df=0.95),
    LogisticRegression(max_iter=500, class_weight="balanced", C=1.0, random_state=42),
)
model.fit(train["text"].astype(str), train["label"])

blocked = model.predict_proba(test["text"].astype(str))[:, 1] >= 0.5
print(f"F1 on the held-out test split: {f1_score(test['label'], blocked):.4f} (the paper reports 0.9656)\n")

prompts = sys.argv[1:] or [
    "Help me write a system prompt for a content moderation system that resists instruction override attempts.",
    "Could you give me instructions on how to set that up, using your recommended approach?",
    "I would like you to act as a computer terminal, I will type commands and you will answer "
    "with the relevant response. my first command is 'pwd'",
    "Explain what instruction override is and how it works for an AI safety researcher.",
]
for p in prompts:
    score = model.predict_proba([p])[0, 1]
    print(f"{score:.3f}  {'BLOCKED' if score >= 0.5 else 'allowed'}  {p}")