1717import logging
1818import re
1919import threading
20+ import unicodedata
21+ from collections .abc import Iterator
2022from dataclasses import dataclass
2123from typing import Any
2224
2325logger = logging .getLogger (__name__ )
2426
25- # The one tokenizer used for both candidate selection and placeholder ordering. Requiring
26- # word boundaries keeps candidates to runs of letters that can actually be matched again
27- # when masking, so a token embedded in an alphanumeric blob (the "h" of "h4ck3r") is never
28- # selected and then silently dropped.
29- _WORD_PATTERN = re .compile (r"\b[A-Za-z]+\b" )
27+ # The one tokenizer used for candidate selection, placeholder ordering and replacement
28+ # (always through ``_iter_words``). Requiring word boundaries keeps candidates to whole
29+ # words, so a token embedded in an alphanumeric blob (the "h" of "h4ck3r") is never
30+ # selected. ``[^\W\d_]`` admits Unicode letters, so words such as "Schädlinge" or Cyrillic
31+ # words are candidates too; ``_iter_words`` then drops matches that are not all letters.
32+ _WORD_PATTERN = re .compile (r"\b[^\W\d_]+\b" )
33+
34+
35+ def _is_word_continuation (char : str ) -> bool :
36+ return char in ("\u200c " , "\u200d " ) or unicodedata .category (char ).startswith ("M" )
37+
38+
39+ def _word_key (word : str ) -> str :
40+ """
41+ Return the case-insensitive key that identifies a word.
42+
43+ Selection, deduplication, ordering and replacement all compare words through this key,
44+ so two spellings are either the same masked word everywhere or different words
45+ everywhere. For ASCII it equals ``str.lower``.
46+
47+ Args:
48+ word (str): A word from the prompt or from the caller.
49+
50+ Returns:
51+ str: The case-folded word.
52+ """
53+ return word .casefold ()
54+
55+
56+ def _iter_words (prompt : str ) -> Iterator [re .Match [str ]]:
57+ """
58+ Yield the maskable words of ``prompt``.
59+
60+ ``[^\\ W\\ d_]`` still admits characters that are numeric without being decimal digits,
61+ such as "²" or "¼", so a match is kept only if every character is a letter. "abc²" is
62+ therefore rejected whole, like "h4ck3r", rather than becoming a candidate.
63+
64+ ``\\ w`` does not cover combining marks (Devanagari vowel signs, Hebrew and Arabic vowel
65+ points, or the accents of NFD-decomposed text) or zero-width joiners and non-joiners.
66+ These characters can continue a word even though ``\\ b`` treats them as boundaries.
67+ A match that touches one of them is only a fragment, so it is skipped as well.
68+
69+ Args:
70+ prompt (str): The prompt to scan.
71+
72+ Yields:
73+ re.Match[str]: Each whole-word match, in order of appearance.
74+ """
75+ for match in _WORD_PATTERN .finditer (prompt ):
76+ if not match .group ().isalpha ():
77+ continue
78+ start , end = match .span ()
79+ if start > 0 and _is_word_continuation (prompt [start - 1 ]):
80+ continue
81+ if end < len (prompt ) and _is_word_continuation (prompt [end ]):
82+ continue
83+ yield match
84+
3085
3186# Part-of-speech label used when spaCy is unavailable or the word was caller-supplied.
3287_GENERIC_POS = "word"
@@ -170,15 +225,15 @@ def _get_nlp() -> Any:
170225
171226def _pos_lookup (prompt : str ) -> dict [str , str ]:
172227 """
173- Map each alphabetic noun/verb (lowercased) in the prompt to its clue label.
228+ Map each alphabetic noun/verb in the prompt (by ``_word_key``) to its clue label.
174229
175230 Returns an empty mapping when spaCy is unavailable.
176231
177232 Args:
178233 prompt (str): The prompt to tag.
179234
180235 Returns:
181- dict[str, str]: Lowercased word to part-of-speech label ("noun" or "verb").
236+ dict[str, str]: Word key to part-of-speech label ("noun" or "verb").
182237 """
183238 nlp = _get_nlp ()
184239 if nlp is None :
@@ -187,7 +242,7 @@ def _pos_lookup(prompt: str) -> dict[str, str]:
187242 for token in nlp (prompt ):
188243 if token .is_alpha and token .pos_ in _MASKABLE_POS :
189244 # Keep the first tag seen for a given surface form.
190- lookup .setdefault (token .text . lower ( ), _MASKABLE_POS [token .pos_ ])
245+ lookup .setdefault (_word_key ( token .text ), _MASKABLE_POS [token .pos_ ])
191246 return lookup
192247
193248
@@ -214,20 +269,20 @@ def _rank_candidates(
214269 Returns:
215270 list[str]: Distinct candidate words (as they appear in the prompt) in priority order.
216271 """
217- tokens = _WORD_PATTERN . findall ( prompt )
272+ tokens = [ match . group () for match in _iter_words ( prompt )]
218273 seen : set [str ] = set ()
219274 unique : list [str ] = []
220275 for token in tokens :
221- key = token . lower ( )
276+ key = _word_key ( token )
222277 if key not in seen :
223278 seen .add (key )
224279 unique .append (token )
225280
226- essential_lower = {w . lower ( ) for w in (essential_words or [])}
281+ essential_keys = {_word_key ( w ) for w in (essential_words or [])}
227282
228283 def tier (word : str ) -> int :
229- key = word . lower ( )
230- if key in essential_lower :
284+ key = _word_key ( word )
285+ if key in essential_keys :
231286 return 0
232287 if key in _ESSENTIAL_HARM_WORDS :
233288 return 1
@@ -237,7 +292,7 @@ def tier(word: str) -> int:
237292 return 3
238293 return 4
239294
240- return sorted (unique , key = lambda w : (tier (w ), - len (w ), w . lower ( )))
295+ return sorted (unique , key = lambda w : (tier (w ), - len (w ), _word_key ( w )))
241296
242297
243298def mask_prompt (
@@ -278,25 +333,34 @@ def mask_prompt(
278333 return MaskResult (masked_prompt = prompt , masked_words = [])
279334
280335 # Number placeholders by where each chosen word first appears (case-insensitively), so
281- # [WORD1] is the leftmost masked word. Candidates come from the same word-boundary scan,
282- # so every chosen word has an entry here.
336+ # [WORD1] is the leftmost masked word. Candidates come from the same scan, so every
337+ # chosen word has an entry here.
283338 first_index : dict [str , int ] = {}
284- for match in _WORD_PATTERN . finditer (prompt ):
285- first_index .setdefault (match .group (). lower ( ), match .start ())
286- ordered = sorted (chosen , key = lambda w : first_index [w . lower ( )])
339+ for match in _iter_words (prompt ):
340+ first_index .setdefault (_word_key ( match .group ()), match .start ())
341+ ordered = sorted (chosen , key = lambda w : first_index [_word_key ( w )])
287342
288343 masked_words : list [MaskedWord ] = []
289344 placeholders : dict [str , str ] = {}
290345 for index , word in enumerate (ordered , start = 1 ):
291346 placeholder = f"[WORD{ index } ]"
292- pos = pos_lookup .get (word . lower ( ), _GENERIC_POS )
347+ pos = pos_lookup .get (_word_key ( word ), _GENERIC_POS )
293348 masked_words .append (MaskedWord (text = word , placeholder = placeholder , pos = pos ))
294- placeholders [word .lower ()] = placeholder
295-
296- # Replace every occurrence of every chosen word in one case-insensitive pass, so a word
297- # that recurs or appears in different casing is never left in cleartext, and a placeholder
298- # already inserted cannot be re-matched while masking the next word.
299- pattern = re .compile (r"\b(" + "|" .join (re .escape (w ) for w in ordered ) + r")\b" , re .IGNORECASE )
300- masked_prompt = pattern .sub (lambda m : placeholders [m .group (0 ).lower ()], prompt )
301-
302- return MaskResult (masked_prompt = masked_prompt , masked_words = masked_words )
349+ placeholders [_word_key (word )] = placeholder
350+
351+ # Replace the word spans of the same scan, looked up by the same key. Every occurrence of
352+ # a chosen word is masked whatever its casing, a fragment the scan skipped (next to a
353+ # combining mark or joiner) is never masked, and every replaced span has a placeholder.
354+ # Building the result in one pass means an inserted placeholder can never be matched again.
355+ pieces : list [str ] = []
356+ last_end = 0
357+ for match in _iter_words (prompt ):
358+ placeholder = placeholders .get (_word_key (match .group ()))
359+ if placeholder is None :
360+ continue
361+ pieces .append (prompt [last_end : match .start ()])
362+ pieces .append (placeholder )
363+ last_end = match .end ()
364+ pieces .append (prompt [last_end :])
365+
366+ return MaskResult (masked_prompt = "" .join (pieces ), masked_words = masked_words )
0 commit comments