Fix(LLMLingua-2): fix wrong special tokens being used (microsoft#181)

cornzz · Sep 13, 2024 · c8709e6 · c8709e6
1 parent 2dbdbd3
commit c8709e6
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/llmlingua/prompt_compressor.py b/llmlingua/prompt_compressor.py
@@ -2161,7 +2161,7 @@ def __get_context_prob(
                 chunk_list.append(c)
 
         dataset = TokenClfDataset(
-            chunk_list, tokenizer=self.tokenizer, max_len=self.max_seq_len
+            chunk_list, tokenizer=self.tokenizer, max_len=self.max_seq_len, model_name=self.model_name
         )
         dataloader = DataLoader(
             dataset, batch_size=self.max_batch_size, shuffle=False, drop_last=False
@@ -2339,7 +2339,7 @@ def split_string_to_words(input_string):
                 chunk_list.append(c)
 
         dataset = TokenClfDataset(
-            chunk_list, tokenizer=self.tokenizer, max_len=self.max_seq_len
+            chunk_list, tokenizer=self.tokenizer, max_len=self.max_seq_len, model_name=self.model_name
         )
         dataloader = DataLoader(
             dataset, batch_size=self.max_batch_size, shuffle=False, drop_last=False