File size: 2,199 Bytes
34e621c | 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 | """
Dataset module -- Member A (Huy)
PyTorch Dataset for prompt classification (Jailbreak vs Benign).
"""
import torch
from torch.utils.data import Dataset
class PromptDataset(Dataset):
"""
Dataset cho bai toan phan loai prompt jailbreak.
- texts: list of raw text strings
- labels: list of int (0=benign, 1=jailbreak) or str labels; None for inference
- tokenizer: HuggingFace tokenizer; None returns raw text dict
- max_length: max token length for truncation
"""
LABEL_MAP = {
"benign": 0,
"jailbreak": 1,
"prompt_injection": 2,
}
ID2LABEL = {0: "BENIGN", 1: "JAILBREAK", 2: "PROMPT_INJECTION"}
LABEL2ID = {"BENIGN": 0, "JAILBREAK": 1, "PROMPT_INJECTION": 2}
def __init__(self, texts, labels=None, tokenizer=None, max_length: int = 256):
self.texts = [str(t) for t in texts]
self.tokenizer = tokenizer
self.max_length = max_length
if labels is not None:
processed_labels = []
for lab in labels:
if isinstance(lab, str):
processed_labels.append(self.LABEL_MAP.get(lab.lower(), 0))
else:
processed_labels.append(int(lab))
self.labels = processed_labels
else:
self.labels = None
def __len__(self) -> int:
return len(self.texts)
def __getitem__(self, idx: int) -> dict:
text = self.texts[idx]
if self.tokenizer is not None:
encoding = self.tokenizer(
text,
truncation=True,
max_length=self.max_length,
padding="max_length",
return_tensors="pt",
)
item = {
"input_ids": encoding["input_ids"].squeeze(0),
"attention_mask": encoding["attention_mask"].squeeze(0),
}
if "token_type_ids" in encoding:
item["token_type_ids"] = encoding["token_type_ids"].squeeze(0)
else:
item = {"text": text}
if self.labels is not None:
item["label"] = torch.tensor(self.labels[idx], dtype=torch.long)
return item
|