Default to trust_remote_code=False, fall back only when required (addresses #293)

air_llm/airllm/airllm_base.py

@@ -16,8 +16,8 @@
 from .profiler import LayeredProfiler
 
 from .utils import clean_memory, load_layer, layer_tensor_names, load_layer_subset, \
-    find_or_create_local_splitted_path, load_merged_ngram_embedding, \
-    open_ngram_mmap_table, MmapEmbedding, _force_meta_embeddings
+    find_or_create_local_splitted_path, load_prefer_no_remote_code, \
+    load_merged_ngram_embedding, open_ngram_mmap_table, MmapEmbedding, _force_meta_embeddings
 from .persist import ModelPersister
 
 try:
@@ -211,10 +211,11 @@ def get_generation_config(self):
             return GenerationConfig()
 
     def get_tokenizer(self, hf_token=None):
-        if hf_token is not None:
-            return AutoTokenizer.from_pretrained(self.model_local_path, token=hf_token, trust_remote_code=True)
-        else:
-            return AutoTokenizer.from_pretrained(self.model_local_path, trust_remote_code=True)
+        token_kwargs = {'token': hf_token} if hf_token is not None else {}
+        # Prefer transformers' native tokenizer; only trust the repo's remote code if it's required
+        # (custom tokenizers). Matches how the config/model are loaded above.
+        return load_prefer_no_remote_code(
+            AutoTokenizer.from_pretrained, self.model_local_path, **token_kwargs)
 
     # ---- model construction -----------------------------------------------------------------
 

air_llm/airllm/airllm_llama_mlx.py

@@ -15,7 +15,7 @@
 import psutil
 from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, AutoModel, GenerationMixin, LlamaForCausalLM, GenerationConfig
 from .utils import clean_memory, load_layer, \
-    find_or_create_local_splitted_path
+    find_or_create_local_splitted_path, load_prefer_no_remote_code
 
 
 
@@ -227,10 +227,9 @@ def __init__(self, model_local_path_or_repo_id, device="cuda:0", dtype=None, max
                                                                                          layer_names=self.layer_names_dict,
                                                                                          hf_token=hf_token,
                                                                                          delete_original=delete_original)
-        if hf_token is not None:
-            self.config = AutoConfig.from_pretrained(self.model_local_path, token=hf_token, trust_remote_code=True)
-        else:
-            self.config = AutoConfig.from_pretrained(self.model_local_path, trust_remote_code=True)
+        token_kwargs = {'token': hf_token} if hf_token is not None else {}
+        self.config = load_prefer_no_remote_code(
+            AutoConfig.from_pretrained, self.model_local_path, **token_kwargs)
 
 
         self.model_args = get_model_args_from_config(self.config)
@@ -243,10 +242,9 @@ def __init__(self, model_local_path_or_repo_id, device="cuda:0", dtype=None, max
 
 
     def get_tokenizer(self, hf_token=None):
-        if hf_token is not None:
-            return AutoTokenizer.from_pretrained(self.model_local_path, token=hf_token, trust_remote_code=True)
-        else:
-            return AutoTokenizer.from_pretrained(self.model_local_path, trust_remote_code=True)
+        token_kwargs = {'token': hf_token} if hf_token is not None else {}
+        return load_prefer_no_remote_code(
+            AutoTokenizer.from_pretrained, self.model_local_path, **token_kwargs)
 
 
     def generate(self, x, temperature=0, max_new_tokens=None, **kwargs):

air_llm/airllm/auto_model.py

@@ -2,6 +2,8 @@
 from transformers import AutoConfig
 from sys import platform
 
+from .utils import load_prefer_no_remote_code
+
 is_on_mac_os = False
 
 if platform == "darwin":
@@ -36,11 +38,11 @@ def __init__(self):
 
     @classmethod
     def get_module_class(cls, pretrained_model_name_or_path, *inputs, **kwargs):
-        if 'hf_token' in kwargs:
-            config = AutoConfig.from_pretrained(pretrained_model_name_or_path, trust_remote_code=True,
-                                                token=kwargs['hf_token'])
-        else:
-            config = AutoConfig.from_pretrained(pretrained_model_name_or_path, trust_remote_code=True)
+        token_kwargs = {'token': kwargs['hf_token']} if 'hf_token' in kwargs else {}
+        # Prefer transformers' native config; only run the repo's remote code if it's actually
+        # required to parse the config (custom architectures). See load_prefer_no_remote_code.
+        config = load_prefer_no_remote_code(
+            AutoConfig.from_pretrained, pretrained_model_name_or_path, **token_kwargs)
 
         architectures = getattr(config, "architectures", None) or []
         arch = architectures[0] if architectures else ""

air_llm/airllm/utils.py

@@ -85,6 +85,24 @@ def clean_memory():
     torch.cuda.empty_cache()
 
 
+def load_prefer_no_remote_code(loader, *args, **kwargs):
+    """
+    Call a transformers ``from_pretrained``-style loader, preferring transformers' native
+    implementation and only trusting the model's bundled remote code when the native load fails.
+
+    AirLLM's advertised entry point is an arbitrary Hugging Face repo id, so a normal load of a
+    standard architecture should not execute Python shipped in the repo. Standard models load fine
+    with ``trust_remote_code=False``; only models that genuinely ship custom code (ChatGLM, Baichuan,
+    some Qwen) raise without it, so we fall back to ``trust_remote_code=True`` for those. This mirrors
+    how ``AirLLMBaseModel`` already loads the model config/weights, and keeps every load boundary
+    (config, tokenizer, model) consistent instead of hard-coding ``trust_remote_code=True``.
+    """
+    try:
+        return loader(*args, trust_remote_code=False, **kwargs)
+    except Exception:
+        return loader(*args, trust_remote_code=True, **kwargs)
+
+
 def uncompress_layer_state_dict(layer_state_dict):
     uncompressed_layer_state_dict = None
     if any(['4bit' in k for k in layer_state_dict.keys()]):

air_llm/tests/test_trust_remote_code.py

@@ -0,0 +1,95 @@
+import unittest
+
+import airllm.auto_model as auto_model_mod
+from airllm.auto_model import AutoModel
+from airllm.utils import load_prefer_no_remote_code
+
+
+class _FakeConfig:
+    def __init__(self, architectures):
+        self.architectures = architectures
+
+
+class TestLoadPreferNoRemoteCode(unittest.TestCase):
+    def test_uses_false_and_does_not_retry_when_native_load_succeeds(self):
+        calls = []
+
+        def loader(path, trust_remote_code=None, **kwargs):
+            calls.append(trust_remote_code)
+            return 'ok'
+
+        result = load_prefer_no_remote_code(loader, 'some/repo')
+
+        self.assertEqual(result, 'ok')
+        self.assertEqual(calls, [False])  # tried False, never fell back to True
+
+    def test_falls_back_to_true_when_native_load_raises(self):
+        calls = []
+
+        def loader(path, trust_remote_code=None, **kwargs):
+            calls.append(trust_remote_code)
+            if trust_remote_code is False:
+                raise ValueError('requires trust_remote_code=True')
+            return 'ok-remote'
+
+        result = load_prefer_no_remote_code(loader, 'custom/repo')
+
+        self.assertEqual(result, 'ok-remote')
+        self.assertEqual(calls, [False, True])  # tried False first, then True
+
+    def test_forwards_extra_kwargs(self):
+        seen = {}
+
+        def loader(path, trust_remote_code=None, **kwargs):
+            seen.update(kwargs)
+            seen['path'] = path
+            return 'ok'
+
+        load_prefer_no_remote_code(loader, 'some/repo', token='abc')
+
+        self.assertEqual(seen['path'], 'some/repo')
+        self.assertEqual(seen['token'], 'abc')
+
+
+class TestGetModuleClassTrust(unittest.TestCase):
+    """AutoModel.get_module_class must not enable remote code unless it is actually required."""
+
+    def setUp(self):
+        self._orig = auto_model_mod.AutoConfig.from_pretrained
+
+    def tearDown(self):
+        auto_model_mod.AutoConfig.from_pretrained = self._orig
+
+    def test_standard_arch_loads_config_without_remote_code(self):
+        calls = []
+
+        def fake(path, trust_remote_code=None, **kwargs):
+            calls.append(trust_remote_code)
+            return _FakeConfig(['LlamaForCausalLM'])
+
+        auto_model_mod.AutoConfig.from_pretrained = staticmethod(fake)
+
+        module, cls = AutoModel.get_module_class('some/llama-repo')
+
+        self.assertEqual((module, cls), ('airllm', 'AirLLMBaseModel'))
+        self.assertEqual(calls, [False])  # native path only; remote code never enabled
+
+    def test_custom_arch_falls_back_to_remote_code(self):
+        calls = []
+
+        def fake(path, trust_remote_code=None, **kwargs):
+            calls.append(trust_remote_code)
+            if trust_remote_code is False:
+                raise ValueError('requires trust_remote_code=True')
+            return _FakeConfig(['ChatGLMModel'])
+
+        auto_model_mod.AutoConfig.from_pretrained = staticmethod(fake)
+
+        module, cls = AutoModel.get_module_class('some/chatglm-repo')
+
+        self.assertEqual((module, cls), ('airllm', 'AirLLMChatGLM'))
+        self.assertEqual(calls, [False, True])  # only trusts remote code as a fallback
+
+
+if __name__ == '__main__':
+    unittest.main()