From 5aab87fb85f30c3b8a2ccfe64665cc703afd7c61 Mon Sep 17 00:00:00 2001 From: Frank Liu Date: Sun, 18 Aug 2024 17:14:11 -0700 Subject: [PATCH] [djl-convert] Sets default max model size limit for importing --- .../djl_converter/huggingface_converter.py | 21 +++++++++++++++++++ .../sentence_similarity_converter.py | 3 ++- 2 files changed, 23 insertions(+), 1 deletion(-) diff --git a/extensions/tokenizers/src/main/python/djl_converter/huggingface_converter.py b/extensions/tokenizers/src/main/python/djl_converter/huggingface_converter.py index c57148fc8eb..2882fe94208 100644 --- a/extensions/tokenizers/src/main/python/djl_converter/huggingface_converter.py +++ b/extensions/tokenizers/src/main/python/djl_converter/huggingface_converter.py @@ -50,6 +50,7 @@ def __init__(self): self.translator = None self.inputs = None self.outputs = None + self.max_model_size = int(os.getenv("MAX_MODEL_SIZE", "2_000_000_000")) self.api = HfApi() def save_model(self, model_info, task: str, args: Namespace, temp_dir: str, @@ -103,6 +104,9 @@ def save_onnx_model(self, model_info, task: str, args: Namespace, temp_dir, hf_pipeline, include_types) if model_zoo: + model_size = self.get_dir_size(temp_dir) + if model_size > self.max_model_size: + return False, f"Model size too large: {model_size}", -1 size = self.save_to_model_zoo(model_info, args.output_dir, "OnnxRuntime", temp_dir, arguments) else: @@ -183,6 +187,9 @@ def save_rust_model(self, model_info, args: Namespace, temp_dir: str, arguments = self.save_serving_properties(model_info, "Rust", temp_dir, hf_pipeline, include_types) if model_zoo: + model_size = self.get_dir_size(temp_dir) + if model_size > self.max_model_size: + return False, f"Model size too large: {model_size}", -1 size = self.save_to_model_zoo(model_info, args.output_dir, "Rust", temp_dir, arguments) else: @@ -231,6 +238,9 @@ def save_pytorch_model(self, model_info, args: Namespace, temp_dir: str, temp_dir, hf_pipeline, include_types) if model_zoo: + model_size = self.get_dir_size(temp_dir) + if model_size > self.max_model_size: + return False, f"Model size too large: {model_size}", -1 size = self.save_to_model_zoo(model_info, args.output_dir, "PyTorch", temp_dir, arguments) else: @@ -332,6 +342,17 @@ def save_to_model_zoo(self, model_info, output_dir: str, engine: str, return file_size + @staticmethod + def get_dir_size(path: str) -> int: + total = 0 + with os.scandir(path) as it: + for entry in it: + if entry.is_file(): + total += entry.stat().st_size + elif entry.is_dir(): + total += HuggingfaceConverter.get_dir_size(entry.path) + return total + def verify_jit_model(self, hf_pipeline, model_file: str, include_types: bool, cpu_only: bool): logging.info( diff --git a/extensions/tokenizers/src/main/python/djl_converter/sentence_similarity_converter.py b/extensions/tokenizers/src/main/python/djl_converter/sentence_similarity_converter.py index 8c3d394bc15..b0ac589d67c 100644 --- a/extensions/tokenizers/src/main/python/djl_converter/sentence_similarity_converter.py +++ b/extensions/tokenizers/src/main/python/djl_converter/sentence_similarity_converter.py @@ -163,5 +163,6 @@ def get_extra_arguments(self, hf_pipeline, model_id: str, @staticmethod def save_module_weight(model_id: str, temp_dir: str, layer: str, name: str): - file = self.get_file(model_id, f"{layer}/model.safetensors") + file = SentenceSimilarityConverter.get_file( + model_id, f"{layer}/model.safetensors") shutil.copyfile(file, os.path.join(temp_dir, f"{name}.safetensors"))