| ---
|
| license: apache-2.0
|
| base_model: distilbert-base-uncased
|
| tags:
|
| - text-classification
|
| - dialogue-act-classification
|
| - multi-task-learning
|
| - onnx
|
| - roleplay
|
| - game-ai
|
| 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 in interactive fiction and games. Performs two classification tasks in a single forward pass:
|
|
|
| | Task | Output | Notes |
|
| |---|---|---|
|
| | Dialogue act | 21-class label | Classifies player utterance type |
|
| | Manipulation detection | Binary probability | Detects prompt injection / NPC takeover attempts |
|
|
|
| ## Dialogue Act Labels
|
|
|
| `accusation`, `acknowledgment`, `action`, `agree`, `command`, `conditional`, `confession`, `disagree`, `emote`, `farewell`, `flirt`, `greeting`, `hedge`, `hostile`, `intent`, `offer`, `opinion`, `out_of_character`, `question`, `statement`, `yes_no_question`
|
|
|
| ## Model Details
|
|
|
| - **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 detection probability
|
|
|
| ## Usage
|
|
|
| ```python
|
| import json
|
| import numpy as np
|
| import onnxruntime
|
| from huggingface_hub import snapshot_download
|
| from transformers import AutoTokenizer
|
|
|
| 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("./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))
|
|
|
| da_label = labels[int(np.argmax(logits_da))] # "question"
|
| manip_prob = float(1 / (1 + np.exp(-logits_manip[0][0]))) # sigmoid
|
| ```
|
|
|
| ## 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.
|
|
|
| ### Manipulation / Prompt Injection Detection
|
| Fine-tuned on a combination of three public datasets plus domain-specific negative examples (in-character RPG dialogue):
|
|
|
| | 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 |
|
|
|
