"""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 ``\\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,)