Added new argument trust_remote_code in infer which will help for emote code execution

Signed-off-by: Amit Raj <[email protected]>

Author: quic-amitraj

QEfficient/base/common.py

@@ -40,9 +40,8 @@ def from_pretrained(cls, pretrained_model_name_or_path: str, *args, **kwargs) ->
         """
         Downloads HuggingFace model if already doesn't exist locally, returns QEFFAutoModel object based on type of model.
         """
-        config = AutoConfig.from_pretrained(pretrained_model_name_or_path, trust_remote_code=True)
+        config = AutoConfig.from_pretrained(pretrained_model_name_or_path, *args, **kwargs)
 
-        # class_name = MODEL_CLASS_MAPPING.get(config.__class__.__name__, None) OR MODEL_EXTERNAL_CLASS_MAPPING(config.__class__.__name__)
         class_name = (
             MODEL_CLASS_MAPPING.get(config.__class__.__name__, None)
             or EXTERNAL_MODEL_CLASS_MAPPING[config.__class__.__name__]
@@ -65,7 +64,6 @@ def from_pretrained(cls, pretrained_model_name_or_path: str, *args, **kwargs) ->
             pretrained_model_name_or_path=(local_model_dir if local_model_dir else pretrained_model_name_or_path),
             token=hf_token,
             continuous_batching=continuous_batching,
-            trust_remote_code=True,
             **kwargs,
         )
         return qeff_model

QEfficient/cloud/infer.py

@@ -111,6 +111,7 @@ def main(
     allow_mxint8_mdp_io: bool = False,
     enable_qnn: Optional[bool] = False,
     qnn_config: Optional[str] = None,
+    trust_remote_code: Optional[bool] = False,
     **kwargs,
 ) -> None:
     """
@@ -140,6 +141,7 @@ def main(
         :allow_mxint8_mdp_io (bool): Allows MXINT8 compression of MDP IO traffic. ``Defaults to False.``
         :enable_qnn (bool): Enables QNN Compilation. ``Defaults to False.``
         :qnn_config (str): Path of QNN Config parameters file. ``Defaults to None.``
+        :trust_remote_code (bool): Trust remote code execution. ``Defaults to False.``
         :kwargs: Pass any compiler option as input. Any flag that is supported by `qaic-exec` can be passed. Params are converted to flags as below:
                 -allocator_dealloc_delay=1 -> -allocator-dealloc-delay=1
                 -qpc_crc=True -> -qpc-crc
@@ -164,6 +166,7 @@ def main(
         hf_token=hf_token,
         full_batch_size=full_batch_size,
         local_model_dir=local_model_dir,
+        trust_remote_code=trust_remote_code,
     )
 
     image_path = kwargs.pop("image_path", None)
@@ -264,6 +267,12 @@ def main(
         action="store_true",
         help="Compress constant MatMul weights to MXFP6 E2M3, default is no compression",
     )
+    parser.add_argument(
+        "--trust_remote_code",
+        action="store_true",
+        default=False,
+        help="Enable trusting remote code when loading models. Default is False; set to True by passing this flag.",
+    )
     parser.add_argument(
         "--mxint8",
         "--mxint8_kv_cache",

QEfficient/transformers/modeling_utils.py

@@ -284,6 +284,7 @@ def build_model_class_mapping(auto_model_class, qeff_class_name):
 
 
 EXTERNAL_MODEL_CLASS_MAPPING = {"Grok1Config": "QEFFAutoModelForCausalLM"}
+
 MODEL_CLASS_MAPPING = {
     **build_model_class_mapping(mapping.AutoModelForCausalLM, "QEFFAutoModelForCausalLM"),
     **build_model_class_mapping(mapping.AutoModelForImageTextToText, "QEFFAutoModelForImageTextToText"),