download model without loading

src/xllm/cli/prepare.py

@@ -12,9 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Tuple, Type
+from typing import Type
 
-from transformers import HfArgumentParser, PreTrainedModel, PreTrainedTokenizer
+from transformers import HfArgumentParser
 
 from ..core.config import Config
 from ..run.prepare import prepare
@@ -23,7 +23,7 @@
 
 def cli_run_prepare(
     config_cls: Type[Config] = Config,
-) -> Tuple[PreTrainedTokenizer, PreTrainedModel]:
+) -> None:
     """
     Provides a command-line interface (CLI) entry point for setting up a tokenizer and model based on a configuration.
 
@@ -64,8 +64,7 @@ def cli_run_prepare(
     parser = HfArgumentParser(config_cls)
     config = parser.parse_args_into_dataclasses()[0]
     setup_cli(config=config, logger_path="./xllm_prepare.log")
-    tokenizer, model = prepare(config=config)
-    return tokenizer, model
+    prepare(config=config)
 
 
 if __name__ == "__main__":

src/xllm/run/prepare.py

@@ -13,21 +13,18 @@
 # limitations under the License.
 
 import json
-from typing import Tuple
 
 from loguru import logger
 from transformers import (
-    AutoModelForCausalLM,
     AutoTokenizer,
-    PreTrainedModel,
-    PreTrainedTokenizer,
 )
+from transformers.modeling_utils import CONFIG_NAME, cached_file
 
 from ..core.config import Config
 from ..datasets.registry import datasets_registry
 
 
-def prepare(config: Config) -> Tuple[PreTrainedTokenizer, PreTrainedModel]:
+def prepare(config: Config) -> None:
     """
     Prepares the tokenizer and model for use from the provided configuration, and optionally prepares a dataset
     if specified.
@@ -89,10 +86,23 @@ def prepare(config: Config) -> Tuple[PreTrainedTokenizer, PreTrainedModel]:
     else:
         logger.warning("Dataset is not prepared because this set in config")
 
-    tokenizer = AutoTokenizer.from_pretrained(config.correct_tokenizer_name_or_path)
+    # tokenizer
+    _ = AutoTokenizer.from_pretrained(config.correct_tokenizer_name_or_path)
     logger.info(f"Tokenizer {config.correct_tokenizer_name_or_path} loaded")
 
-    model = AutoModelForCausalLM.from_pretrained(config.model_name_or_path)
+    # model
+    cached_file(
+        config.model_name_or_path,
+        CONFIG_NAME,
+        cache_dir=None,
+        force_download=False,
+        resume_download=False,
+        proxies=None,
+        local_files_only=False,
+        token=None,
+        revision="main",
+        subfolder="",
+        _raise_exceptions_for_missing_entries=False,
+        _raise_exceptions_for_connection_errors=False,
+    )
     logger.info(f"Model {config.model_name_or_path} loaded")
-
-    return tokenizer, model