initial upload

utils/whitelist.py

@@ -0,0 +1,163 @@
+# utils/whitelist.py
+"""
+Constructs a global whitelist of strings that must never be placed in
+any ban list.
+
+Sources
+-------
+1. All special-token texts exposed by the model’s tokenizer.
+2. Every phrase (entire line) – and every word inside those phrases –
+   that appears *after* the assistant-message placeholder in a single
+   user→assistant chat-template example.
+3. Optional user-supplied strings from the YAML / CLI configuration.
+"""
+
+from __future__ import annotations
+import json
+import threading
+from pathlib import Path
+from typing import Iterable, Set
+
+from transformers import AutoTokenizer
+from slop_forensics.utils import normalize_text as normalise_keep_marks
+from slop_forensics.utils import extract_words
+
+# ------------------------------------------------------------------------------
+# Helper class
+# ------------------------------------------------------------------------------
+
+class WhitelistBuilder:
+    """
+    Static helper for creating and persisting the whitelist.
+
+    All strings are:
+
+    * converted to lowercase
+    * normalised via `normalise_keep_marks`
+    * deduplicated
+    """
+
+    _tokenizer_cache: dict[str, "AutoTokenizer"] = {}
+    _cache_lock = threading.Lock()
+
+    # ------------------------------------------------------------------ #
+    # Public API                                                         #
+    # ------------------------------------------------------------------ #
+
+    @classmethod
+    def build(
+        cls,
+        model_id: str,
+        *,
+        extra_user_items: Iterable[str] | None = None,
+    ) -> Set[str]:
+        """
+        Parameters
+        ----------
+        model_id
+            Hugging Face model ID or local checkpoint directory.
+        extra_user_items
+            Arbitrary strings provided by the user that must also be whitelisted.
+
+        Returns
+        -------
+        Set[str]
+            Normalised whitelist entries (lower-cased, no duplicates, no blanks).
+        """
+        tokenizer = cls._get_tokenizer(model_id)
+        whitelist: set[str] = set()
+
+        # 1. Special token texts -----------------------------------------
+        special_token_texts = [
+            tokenizer.bos_token,
+            tokenizer.eos_token,
+            tokenizer.unk_token,
+            tokenizer.pad_token,
+            tokenizer.cls_token,
+            tokenizer.sep_token,
+            *(tokenizer.additional_special_tokens or []),
+        ]
+        for raw_text in special_token_texts:
+            if not raw_text:
+                continue
+            cls._add_phrase_and_words(whitelist, raw_text)
+
+        # 2. Tail of the chat template -----------------------------------
+        template_tail_text = cls._get_chat_template_tail(tokenizer)
+        for line in template_tail_text.splitlines():
+            cls._add_phrase_and_words(whitelist, line)
+
+        # 3. User-supplied extras ----------------------------------------
+        if extra_user_items:
+            for item in extra_user_items:
+                cls._add_phrase_and_words(whitelist, str(item))
+
+        # Final clean-up: remove any empty strings that might have slipped in
+        whitelist.discard("")
+        return whitelist
+
+    @classmethod
+    def write(cls, file_path: Path, whitelist: Iterable[str]) -> None:
+        """Write the whitelist to *file_path* as pretty-printed JSON."""
+        file_path.parent.mkdir(parents=True, exist_ok=True)
+        file_path.write_text(
+            json.dumps(sorted(whitelist), indent=2, ensure_ascii=False),
+            encoding="utf-8",
+        )
+
+    # ------------------------------------------------------------------ #
+    # Internal helpers                                                   #
+    # ------------------------------------------------------------------ #
+
+    @classmethod
+    def _get_tokenizer(cls, model_id: str):
+        """Thread-safe one-time load of the tokenizer."""
+        with cls._cache_lock:
+            tokenizer = cls._tokenizer_cache.get(model_id)
+            if tokenizer is None:
+                tokenizer = AutoTokenizer.from_pretrained(
+                    model_id, trust_remote_code=True
+                )
+                cls._tokenizer_cache[model_id] = tokenizer
+            return tokenizer
+
+    @staticmethod
+    def _add_phrase_and_words(target_set: set[str], raw_text: str) -> None:
+        """
+        Normalise *raw_text*, add the whole phrase, then add each individual
+        word extracted from the phrase.
+        """
+        normalised = normalise_keep_marks(raw_text)
+        if not normalised:
+            return
+        target_set.add(normalised)
+        target_set.update(extract_words(normalised))
+
+    @staticmethod
+    def _get_chat_template_tail(tokenizer) -> str:
+        """
+        Build one user→assistant chat-template instance and return only the
+        text *after* the assistant placeholder.  That is the scaffold the
+        model tends to emit, so its words must be whitelisted.
+        """
+        placeholder_user = "__USER__"
+        placeholder_assistant = "__ASSISTANT__"
+
+        messages = [
+            {"role": "user",      "content": placeholder_user},
+            {"role": "assistant", "content": placeholder_assistant},
+        ]
+
+        full_template: str = tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=False,
+        )
+
+        assistant_pos = full_template.find(placeholder_assistant)
+        if assistant_pos == -1:
+            # Fallback: return the whole template if the placeholder wasn't found
+            return full_template.strip()
+
+        tail = full_template[assistant_pos + len(placeholder_assistant):].strip()
+        return tail