164 lines
5.6 KiB
Python
164 lines
5.6 KiB
Python
# 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,
|
||
*(getattr(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
|