From 29230777378f798136d814b8b84935bdbefa396c Mon Sep 17 00:00:00 2001 From: Byran Liu Date: Fri, 27 Dec 2024 13:31:25 -0800 Subject: [PATCH] [tokenizer] fixes trust_remote_code issues for DJL converter --- .../python/djl_converter/huggingface_converter.py | 10 ++++++---- .../djl_converter/sentence_similarity_converter.py | 12 ++++++++---- 2 files changed, 14 insertions(+), 8 deletions(-) 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 b0c4ed25f8d..17353b66b5f 100644 --- a/extensions/tokenizers/src/main/python/djl_converter/huggingface_converter.py +++ b/extensions/tokenizers/src/main/python/djl_converter/huggingface_converter.py @@ -136,7 +136,8 @@ def save_rust_model(self, model_info, args: Namespace, temp_dir: str, if not os.path.exists(temp_dir): os.makedirs(temp_dir) - tokenizer = AutoTokenizer.from_pretrained(model_id) + tokenizer = AutoTokenizer.from_pretrained( + model_id, trust_remote_code=args.trust_remote_code) include_types = config.model_type not in [ "distilbert", "mistral", "qwen2", "gemma2" ] @@ -208,7 +209,7 @@ def save_pytorch_model(self, model_info, args: Namespace, temp_dir: str, os.makedirs(temp_dir) try: - hf_pipeline = self.load_model(model_id) + hf_pipeline = self.load_model(model_id, args.trust_remote_code) except Exception as e: logging.warning(f"Failed to load model: {model_id}.") logging.warning(e, exc_info=True) @@ -410,11 +411,12 @@ def verify_jit_output(self, hf_pipeline, encoding, out): return True, None - def load_model(self, model_id: str): + def load_model(self, model_id: str, trust_remote_code: bool): logging.info(f"Loading model: {model_id} ...") kwargs = { "tokenizer": model_id, - "device": -1 # always use CPU to trace the model + "device": -1, # always use CPU to trace the model + "trust_remote_code": trust_remote_code } return pipeline(task=self.task, model=model_id, 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 b0ac589d67c..f46925aeaea 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 @@ -32,10 +32,13 @@ def __init__(self): self.inputs = "This is an example sentence" self.outputs = 0 - def load_model(self, model_id: str): + def load_model(self, model_id: str, trust_remote_code: bool): logging.info(f"Loading model: {model_id} ...") - tokenizer = AutoTokenizer.from_pretrained(model_id) - model = AutoModel.from_pretrained(model_id) + tokenizer = AutoTokenizer.from_pretrained( + model_id, trust_remote_code=trust_remote_code) + + model = AutoModel.from_pretrained(model_id, + trust_remote_code=trust_remote_code) return PipelineHolder(tokenizer, model) @@ -78,7 +81,8 @@ def get_extra_arguments(self, hf_pipeline, model_id: str, if hasattr(hf_pipeline.model, "config"): config = hf_pipeline.model.config else: - config = AutoConfig.from_pretrained(model_id) + config = AutoConfig.from_pretrained( + model_id, trust_remote_code=trust_remote_code) tokenizer = hf_pipeline.tokenizer if hasattr(config, "max_position_embeddings") and hasattr( tokenizer, "model_max_length"):