Update model card: dataset citations, remove private repo refs

README.md

@@ -11,18 +11,20 @@ tags:
 language:
   - en
 pipeline_tag: text-classification
+datasets:
+  - deepset/prompt-injections
+  - hackaprompt/hackaprompt-dataset
+  - lakera-ai/gandalf_ignore_instructions
 ---
 
 # distilbert-multitask
 
-Multi-task DistilBERT classifier for conversational AI pipelines. Trained for use with the [rpchat engine](https://github.com/myemfar/rpchat) — a slot-based NPC dialogue pipeline for interactive fiction and games.
+Multi-task DistilBERT classifier for conversational AI pipelines in interactive fiction and games. Performs two classification tasks in a single forward pass:
 
-Performs two classification tasks in a single forward pass:
-
-| Task | Classes | Notes |
+| Task | Output | Notes |
 |---|---|---|
-| Dialogue act | 21 categories | See label list below |
-| Manipulation detection | Binary | Detects prompt injection / NPC takeover attempts |
+| Dialogue act | 21-class label | Classifies player utterance type |
+| Manipulation detection | Binary probability | Detects prompt injection / NPC takeover attempts |
 
 ## Dialogue Act Labels
 
@@ -30,44 +32,46 @@ Performs two classification tasks in a single forward pass:
 
 ## Model Details
 
-- **Base model:** `distilbert-base-uncased`
+- **Base model:** [distilbert-base-uncased](https://huggingface.co/distilbert/distilbert-base-uncased)
 - **Format:** ONNX (CPU inference via `onnxruntime`)
 - **Inference time:** ~10–15ms per input on CPU
 - **Input:** Single sentence (player utterance in a conversation)
-- **Output:** Dialogue act label + manipulation probability
+- **Output:** Dialogue act label + manipulation detection probability
 
 ## Usage
 
-The model is consumed automatically by the rpchat engine. On first use, it downloads here if not present locally:
-
-```python
-from huggingface_hub import snapshot_download
-snapshot_download(repo_id="myemfar/distilbert-multitask", local_dir="models/distilbert_multitask")
-```
-
-Then load via `onnxruntime`:
-
 ```python
-import onnxruntime
 import json
+import numpy as np
+import onnxruntime
+from huggingface_hub import snapshot_download
 from transformers import AutoTokenizer
 
-session = onnxruntime.InferenceSession("models/distilbert_multitask/model.onnx")
-tokenizer = AutoTokenizer.from_pretrained("models/distilbert_multitask")
+snapshot_download(repo_id="myemfar/distilbert-multitask", local_dir="./distilbert_multitask")
+
+session = onnxruntime.InferenceSession("./distilbert_multitask/model.onnx")
+tokenizer = AutoTokenizer.from_pretrained("./distilbert_multitask")
 
-with open("models/distilbert_multitask/label_map_da.json") as f:
+with open("./distilbert_multitask/label_map_da.json") as f:
     labels = {int(k): v for k, v in json.load(f).items()}
 
 inputs = tokenizer("Where is the tavern?", return_tensors="np")
 logits_da, logits_manip = session.run(None, dict(inputs))
 
-import numpy as np
-da_label = labels[int(np.argmax(logits_da))]      # "question"
+da_label = labels[int(np.argmax(logits_da))]           # "question"
 manip_prob = float(1 / (1 + np.exp(-logits_manip[0][0])))  # sigmoid
 ```
 
-## Training
+## Training Data
+
+### Dialogue Act Classification
+Synthetic training data generated via Claude across 21 conversational categories, curated for interactive fiction and RPG dialogue contexts. Approximately 2,000 labeled examples with targeted augmentation at category boundaries.
 
-Fine-tuned from `distilbert-base-uncased` on a curated dataset of ~2,000 conversational examples across the 21 dialogue act categories, plus a separate manipulation detection dataset targeting prompt injection patterns in NPC conversation contexts.
+### Manipulation / Prompt Injection Detection
+Fine-tuned on a combination of three public datasets plus domain-specific negative examples (in-character RPG dialogue):
 
-Training script: `data/train_distilbert.py` in the rpchat repository.
+| Dataset | License | Description |
+|---|---|---|
+| [deepset/prompt-injections](https://huggingface.co/datasets/deepset/prompt-injections) | CC BY 4.0 | Benign queries + prompt injection examples |
+| [hackaprompt/hackaprompt-dataset](https://huggingface.co/datasets/hackaprompt/hackaprompt-dataset) | Apache 2.0 | Red-teaming competition submissions |
+| [lakera-ai/gandalf_ignore_instructions](https://huggingface.co/datasets/lakera-ai/gandalf_ignore_instructions) | CC BY 4.0 | Instruction-override attempts from Lakera's Gandalf challenge |