Support any LLMLingua-2-style token classifier (detect subword style from the tokenizer)

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[:])

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()