tokenization_ravenguard.py

"""RavenGuard tokenizer (trust_remote_code).

Thin ``PreTrainedTokenizer`` wrapper around the Netis Amniota rustbpe/tiktoken encoding
pickled in ``tokenizer.pkl``. ``encode``/``decode`` delegate straight to the
tiktoken ``Encoding`` so text round-trips byte-for-byte, matching how the guard
was trained. Vocab size is the tiktoken ``n_vocab`` (65536).
"""
from __future__ import annotations

import os
import pickle
import shutil

from transformers import PreTrainedTokenizer

SPECIAL_TOKENS = [
    "<|bos|>",
    "<|user_start|>", "<|user_end|>",
    "<|assistant_start|>", "<|assistant_end|>",
    "<|python_start|>", "<|python_end|>",
    "<|output_start|>", "<|output_end|>",
]

# Bundled canonical policy prompt (the exact v2 system prompt the guard was trained
# and evaluated with). Used as the default system message by apply_chat_template.
POLICY_PROMPT_FILE = "policy_prompt_v2.txt"

_ROLE_LABEL = {"user": "User", "assistant": "Assistant",
               "tool": "Tool", "tool_result": "Tool"}


class RavenGuardTokenizer(PreTrainedTokenizer):
    vocab_files_names = {"vocab_file": "tokenizer.pkl"}
    model_input_names = ["input_ids", "attention_mask"]

    def __init__(self, vocab_file=None, **kwargs):
        if vocab_file is None or not os.path.isfile(vocab_file):
            raise ValueError(f"tokenizer.pkl not found (vocab_file={vocab_file!r})")
        with open(vocab_file, "rb") as f:
            self.enc = pickle.load(f)
        self._vocab_file = vocab_file
        super().__init__(**kwargs)
        # Special-token ids used to wrap a chat turn (resolved once).
        self._sp = {t: self.enc.encode_single_token(t) for t in (
            "<|bos|>", "<|user_start|>", "<|user_end|>", "<|assistant_start|>",
            "<|assistant_end|>")}
        # eos = assistant_end, so .generate() stops at the end of the label block.
        self._assistant_end_id = self._sp["<|assistant_end|>"]
        # Default system prompt: the bundled canonical policy (single source of
        # truth, shipped next to this file); None if the asset is absent.
        self._default_policy = None
        pol = os.path.join(os.path.dirname(os.path.abspath(vocab_file)), POLICY_PROMPT_FILE)
        if os.path.isfile(pol):
            with open(pol, encoding="utf-8") as f:
                self._default_policy = f.read()

    # ---- core sizes ---------------------------------------------------------
    @property
    def vocab_size(self):
        return self.enc.n_vocab

    def get_vocab(self):
        # tiktoken is byte-level with integer ids; expose an id<->str(id) view so
        # the base-class machinery stays consistent with _tokenize below.
        return {str(i): i for i in range(self.vocab_size)}

    # ---- fast path: delegate straight to tiktoken ---------------------------
    def encode(self, text, add_special_tokens=False, **kwargs):
        if not isinstance(text, str):
            return super().encode(text, add_special_tokens=add_special_tokens, **kwargs)
        return self.enc.encode_ordinary(text)

    def decode(self, token_ids, skip_special_tokens=False, **kwargs):
        if hasattr(token_ids, "tolist"):
            token_ids = token_ids.tolist()
        if isinstance(token_ids, int):
            token_ids = [token_ids]
        return self.enc.decode(list(token_ids))

    def encode_special(self, text):
        return self.enc.encode_single_token(text)

    # ---- chat interface -----------------------------------------------------
    def apply_chat_template(self, conversation, *, tokenize=True,
                            add_generation_prompt=True, return_tensors=None, **kwargs):
        """Build the guard's prompt from chat messages, injection-safely.

        Renders ``<system policy>\\n\\n`` + role-tagged turns (``User: ... /
        Assistant: ...``), then wraps it as
        ``[<|bos|> <|user_start|>] + encode_ordinary(text) + [<|user_end|>]``
        (+ ``<|assistant_start|>`` when ``add_generation_prompt``), byte-identical
        to how the guard was trained. Message CONTENT is encoded with
        ``encode_ordinary`` (special tokens disallowed), so a message can never
        smuggle a control token such as ``<|assistant_end|>`` into the stream — a
        string Jinja template tokenized uniformly could not guarantee that.

        If no ``system`` message is supplied, the bundled canonical policy prompt
        is used. Batched input (list of conversations) is rejected when
        ``tokenize=True`` because the model has no padding-mask path — tokenize one
        conversation at a time.
        """
        if conversation and isinstance(conversation[0], (list, tuple)):
            if tokenize:
                raise ValueError(
                    "RavenGuard tokenizes one conversation at a time (no padding-mask "
                    "path for batches); call apply_chat_template per conversation.")
            return [self.apply_chat_template(c, tokenize=False,
                    add_generation_prompt=add_generation_prompt, **kwargs)
                    for c in conversation]

        system = self._default_policy or ""
        body = []
        for m in conversation:
            role, content = m.get("role"), (m.get("content") or "")
            if role == "system":
                system = content
            else:
                body.append(f"{_ROLE_LABEL.get(role, str(role).capitalize())}: {content}")
        full = (system + "\n\n" + "\n".join(body)) if system else "\n".join(body)
        if not tokenize:
            return full

        ids = ([self._sp["<|bos|>"], self._sp["<|user_start|>"]]
               + self.enc.encode_ordinary(full)
               + [self._sp["<|user_end|>"]])
        if add_generation_prompt:
            ids.append(self._sp["<|assistant_start|>"])
        if return_tensors == "pt":
            import torch
            return torch.tensor([ids], dtype=torch.long)
        return ids

    # ---- base-class hooks (kept consistent with get_vocab) -------------------
    def _tokenize(self, text, **kwargs):
        return [str(i) for i in self.enc.encode_ordinary(text)]

    def _convert_token_to_id(self, token):
        try:
            return int(token)
        except (TypeError, ValueError):
            return self.enc.encode_single_token(token)

    def _convert_id_to_token(self, index):
        return str(index)

    def convert_tokens_to_string(self, tokens):
        ids = [int(t) for t in tokens]
        return self.enc.decode(ids)

    def save_vocabulary(self, save_directory, filename_prefix=None):
        os.makedirs(save_directory, exist_ok=True)
        name = (filename_prefix + "-" if filename_prefix else "") + "tokenizer.pkl"
        out = os.path.join(save_directory, name)
        if os.path.abspath(out) != os.path.abspath(self._vocab_file):
            shutil.copyfile(self._vocab_file, out)
        return (out,)