distilbert-multitask

distilbert-multitask

Multi-task DistilBERT classifier for conversational AI pipelines in interactive fiction and games. Performs two classification tasks in a single forward pass:

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
• 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

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):