Special token bug fix

cfacbcc verified 27 days ago

19.1 kB

from __future__ import annotations

from functools import partial
from typing import Optional

import torch
from torch.utils.data import DataLoader, Dataset
from transformers import PreTrainedTokenizer

from .prompt import (
    build_prompt,
    build_prompt_parts,
    build_target,
    canonicalize_violated_categories,
)

PROMPT_BLOCK_ID = 1
RESPONSE_BLOCK_ID = 2
MASK_TOKEN_TEXT = "<|mdm_mask|>"


class SafetyDataset(Dataset):
    """Materialize the safe/unsafe prompt pairs used by masked-scope scoring."""

    def __init__(
        self,
        examples: list[dict],
        tokenizer: PreTrainedTokenizer,
        max_seq_length: int = 2048,
        category_augmentation: bool = False,
        category_augmentation_prob: float = 0.0,
        prompt_only_augmentation_prob: float = 0.0,
        prompt_format: str = "aegis_taxonomy",
        gen_only: bool = False,
        skip_cls: bool = False,
    ):
        del category_augmentation, category_augmentation_prob, prompt_only_augmentation_prob
        del skip_cls
        self.examples = examples
        self.tokenizer = tokenizer
        self.max_seq_length = max_seq_length
        self.prompt_format = prompt_format
        self.gen_only = gen_only

    def __len__(self):
        return len(self.examples)

    def _chat_token_count(
        self, content: str, trusted_span: Optional[tuple[int, int]] = None
    ) -> int:
        prompt_str = self.tokenizer.apply_chat_template(
            [{"role": "user", "content": content}],
            add_generation_prompt=True,
            tokenize=False,
        )
        ids, _ = self._encode_chat(prompt_str, content, trusted_span)
        return len(ids) + len(self.tokenizer("")["input_ids"])

    def _encode_chat(
        self,
        prompt_str: str,
        content: str,
        trusted_span: Optional[tuple[int, int]] = None,
    ) -> tuple[list[int], list[tuple[int, int]]]:
        """Tokenize a chat-formatted prompt without parsing special tokens in content.
        Special-token strings typed by a user or agent (e.g. ``<|eot_id|>``,
        ``<|mdm_mask|>``) stay plain text; only the chat template's own control
        tokens are parsed. ``trusted_span`` (offsets within ``content``) marks
        release-generated text, such as the synthetic mask response, whose
        control tokens are kept.
        """
        content_start = prompt_str.find(content)
        if content_start < 0:
            raise ValueError("Classification prompt not found inside chat template")
        content_end = content_start + len(content)
        segments = [(0, content_start, False)]
        if trusted_span is None:
            segments.append((content_start, content_end, True))
        else:
            trusted_start = content_start + trusted_span[0]
            trusted_end = content_start + trusted_span[1]
            segments += [
                (content_start, trusted_start, True),
                (trusted_start, trusted_end, False),
                (trusted_end, content_end, True),
            ]
        segments.append((content_end, len(prompt_str), False))

        input_ids: list[int] = []
        offsets: list[tuple[int, int]] = []
        for start, end, split_special in segments:
            if start == end:
                continue
            encoding = self.tokenizer(
                prompt_str[start:end],
                add_special_tokens=False,
                return_offsets_mapping=True,
                split_special_tokens=split_special,
            )
            input_ids.extend(encoding["input_ids"])
            offsets.extend((s + start, e + start) for s, e in encoding["offset_mapping"])
        return input_ids, offsets

    def _gen_total_len(
        self,
        verdict_text: str,
        user_message: str,
        agent_response: Optional[str],
        categories: Optional[list[str]],
        violated_categories: Optional[list[str]] = None,
        synthetic_response: bool = False,
    ) -> int:
        prefix, conversation, suffix = build_prompt_parts(
            user_message,
            agent_response,
            categories=categories,
            prompt_format=self.prompt_format,
        )
        trusted_span = None
        if synthetic_response and agent_response is not None:
            trusted_end = len(prefix) + len(conversation)
            trusted_span = (trusted_end - len(agent_response), trusted_end)
        prompt_len = self._chat_token_count(prefix + conversation + suffix, trusted_span)
        if verdict_text == "unsafe" and violated_categories:
            label_text = build_target(
                False,
                violated_categories,
                prompt_format=self.prompt_format,
            )
        else:
            label_text = verdict_text
        label_ids = self.tokenizer(label_text, add_special_tokens=False)["input_ids"]
        eos_extra = 1 if self.tokenizer.eos_token_id is not None else 0
        return prompt_len + len(label_ids) + eos_extra

    def _truncate_to_budget(
        self,
        user_message: str,
        agent_response: Optional[str],
        categories: Optional[list[str]],
        violated_categories: Optional[list[str]] = None,
        allow_agent_drop: bool = True,
    ) -> tuple[str, Optional[str]]:
        """Trim the response first, then the user text, until the pair fits."""
        while True:
            gen_safe_total = self._gen_total_len(
                "safe",
                user_message,
                agent_response,
                categories,
                synthetic_response=not allow_agent_drop,
            )
            gen_unsafe_total = self._gen_total_len(
                "unsafe",
                user_message,
                agent_response,
                categories,
                violated_categories=violated_categories,
                synthetic_response=not allow_agent_drop,
            )
            overflow = max(gen_safe_total, gen_unsafe_total) - self.max_seq_length
            if overflow <= 0:
                return user_message, agent_response

            if agent_response and allow_agent_drop:
                agent_ids = self.tokenizer(agent_response, add_special_tokens=False)["input_ids"]
                drop = min(max(overflow, 1), len(agent_ids))
                agent_ids = agent_ids[:-drop]
                agent_response = (
                    self.tokenizer.decode(agent_ids, skip_special_tokens=True) if agent_ids else None
                )
                continue

            user_ids = self.tokenizer(user_message, add_special_tokens=False)["input_ids"]
            drop = min(max(overflow, 1), len(user_ids))
            user_ids = user_ids[:-drop]
            user_message = self.tokenizer.decode(user_ids, skip_special_tokens=True) if user_ids else ""

    def _build_gen_sequence(
        self,
        verdict_text: str,
        user_message: str,
        agent_response: Optional[str],
        categories: Optional[list[str]],
        violated_categories: Optional[list[str]] = None,
        synthetic_response: bool = False,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        prefix, conversation_text, suffix = build_prompt_parts(
            user_message,
            agent_response,
            categories=categories,
            prompt_format=self.prompt_format,
        )
        prompt_text = prefix + conversation_text + suffix
        prompt_str = self.tokenizer.apply_chat_template(
            [{"role": "user", "content": prompt_text}],
            add_generation_prompt=True,
            tokenize=False,
        )
        content_start = prompt_str.find(prompt_text)
        if content_start < 0:
            raise ValueError("Classification prompt not found inside chat template")
        conversation_start = content_start + len(prefix)

        user_span = f"User: {user_message}"
        user_start = conversation_start
        user_end = user_start + len(user_span)

        response_start = None
        response_end = None
        if agent_response is not None:
            response_span = f"\nAgent: {agent_response}"
            response_start = user_end
            response_end = response_start + len(response_span)

        trusted_span = None
        if synthetic_response and agent_response is not None:
            trusted_end = response_end - content_start
            trusted_span = (trusted_end - len(agent_response), trusted_end)
        prompt_ids, prompt_offsets = self._encode_chat(
            prompt_str, prompt_text, trusted_span
        )

        prompt_block_ids = []
        for start, end in prompt_offsets:
            in_user = end > user_start and start < user_end
            in_response = (
                response_start is not None
                and response_end is not None
                and end > response_start
                and start < response_end
            )
            if in_user:
                prompt_block_ids.append(PROMPT_BLOCK_ID)
            elif in_response:
                prompt_block_ids.append(RESPONSE_BLOCK_ID)
            else:
                prompt_block_ids.append(0)

        include_cats = verdict_text == "unsafe" and violated_categories
        if include_cats:
            label_text = build_target(
                False,
                violated_categories,
                prompt_format=self.prompt_format,
            )
        else:
            label_text = verdict_text

        label_encoding = self.tokenizer(
            label_text,
            add_special_tokens=False,
            return_offsets_mapping=True,
        )
        label_ids = label_encoding["input_ids"]
        label_offsets = label_encoding["offset_mapping"]

        if include_cats:
            char_blocks = [0] * len(label_text)
            pos = len("unsafe") + 1
            for cat_idx, cat in enumerate(violated_categories):
                for offset in range(len(cat)):
                    if pos + offset < len(char_blocks):
                        char_blocks[pos + offset] = 3 + cat_idx
                pos += len(cat) + 1
            label_block_ids = []
            for start, end in label_offsets:
                block_id = 0
                for char_idx in range(start, min(end, len(char_blocks))):
                    block_id = max(block_id, char_blocks[char_idx])
                label_block_ids.append(block_id)
        else:
            label_block_ids = [0] * len(label_ids)

        input_ids = list(prompt_ids) + list(label_ids)
        block_ids = prompt_block_ids + label_block_ids

        eos_id = self.tokenizer.eos_token_id
        if eos_id is not None:
            input_ids.append(eos_id)
            block_ids.append(0)

        return (
            torch.tensor(input_ids, dtype=torch.long),
            torch.tensor(block_ids, dtype=torch.long),
        )

    def __getitem__(self, idx):
        ex = self.examples[idx]
        user_message = str(ex["user_message"])
        agent_response = ex.get("agent_response")
        is_safe = bool(ex.get("is_safe", True))
        violated = canonicalize_violated_categories(ex.get("violated_categories", []))

        had_response = agent_response is not None
        rendered_agent_response = str(agent_response) if had_response else MASK_TOKEN_TEXT
        user_message, rendered_agent_response = self._truncate_to_budget(
            user_message,
            rendered_agent_response,
            None,
            violated_categories=violated,
            allow_agent_drop=had_response,
        )
        has_response = int(had_response and rendered_agent_response is not None)
        prompt_text = build_prompt(
            user_message,
            rendered_agent_response,
            prompt_format=self.prompt_format,
        )

        gen_safe_input_ids, gen_safe_block_ids = self._build_gen_sequence(
            "safe",
            user_message,
            rendered_agent_response,
            None,
            synthetic_response=not had_response,
        )
        gen_unsafe_input_ids, gen_unsafe_block_ids = self._build_gen_sequence(
            "unsafe",
            user_message,
            rendered_agent_response,
            None,
            synthetic_response=not had_response,
        )
        hazard_label_mask = int((not is_safe) and bool(violated))
        if hazard_label_mask:
            hazard_input_ids, hazard_block_ids = self._build_gen_sequence(
                "unsafe",
                user_message,
                rendered_agent_response,
                None,
                violated_categories=violated,
                synthetic_response=not had_response,
            )
        else:
            hazard_input_ids = gen_unsafe_input_ids
            hazard_block_ids = gen_unsafe_block_ids

        prompt_label = (ex.get("prompt_label") or "").strip().lower()
        response_label = (ex.get("response_label") or "").strip().lower()
        prompt_label_known = prompt_label in {"safe", "unsafe"}
        response_label_known = had_response and response_label in {"safe", "unsafe"}
        if not prompt_label_known and not had_response:
            prompt_label_known = True
        if not response_label_known and had_response:
            response_label_known = True

        prompt_is_unsafe = (
            int(prompt_label == "unsafe")
            if prompt_label in {"safe", "unsafe"}
            else int((not is_safe) and not had_response)
        )
        response_is_unsafe = (
            int(response_label == "unsafe")
            if response_label in {"safe", "unsafe"}
            else int((not is_safe) and had_response)
        )

        return {
            "gen_safe_input_ids": gen_safe_input_ids,
            "gen_safe_block_ids": gen_safe_block_ids,
            "gen_safe_real_length": len(gen_safe_input_ids),
            "gen_unsafe_input_ids": gen_unsafe_input_ids,
            "gen_unsafe_block_ids": gen_unsafe_block_ids,
            "gen_unsafe_real_length": len(gen_unsafe_input_ids),
            "hazard_input_ids": hazard_input_ids,
            "hazard_block_ids": hazard_block_ids,
            "hazard_real_length": len(hazard_input_ids),
            "hazard_label_mask": hazard_label_mask,
            "hazard_eval_mask": int(is_safe or bool(violated)),
            "hazard_targets": list(violated),
            "prompt_text": prompt_text,
            "is_unsafe": int(not is_safe),
            "conv_label_mask": int(ex.get("conv_label_mask", 1)),
            "prompt_is_unsafe": prompt_is_unsafe,
            "prompt_label_mask": int(prompt_label_known),
            "response_is_unsafe": response_is_unsafe,
            "response_label_mask": int(response_label_known),
            "has_response": has_response,
            "example_index": int(ex.get("example_index", idx)),
            "eval_position": int(ex.get("eval_position", -1)),
        }


def _pad_tensor_batch(
    batch: list[dict],
    ids_key: str,
    blocks_key: str,
    length_key: str,
    pad_id: int,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    max_len = max(item[length_key] for item in batch)
    ids_list = []
    block_list = []
    lengths = []
    for item in batch:
        ids = item[ids_key]
        blocks = item[blocks_key]
        pad_len = max_len - item[length_key]
        ids_list.append(torch.cat([ids, torch.full((pad_len,), pad_id, dtype=torch.long)]))
        block_list.append(torch.cat([blocks, torch.zeros(pad_len, dtype=torch.long)]))
        lengths.append(item[length_key])
    return (
        torch.stack(ids_list),
        torch.stack(block_list),
        torch.tensor(lengths, dtype=torch.long),
    )


def collate_fn(
    batch,
    pad_id: int = 0,
    mask_id: int = 126336,
    gen_only: bool = False,
    skip_cls: bool = False,
):
    """Collate materialized generative sequence views."""
    del mask_id, gen_only, skip_cls

    gen_safe_ids, gen_safe_blocks, gen_safe_lengths = _pad_tensor_batch(
        batch,
        "gen_safe_input_ids",
        "gen_safe_block_ids",
        "gen_safe_real_length",
        pad_id,
    )
    gen_unsafe_ids, gen_unsafe_blocks, gen_unsafe_lengths = _pad_tensor_batch(
        batch,
        "gen_unsafe_input_ids",
        "gen_unsafe_block_ids",
        "gen_unsafe_real_length",
        pad_id,
    )
    hazard_ids, hazard_blocks, hazard_lengths = _pad_tensor_batch(
        batch,
        "hazard_input_ids",
        "hazard_block_ids",
        "hazard_real_length",
        pad_id,
    )

    return {
        "gen_safe_input_ids": gen_safe_ids,
        "gen_safe_block_ids": gen_safe_blocks,
        "gen_safe_real_lengths": gen_safe_lengths,
        "gen_unsafe_input_ids": gen_unsafe_ids,
        "gen_unsafe_block_ids": gen_unsafe_blocks,
        "gen_unsafe_real_lengths": gen_unsafe_lengths,
        "hazard_input_ids": hazard_ids,
        "hazard_block_ids": hazard_blocks,
        "hazard_real_lengths": hazard_lengths,
        "hazard_label_mask": torch.tensor(
            [item.get("hazard_label_mask", 0) for item in batch],
            dtype=torch.float32,
        ),
        "is_unsafe": torch.tensor([item["is_unsafe"] for item in batch], dtype=torch.float32),
        "conv_label_mask": torch.tensor(
            [item.get("conv_label_mask", 1) for item in batch],
            dtype=torch.float32,
        ),
        "has_response": torch.tensor([item["has_response"] for item in batch], dtype=torch.float32),
        "example_index": torch.tensor(
            [item.get("example_index", -1) for item in batch],
            dtype=torch.long,
        ),
        "prompt_is_unsafe": torch.tensor(
            [item.get("prompt_is_unsafe", 0) for item in batch],
            dtype=torch.float32,
        ),
        "prompt_label_mask": torch.tensor(
            [item.get("prompt_label_mask", 0) for item in batch],
            dtype=torch.float32,
        ),
        "response_is_unsafe": torch.tensor(
            [item.get("response_is_unsafe", 0) for item in batch],
            dtype=torch.float32,
        ),
        "response_label_mask": torch.tensor(
            [item.get("response_label_mask", 0) for item in batch],
            dtype=torch.float32,
        ),
    }


def build_eval_loader(
    examples: list[dict],
    tokenizer,
    max_seq_length: int,
    batch_size: int,
    *,
    gen_only: bool = True,
    prompt_format: str = "aegis_taxonomy",
) -> DataLoader:
    """Build the deterministic eval loader used by public inference."""
    dataset = SafetyDataset(
        examples,
        tokenizer,
        max_seq_length=max_seq_length,
        category_augmentation=False,
        prompt_format=prompt_format,
        gen_only=gen_only,
    )
    items = [dataset[idx] for idx in range(len(dataset))]
    items.sort(
        key=lambda it: (
            int(it.get("has_response", 0)),
            max(it["gen_safe_real_length"], it["gen_unsafe_real_length"]),
        )
    )
    pad_id = tokenizer.eos_token_id or 0
    return DataLoader(
        items,
        batch_size=batch_size,
        shuffle=False,
        collate_fn=partial(collate_fn, pad_id=pad_id, gen_only=gen_only),
        num_workers=0,
        pin_memory=torch.cuda.is_available(),
    )

	from __future__ import annotations

	from functools import partial
	from typing import Optional

	import torch
	from torch.utils.data import DataLoader, Dataset
	from transformers import PreTrainedTokenizer

	from .prompt import (
	build_prompt,
	build_prompt_parts,
	build_target,
	canonicalize_violated_categories,
	)

	PROMPT_BLOCK_ID = 1
	RESPONSE_BLOCK_ID = 2
	MASK_TOKEN_TEXT = "<\|mdm_mask\|>"


	class SafetyDataset(Dataset):
	"""Materialize the safe/unsafe prompt pairs used by masked-scope scoring."""

	def __init__(
	self,
	examples: list[dict],
	tokenizer: PreTrainedTokenizer,
	max_seq_length: int = 2048,
	category_augmentation: bool = False,
	category_augmentation_prob: float = 0.0,
	prompt_only_augmentation_prob: float = 0.0,
	prompt_format: str = "aegis_taxonomy",
	gen_only: bool = False,
	skip_cls: bool = False,
	):
	del category_augmentation, category_augmentation_prob, prompt_only_augmentation_prob
	del skip_cls
	self.examples = examples
	self.tokenizer = tokenizer
	self.max_seq_length = max_seq_length
	self.prompt_format = prompt_format
	self.gen_only = gen_only

	def __len__(self):
	return len(self.examples)

	def _chat_token_count(
	self, content: str, trusted_span: Optional[tuple[int, int]] = None
	) -> int:
	prompt_str = self.tokenizer.apply_chat_template(
	[{"role": "user", "content": content}],
	add_generation_prompt=True,
	tokenize=False,
	)
	ids, _ = self._encode_chat(prompt_str, content, trusted_span)
	return len(ids) + len(self.tokenizer("")["input_ids"])

	def _encode_chat(
	self,
	prompt_str: str,
	content: str,
	trusted_span: Optional[tuple[int, int]] = None,
	) -> tuple[list[int], list[tuple[int, int]]]:
	"""Tokenize a chat-formatted prompt without parsing special tokens in content.

	Special-token strings typed by a user or agent (e.g. ``<\|eot_id\|>``,
	``<\|mdm_mask\|>``) stay plain text; only the chat template's own control
	tokens are parsed. ``trusted_span`` (offsets within ``content``) marks
	release-generated text, such as the synthetic mask response, whose
	control tokens are kept.
	"""
	content_start = prompt_str.find(content)
	if content_start < 0:
	raise ValueError("Classification prompt not found inside chat template")
	content_end = content_start + len(content)
	segments = [(0, content_start, False)]
	if trusted_span is None:
	segments.append((content_start, content_end, True))
	else:
	trusted_start = content_start + trusted_span[0]
	trusted_end = content_start + trusted_span[1]
	segments += [
	(content_start, trusted_start, True),
	(trusted_start, trusted_end, False),
	(trusted_end, content_end, True),
	]
	segments.append((content_end, len(prompt_str), False))

	input_ids: list[int] = []
	offsets: list[tuple[int, int]] = []
	for start, end, split_special in segments:
	if start == end:
	continue
	encoding = self.tokenizer(
	prompt_str[start:end],
	add_special_tokens=False,
	return_offsets_mapping=True,
	split_special_tokens=split_special,
	)
	input_ids.extend(encoding["input_ids"])
	offsets.extend((s + start, e + start) for s, e in encoding["offset_mapping"])
	return input_ids, offsets

	def _gen_total_len(
	self,
	verdict_text: str,
	user_message: str,
	agent_response: Optional[str],
	categories: Optional[list[str]],
	violated_categories: Optional[list[str]] = None,
	synthetic_response: bool = False,
	) -> int:
	prefix, conversation, suffix = build_prompt_parts(
	user_message,
	agent_response,
	categories=categories,
	prompt_format=self.prompt_format,
	)
	trusted_span = None
	if synthetic_response and agent_response is not None:
	trusted_end = len(prefix) + len(conversation)
	trusted_span = (trusted_end - len(agent_response), trusted_end)
	prompt_len = self._chat_token_count(prefix + conversation + suffix, trusted_span)
	if verdict_text == "unsafe" and violated_categories:
	label_text = build_target(
	False,
	violated_categories,
	prompt_format=self.prompt_format,
	)
	else:
	label_text = verdict_text
	label_ids = self.tokenizer(label_text, add_special_tokens=False)["input_ids"]
	eos_extra = 1 if self.tokenizer.eos_token_id is not None else 0
	return prompt_len + len(label_ids) + eos_extra

	def _truncate_to_budget(
	self,
	user_message: str,
	agent_response: Optional[str],
	categories: Optional[list[str]],
	violated_categories: Optional[list[str]] = None,
	allow_agent_drop: bool = True,
	) -> tuple[str, Optional[str]]:
	"""Trim the response first, then the user text, until the pair fits."""
	while True:
	gen_safe_total = self._gen_total_len(
	"safe",
	user_message,
	agent_response,
	categories,
	synthetic_response=not allow_agent_drop,
	)
	gen_unsafe_total = self._gen_total_len(
	"unsafe",
	user_message,
	agent_response,
	categories,
	violated_categories=violated_categories,
	synthetic_response=not allow_agent_drop,
	)
	overflow = max(gen_safe_total, gen_unsafe_total) - self.max_seq_length
	if overflow <= 0:
	return user_message, agent_response

	if agent_response and allow_agent_drop:
	agent_ids = self.tokenizer(agent_response, add_special_tokens=False)["input_ids"]
	drop = min(max(overflow, 1), len(agent_ids))
	agent_ids = agent_ids[:-drop]
	agent_response = (
	self.tokenizer.decode(agent_ids, skip_special_tokens=True) if agent_ids else None
	)
	continue

	user_ids = self.tokenizer(user_message, add_special_tokens=False)["input_ids"]
	drop = min(max(overflow, 1), len(user_ids))
	user_ids = user_ids[:-drop]
	user_message = self.tokenizer.decode(user_ids, skip_special_tokens=True) if user_ids else ""

	def _build_gen_sequence(
	self,
	verdict_text: str,
	user_message: str,
	agent_response: Optional[str],
	categories: Optional[list[str]],
	violated_categories: Optional[list[str]] = None,
	synthetic_response: bool = False,
	) -> tuple[torch.Tensor, torch.Tensor]:
	prefix, conversation_text, suffix = build_prompt_parts(
	user_message,
	agent_response,
	categories=categories,
	prompt_format=self.prompt_format,
	)
	prompt_text = prefix + conversation_text + suffix
	prompt_str = self.tokenizer.apply_chat_template(
	[{"role": "user", "content": prompt_text}],
	add_generation_prompt=True,
	tokenize=False,
	)
	content_start = prompt_str.find(prompt_text)
	if content_start < 0:
	raise ValueError("Classification prompt not found inside chat template")
	conversation_start = content_start + len(prefix)

	user_span = f"User: {user_message}"
	user_start = conversation_start
	user_end = user_start + len(user_span)

	response_start = None
	response_end = None
	if agent_response is not None:
	response_span = f"\nAgent: {agent_response}"
	response_start = user_end
	response_end = response_start + len(response_span)

	trusted_span = None
	if synthetic_response and agent_response is not None:
	trusted_end = response_end - content_start
	trusted_span = (trusted_end - len(agent_response), trusted_end)
	prompt_ids, prompt_offsets = self._encode_chat(
	prompt_str, prompt_text, trusted_span
	)

	prompt_block_ids = []
	for start, end in prompt_offsets:
	in_user = end > user_start and start < user_end
	in_response = (
	response_start is not None
	and response_end is not None
	and end > response_start
	and start < response_end
	)
	if in_user:
	prompt_block_ids.append(PROMPT_BLOCK_ID)
	elif in_response:
	prompt_block_ids.append(RESPONSE_BLOCK_ID)
	else:
	prompt_block_ids.append(0)

	include_cats = verdict_text == "unsafe" and violated_categories
	if include_cats:
	label_text = build_target(
	False,
	violated_categories,
	prompt_format=self.prompt_format,
	)
	else:
	label_text = verdict_text

	label_encoding = self.tokenizer(
	label_text,
	add_special_tokens=False,
	return_offsets_mapping=True,
	)
	label_ids = label_encoding["input_ids"]
	label_offsets = label_encoding["offset_mapping"]

	if include_cats:
	char_blocks = [0] * len(label_text)
	pos = len("unsafe") + 1
	for cat_idx, cat in enumerate(violated_categories):
	for offset in range(len(cat)):
	if pos + offset < len(char_blocks):
	char_blocks[pos + offset] = 3 + cat_idx
	pos += len(cat) + 1
	label_block_ids = []
	for start, end in label_offsets:
	block_id = 0
	for char_idx in range(start, min(end, len(char_blocks))):
	block_id = max(block_id, char_blocks[char_idx])
	label_block_ids.append(block_id)
	else:
	label_block_ids = [0] * len(label_ids)

	input_ids = list(prompt_ids) + list(label_ids)
	block_ids = prompt_block_ids + label_block_ids

	eos_id = self.tokenizer.eos_token_id
	if eos_id is not None:
	input_ids.append(eos_id)
	block_ids.append(0)

	return (
	torch.tensor(input_ids, dtype=torch.long),
	torch.tensor(block_ids, dtype=torch.long),
	)

	def __getitem__(self, idx):
	ex = self.examples[idx]
	user_message = str(ex["user_message"])
	agent_response = ex.get("agent_response")
	is_safe = bool(ex.get("is_safe", True))
	violated = canonicalize_violated_categories(ex.get("violated_categories", []))

	had_response = agent_response is not None
	rendered_agent_response = str(agent_response) if had_response else MASK_TOKEN_TEXT
	user_message, rendered_agent_response = self._truncate_to_budget(
	user_message,
	rendered_agent_response,
	None,
	violated_categories=violated,
	allow_agent_drop=had_response,
	)
	has_response = int(had_response and rendered_agent_response is not None)
	prompt_text = build_prompt(
	user_message,
	rendered_agent_response,
	prompt_format=self.prompt_format,
	)

	gen_safe_input_ids, gen_safe_block_ids = self._build_gen_sequence(
	"safe",
	user_message,
	rendered_agent_response,
	None,
	synthetic_response=not had_response,
	)
	gen_unsafe_input_ids, gen_unsafe_block_ids = self._build_gen_sequence(
	"unsafe",
	user_message,
	rendered_agent_response,
	None,
	synthetic_response=not had_response,
	)
	hazard_label_mask = int((not is_safe) and bool(violated))
	if hazard_label_mask:
	hazard_input_ids, hazard_block_ids = self._build_gen_sequence(
	"unsafe",
	user_message,
	rendered_agent_response,
	None,
	violated_categories=violated,
	synthetic_response=not had_response,
	)
	else:
	hazard_input_ids = gen_unsafe_input_ids
	hazard_block_ids = gen_unsafe_block_ids

	prompt_label = (ex.get("prompt_label") or "").strip().lower()
	response_label = (ex.get("response_label") or "").strip().lower()
	prompt_label_known = prompt_label in {"safe", "unsafe"}
	response_label_known = had_response and response_label in {"safe", "unsafe"}
	if not prompt_label_known and not had_response:
	prompt_label_known = True
	if not response_label_known and had_response:
	response_label_known = True

	prompt_is_unsafe = (
	int(prompt_label == "unsafe")
	if prompt_label in {"safe", "unsafe"}
	else int((not is_safe) and not had_response)
	)
	response_is_unsafe = (
	int(response_label == "unsafe")
	if response_label in {"safe", "unsafe"}
	else int((not is_safe) and had_response)
	)

	return {
	"gen_safe_input_ids": gen_safe_input_ids,
	"gen_safe_block_ids": gen_safe_block_ids,
	"gen_safe_real_length": len(gen_safe_input_ids),
	"gen_unsafe_input_ids": gen_unsafe_input_ids,
	"gen_unsafe_block_ids": gen_unsafe_block_ids,
	"gen_unsafe_real_length": len(gen_unsafe_input_ids),
	"hazard_input_ids": hazard_input_ids,
	"hazard_block_ids": hazard_block_ids,
	"hazard_real_length": len(hazard_input_ids),
	"hazard_label_mask": hazard_label_mask,
	"hazard_eval_mask": int(is_safe or bool(violated)),
	"hazard_targets": list(violated),
	"prompt_text": prompt_text,
	"is_unsafe": int(not is_safe),
	"conv_label_mask": int(ex.get("conv_label_mask", 1)),
	"prompt_is_unsafe": prompt_is_unsafe,
	"prompt_label_mask": int(prompt_label_known),
	"response_is_unsafe": response_is_unsafe,
	"response_label_mask": int(response_label_known),
	"has_response": has_response,
	"example_index": int(ex.get("example_index", idx)),
	"eval_position": int(ex.get("eval_position", -1)),
	}


	def _pad_tensor_batch(
	batch: list[dict],
	ids_key: str,
	blocks_key: str,
	length_key: str,
	pad_id: int,
	) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
	max_len = max(item[length_key] for item in batch)
	ids_list = []
	block_list = []
	lengths = []
	for item in batch:
	ids = item[ids_key]
	blocks = item[blocks_key]
	pad_len = max_len - item[length_key]
	ids_list.append(torch.cat([ids, torch.full((pad_len,), pad_id, dtype=torch.long)]))
	block_list.append(torch.cat([blocks, torch.zeros(pad_len, dtype=torch.long)]))
	lengths.append(item[length_key])
	return (
	torch.stack(ids_list),
	torch.stack(block_list),
	torch.tensor(lengths, dtype=torch.long),
	)


	def collate_fn(
	batch,
	pad_id: int = 0,
	mask_id: int = 126336,
	gen_only: bool = False,
	skip_cls: bool = False,
	):
	"""Collate materialized generative sequence views."""
	del mask_id, gen_only, skip_cls

	gen_safe_ids, gen_safe_blocks, gen_safe_lengths = _pad_tensor_batch(
	batch,
	"gen_safe_input_ids",
	"gen_safe_block_ids",
	"gen_safe_real_length",
	pad_id,
	)
	gen_unsafe_ids, gen_unsafe_blocks, gen_unsafe_lengths = _pad_tensor_batch(
	batch,
	"gen_unsafe_input_ids",
	"gen_unsafe_block_ids",
	"gen_unsafe_real_length",
	pad_id,
	)
	hazard_ids, hazard_blocks, hazard_lengths = _pad_tensor_batch(
	batch,
	"hazard_input_ids",
	"hazard_block_ids",
	"hazard_real_length",
	pad_id,
	)

	return {
	"gen_safe_input_ids": gen_safe_ids,
	"gen_safe_block_ids": gen_safe_blocks,
	"gen_safe_real_lengths": gen_safe_lengths,
	"gen_unsafe_input_ids": gen_unsafe_ids,
	"gen_unsafe_block_ids": gen_unsafe_blocks,
	"gen_unsafe_real_lengths": gen_unsafe_lengths,
	"hazard_input_ids": hazard_ids,
	"hazard_block_ids": hazard_blocks,
	"hazard_real_lengths": hazard_lengths,
	"hazard_label_mask": torch.tensor(
	[item.get("hazard_label_mask", 0) for item in batch],
	dtype=torch.float32,
	),
	"is_unsafe": torch.tensor([item["is_unsafe"] for item in batch], dtype=torch.float32),
	"conv_label_mask": torch.tensor(
	[item.get("conv_label_mask", 1) for item in batch],
	dtype=torch.float32,
	),
	"has_response": torch.tensor([item["has_response"] for item in batch], dtype=torch.float32),
	"example_index": torch.tensor(
	[item.get("example_index", -1) for item in batch],
	dtype=torch.long,
	),
	"prompt_is_unsafe": torch.tensor(
	[item.get("prompt_is_unsafe", 0) for item in batch],
	dtype=torch.float32,
	),
	"prompt_label_mask": torch.tensor(
	[item.get("prompt_label_mask", 0) for item in batch],
	dtype=torch.float32,
	),
	"response_is_unsafe": torch.tensor(
	[item.get("response_is_unsafe", 0) for item in batch],
	dtype=torch.float32,
	),
	"response_label_mask": torch.tensor(
	[item.get("response_label_mask", 0) for item in batch],
	dtype=torch.float32,
	),
	}


	def build_eval_loader(
	examples: list[dict],
	tokenizer,
	max_seq_length: int,
	batch_size: int,
	*,
	gen_only: bool = True,
	prompt_format: str = "aegis_taxonomy",
	) -> DataLoader:
	"""Build the deterministic eval loader used by public inference."""
	dataset = SafetyDataset(
	examples,
	tokenizer,
	max_seq_length=max_seq_length,
	category_augmentation=False,
	prompt_format=prompt_format,
	gen_only=gen_only,
	)
	items = [dataset[idx] for idx in range(len(dataset))]
	items.sort(
	key=lambda it: (
	int(it.get("has_response", 0)),
	max(it["gen_safe_real_length"], it["gen_unsafe_real_length"]),
	)
	)
	pad_id = tokenizer.eos_token_id or 0
	return DataLoader(
	items,
	batch_size=batch_size,
	shuffle=False,
	collate_fn=partial(collate_fn, pad_id=pad_id, gen_only=gen_only),
	num_workers=0,
	pin_memory=torch.cuda.is_available(),
	)