From 1f2b8817dfecb4b7c17c217e2b517d0ce9f06c36 Mon Sep 17 00:00:00 2001 From: riceharvest Date: Thu, 8 Oct 2026 16:00:57 +0200 Subject: [PATCH] Detect subword style from the tokenizer so any LLMLingua-2-style classifier works is_begin_of_new_word/get_pure_token matched model-name substrings and raised NotImplementedError for any other checkpoint. Detect WordPiece (##), SentencePiece (U+2581) or byte-level (U+0120) marking once from the tokenizer and use it when no name rule matches. Existing name-based behaviour is unchanged. Co-Authored-By: Claude Opus 5.5 --- llmlingua/prompt_compressor.py | 22 ++++++++++++++------- llmlingua/utils.py | 35 ++++++++++++++++++++++++++++++++-- 2 files changed, 48 insertions(+), 9 deletions(-) diff --git a/llmlingua/prompt_compressor.py b/llmlingua/prompt_compressor.py index 84e390e..7e4021b 100644 --- a/llmlingua/prompt_compressor.py +++ b/llmlingua/prompt_compressor.py @@ -25,6 +25,7 @@ from .utils import ( TokenClfDataset, get_pure_token, + detect_subword_style, is_begin_of_new_word, process_structured_json_data, remove_consecutive_commas, @@ -89,6 +90,7 @@ def __init__( self.oai_tokenizer = tiktoken.encoding_for_model("gpt-3.5-turbo") self.load_model(model_name, device_map, model_config) + self.subword_style = detect_subword_style(self.tokenizer) if use_llmlingua2 or use_slingua: # slingua use llmlingua2 backend self.init_llmlingua2(**llmlingua2_config) @@ -2207,9 +2209,11 @@ def __get_context_prob( valid_token_probs_no_force, convert_mode=token_to_word ) - if "xlm-roberta-large" in self.model_name: + prefix = "▁" if "xlm-roberta-large" in self.model_name else { + "sentencepiece": "▁", "bytelevel": "Ġ"}.get(self.subword_style) + if prefix: for i in range(len(words)): - words[i] = words[i].lstrip("▁") + words[i] = words[i].lstrip(prefix) chunk_words.append(words) chunk_probs.append(word_probs_no_force) @@ -2263,8 +2267,10 @@ def __merge_token_to_word( if token in self.special_tokens: continue # add a new word - elif is_begin_of_new_word(token, self.model_name, force_tokens, token_map): - pure_token = get_pure_token(token, self.model_name) + elif is_begin_of_new_word( + token, self.model_name, force_tokens, token_map, self.subword_style + ): + pure_token = get_pure_token(token, self.model_name, self.subword_style) prob_no_force = prob if pure_token in force_tokens or pure_token in set(token_map.values()): prob = 1.0 @@ -2280,7 +2286,7 @@ def __merge_token_to_word( word_probs_no_force.append([prob_no_force]) # concatenate with previous token else: - pure_token = get_pure_token(token, self.model_name) + pure_token = get_pure_token(token, self.model_name, self.subword_style) words[-1] += pure_token word_probs[-1].append( 1.0 @@ -2428,9 +2434,11 @@ def split_string_to_words(input_string): else: word_labels.append(0) keep_str = self.tokenizer.convert_tokens_to_string(keep_words) - if "xlm-roberta-large" in self.model_name: + prefix = "▁" if "xlm-roberta-large" in self.model_name else { + "sentencepiece": "▁", "bytelevel": "Ġ"}.get(self.subword_style) + if prefix: for i in range(len(words)): - words[i] = words[i].lstrip("▁") + words[i] = words[i].lstrip(prefix) compressed_chunk_list.append(keep_str) word_list.append(words[:]) diff --git a/llmlingua/utils.py b/llmlingua/utils.py index da9986e..040a8fb 100644 --- a/llmlingua/utils.py +++ b/llmlingua/utils.py @@ -78,7 +78,24 @@ def seed_everything(seed: int): torch.backends.cudnn.benchmark = False -def is_begin_of_new_word(token, model_name, force_tokens, token_map): +def detect_subword_style(tokenizer): + """How the tokenizer marks word boundaries: "wordpiece" (continuations start with ##), + "sentencepiece" (word starts carry ▁) or "bytelevel" (word starts carry Ġ). + Lets any LLMLingua-2-style token classifier work without a model-name special case.""" + tokens = tokenizer.tokenize("hello unbelievably wonderful") + if any(t.startswith("##") for t in tokens): + return "wordpiece" + if any(t.startswith("▁") for t in tokens): + return "sentencepiece" + if any(t.startswith("Ġ") for t in tokens): + return "bytelevel" + return None + + +_STYLE_PREFIX = {"sentencepiece": "▁", "bytelevel": "Ġ"} + + +def is_begin_of_new_word(token, model_name, force_tokens, token_map, style=None): if "bert-base-multilingual-cased" in model_name \ or "tinybert" in model_name.lower() \ or "mobilebert" in model_name.lower(): @@ -97,6 +114,16 @@ def is_begin_of_new_word(token, model_name, force_tokens, token_map): ): return True return token.startswith("▁") + elif style == "wordpiece": + if token.lstrip("##") in force_tokens or token.lstrip("##") in set(token_map.values()): + return True + return not token.startswith("##") + elif style in _STYLE_PREFIX: + # A word starts only where the tokenizer saw whitespace, so identifiers such as + # "ord_8f3a2c91" or "10.0.4.17:5432" stay one word instead of being split at "_" / ".". + if token in force_tokens or token in set(token_map.values()): + return True + return token.startswith(_STYLE_PREFIX[style]) else: raise NotImplementedError() @@ -107,7 +134,7 @@ def replace_added_token(token, token_map): return token -def get_pure_token(token, model_name): +def get_pure_token(token, model_name, style=None): if "bert-base-multilingual-cased" in model_name \ or "tinybert" in model_name.lower() \ or "mobilebert" in model_name.lower(): @@ -116,6 +143,10 @@ def get_pure_token(token, model_name): or 'slingua' in model_name.lower() \ or 'securitylingua' in model_name.lower(): return token.lstrip("▁") + elif style == "wordpiece": + return token.lstrip("##") + elif style in _STYLE_PREFIX: + return token.lstrip(_STYLE_PREFIX[style]) else: raise NotImplementedError()