diff --git a/extensions/tokenizers/src/main/python/djl_converter/arg_parser.py b/extensions/tokenizers/src/main/python/djl_converter/arg_parser.py index b4541b0eae2..a7a0edad382 100644 --- a/extensions/tokenizers/src/main/python/djl_converter/arg_parser.py +++ b/extensions/tokenizers/src/main/python/djl_converter/arg_parser.py @@ -92,6 +92,9 @@ def importer_args(): " repository. This option should only be set for repositories you trust and in which" " you have read the code, as it will execute on your local machine arbitrary code" " present in the model repository.") + parser.add_argument( + "--min-version", + help="Requires a specific version of DJL to load the model.") args = parser.parse_args() if args.output_dir is None: diff --git a/extensions/tokenizers/src/main/python/djl_converter/huggingface_models.py b/extensions/tokenizers/src/main/python/djl_converter/huggingface_models.py index 43a15f74728..59d65a958be 100644 --- a/extensions/tokenizers/src/main/python/djl_converter/huggingface_models.py +++ b/extensions/tokenizers/src/main/python/djl_converter/huggingface_models.py @@ -14,7 +14,7 @@ import logging import os from argparse import Namespace -from typing import List +from typing import List, Optional from huggingface_hub import HfApi from huggingface_hub import hf_hub_download @@ -150,7 +150,8 @@ def list_models(self, args: Namespace) -> List[dict]: return ret def update_progress(self, model_info: ModelInfo, application: str, - result: bool, reason: str, size: int, cpu_only: bool): + result: bool, reason: str, size: int, cpu_only: bool, + min_version: Optional[str]): status = { "result": "success" if result else "failed", "application": application, @@ -162,6 +163,8 @@ def update_progress(self, model_info: ModelInfo, application: str, status["reason"] = reason if cpu_only: status["cpu_only"] = True + if result and min_version: + status["requires"] = min_version self.processed_models[model_info.modelId] = status diff --git a/extensions/tokenizers/src/main/python/djl_converter/model_zoo_importer.py b/extensions/tokenizers/src/main/python/djl_converter/model_zoo_importer.py index 909563ced6b..e6fb7aba1a5 100644 --- a/extensions/tokenizers/src/main/python/djl_converter/model_zoo_importer.py +++ b/extensions/tokenizers/src/main/python/djl_converter/model_zoo_importer.py @@ -54,7 +54,8 @@ def main(): size = -1 huggingface_models.update_progress(model_info, converter.application, - result, reason, size, args.cpu_only) + result, reason, size, args.cpu_only, + args.min_version) if os.path.exists(temp_dir): shutil.rmtree(temp_dir)