Update gptx_tokenizer.py

gptx_tokenizer.py

@@ -38,7 +38,6 @@ class HFGPTXTokenizer(PreTrainedTokenizer):
             text (str): The text to encode.
             return_tokens (bool): If True, returns token strings instead of token IDs.
             is_continuation (bool): If True, uses a continuation tokenizer (if available).
-
         Returns:
             List[int] or List[str]: Encoded text as a list of token IDs or token strings.
         """
@@ -56,7 +55,6 @@ class HFGPTXTokenizer(PreTrainedTokenizer):
         """
         Create a list of special tokens, including the BOS, EOS, PAD, EOD tokens,
         and 256 additional placeholder tokens.
-
         Returns:
             List[str]: List of special tokens.
         """
@@ -64,7 +62,7 @@ class HFGPTXTokenizer(PreTrainedTokenizer):
             f"<placeholder_tok_{i}>" for i in range(256)
         ]
     
-    def find_tokenizer_config(self, config_path: Path, repo_id: str = None) -> Optional[Path]:
+def find_tokenizer_config(self, config_path: Path, repo_id: str = None) -> Optional[Path]:
         if not os.path.isfile(config_path):
             config_path = try_to_load_from_cache(repo_id=repo_id, filename=Path(config_path).name)
             if not config_path:
@@ -89,33 +87,60 @@ class HFGPTXTokenizer(PreTrainedTokenizer):
             OSError: If the model file cannot be loaded or downloaded.
         """
         if not os.path.isfile(model_file_or_name):
-            if repo_id is None:
-                raise ValueError("repo_id must be provided if model_file_or_name is not a local file")
-            
-            try:
-                # List all files in the repo
-                repo_files = list_repo_files(repo_id)
-                
-                # Find the tokenizer model file
-                tokenizer_files = [f for f in repo_files if f.endswith('.model')]
-                if not tokenizer_files:
-                    raise FileNotFoundError(f"No .model file found in repository {repo_id}")
-                
-                # Use the first .model file found
-                model_file = tokenizer_files[0]
-                print(f"Found tokenizer model file: {model_file}")
-                
-                # Download the file
-                model_file_or_name = hf_hub_download(repo_id=repo_id, filename=model_file)
-                print(f"Downloaded tokenizer model to: {model_file_or_name}")
-            except Exception as e:
-                raise OSError(f"Failed to download tokenizer model: {str(e)}")
+            model_file_or_name = try_to_load_from_cache(repo_id=repo_id, filename=Path(model_file_or_name).name)
+            if not model_file_or_name:
+                model_file_or_name = self._download_model_from_hub(repo_id=repo_id)
 
         try:
             return spm.SentencePieceProcessor(model_file=model_file_or_name)
         except Exception as e:
             raise OSError(f"Failed to load tokenizer model: {str(e)}")
-    
+
+    def _download_model_from_hub(self, repo_id: str) -> Optional[str]:
+        try:
+            # List all files in the repo
+            repo_files = list_repo_files(repo_id)
+
+            # Find the tokenizer model file
+            tokenizer_files = [f for f in repo_files if f.endswith('.model')]
+            if not tokenizer_files:
+                raise FileNotFoundError(f"No .model file found in repository {repo_id}")
+
+            # Use the first .model file found
+            model_file = tokenizer_files[0]
+            print(f"Found tokenizer model file: {model_file}")
+
+            # Download the file
+            model_file_or_name = hf_hub_download(repo_id=repo_id, filename=model_file)
+            print(f"Downloaded tokenizer model to: {model_file_or_name}")
+        except Exception as e:
+            raise OSError(f"Failed to download tokenizer model: {str(e)}")
+
+        return model_file_or_name
+
+    def _download_config_from_hub(self, repo_id: str):
+        if repo_id is None:
+            raise ValueError("repo_id must be provided if config_path is not a local file")
+
+        try:
+            # List all files in the repo
+            repo_files = list_repo_files(repo_id)
+
+            # Find the tokenizer config file
+            tokenizer_files = [f for f in repo_files if f.endswith('tokenizer_config.json')]
+            if not tokenizer_files:
+                raise FileNotFoundError(f"No tokenizer_config.json file found in repository {repo_id}")
+
+            # Use the first tokenizer_config.json file found
+            tokenizer_config_file = tokenizer_files[0]
+            print(f"Found tokenizer config file: {tokenizer_config_file}")
+
+            # Download the file
+            tokenizer_config_file_or_name = hf_hub_download(repo_id=repo_id, filename=tokenizer_config_file)
+            print(f"Downloaded tokenizer config file to: {tokenizer_config_file_or_name}")
+            return tokenizer_config_file_or_name
+        except Exception as e:
+            raise OSError(f"Failed to download tokenizer model: {str(e)}")    
     def __init__(
         self,
         model_path: Optional[str] = None,
@@ -124,12 +149,10 @@ class HFGPTXTokenizer(PreTrainedTokenizer):
     ) -> None:
         """
         Initialize the tokenizer.
-
         Args:
             model_path (Optional[str]): Path to the tokenizer model file.
             config_path (Optional[str]): Path to the tokenizer configuration file.
             **kwargs: Additional keyword arguments passed to the superclass.
-
         This method also ensures backward compatibility by setting
         `clean_up_tokenization_spaces` to False by default.
         """
@@ -176,7 +199,6 @@ class HFGPTXTokenizer(PreTrainedTokenizer):
     def vocab_size(self) -> int:
         """
         Get the size of the tokenizer vocabulary.
-
         Returns:
             int: The size of the vocabulary.
         """
@@ -185,7 +207,6 @@ class HFGPTXTokenizer(PreTrainedTokenizer):
     def get_vocab(self) -> Dict[str, int]:
         """
         Get the vocabulary as a dictionary mapping token strings to their IDs.
-
         Returns:
             Dict[str, int]: Vocabulary mapping.
         """
@@ -196,11 +217,9 @@ class HFGPTXTokenizer(PreTrainedTokenizer):
     def _tokenize(self, text: str, **kwargs) -> List[int]:
         """
         Tokenize the input text.
-
         Args:
             text (str): Text to tokenize.
             **kwargs: Additional keyword arguments.
-
         Returns:
             List[int]: List of token IDs.
         """
@@ -210,13 +229,10 @@ class HFGPTXTokenizer(PreTrainedTokenizer):
     def _convert_token_to_id(self, token: str) -> int:
         """
         Convert a token string to its corresponding ID.
-
         Args:
             token (str): The token to convert.
-
         Returns:
             int: The token's ID.
-
         Raises:
             ValueError: If the token is unknown and cannot be encoded to a single ID.
         """
@@ -230,11 +246,9 @@ class HFGPTXTokenizer(PreTrainedTokenizer):
     ) -> str:
         """
         Decode a list of token IDs into a string.
-
         Args:
             token_ids (Union[List[int], List[List[int]]]): List of token IDs or lists of token IDs.
             num_threads (Optional[int]): Number of threads to use for decoding.
-
         Returns:
             str: Decoded string.
         """
@@ -243,10 +257,8 @@ class HFGPTXTokenizer(PreTrainedTokenizer):
     def _convert_id_to_token(self, index: int) -> str:
         """
         Convert a token ID to its corresponding token string.
-
         Args:
             index (int): Token ID.
-
         Returns:
             str: Corresponding token string.
         """
@@ -255,10 +267,8 @@ class HFGPTXTokenizer(PreTrainedTokenizer):
     def convert_tokens_to_string(self, tokens: List[str]) -> str:
         """
         Convert a list of tokens into a single string.
-
         Args:
             tokens (List[str]): List of token strings.
-
         Returns:
             str: Concatenated string of tokens.
         """
@@ -267,14 +277,11 @@ class HFGPTXTokenizer(PreTrainedTokenizer):
     def _tok_decode(self, token_ids: List[int], **kwargs: Any) -> str:
         """
         Internal method to decode token IDs with additional arguments.
-
         Args:
             token_ids (List[int]): List of token IDs.
             **kwargs: Additional arguments to pass to the decode method.
-
         Returns:
             str: Decoded string.
-
         This method also issues a warning if unsupported arguments are provided.
         """
         passed_kwargs = {key: value for (key, value) in kwargs.items() if key in self.decode_kwargs}
@@ -440,6 +447,4 @@ class SPTokenizer(HFGPTXTokenizer):
         self.chat_template = {
             lang: f"System: {sys_msg}" + "{{- '\\n'}}\n" + chat_template
             for lang, sys_msg in self.system_messages_by_lang.items()
-        }
-
-    
+        }