fix

Signed-off-by: guangli.bao <[email protected]>

Author: tukwila

src/guidellm/utils/preprocessing_sharegpt_data.py

@@ -7,7 +7,7 @@
 
 import numpy as np
 from datasets import load_dataset
-from transformers import AutoTokenizer
+from transformers import AutoTokenizer, PreTrainedTokenizerBase
 
 MIN_CHAR = 10
 MAX_CHAR = 1000
@@ -16,7 +16,7 @@
 class TokenCounter:
     def __init__(self, model_name: str = "mistralai/Mistral-7B-Instruct-v0.2"):
         self.model_name = model_name
-        self._tokenizer: Optional[AutoTokenizer] = None
+        self._tokenizer: Optional[PreTrainedTokenizerBase] = None
 
     def _initialize_tokenizer(self) -> None:
         if self._tokenizer is None:
@@ -33,7 +33,7 @@ def estimate_num_tokens(self, text: str) -> int:
             return 0
 
         try:
-            encoding = self._tokenizer(text, return_tensors=None)
+            encoding = self._tokenizer.__call__(text, return_tensors=None)
             return len(encoding["input_ids"])
         except (AttributeError, TypeError, RuntimeError) as e:
             raise ValueError(f"Error processing text: {e}") from e