feat(autogpt/transformers): consume trust_remote_code (#1799)

trusting remote code by default is a danger to our users
mudler · Mar 5, 2024 · 5c69dd1 · 5c69dd1
1 parent 504f2e8
commit 5c69dd1
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/backend/python/autogptq/autogptq.py b/backend/python/autogptq/autogptq.py
@@ -33,7 +33,7 @@ def LoadModel(self, request, context):
             model = AutoGPTQForCausalLM.from_quantized(request.Model,
                     model_basename=request.ModelBaseName,
                     use_safetensors=True,
-                    trust_remote_code=True,
+                    trust_remote_code=request.TrustRemoteCode,
                     device=device,
                     use_triton=request.UseTriton,
                     quantize_config=None)

diff --git a/backend/python/transformers/transformers_server.py b/backend/python/transformers/transformers_server.py
@@ -69,9 +69,9 @@ def LoadModel(self, request, context):
         model_name = request.Model
         try:
             if request.Type == "AutoModelForCausalLM":
-                self.model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True)
+                self.model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=request.TrustRemoteCode)
             else:
-                self.model = AutoModel.from_pretrained(model_name, trust_remote_code=True)
+                self.model = AutoModel.from_pretrained(model_name, trust_remote_code=request.TrustRemoteCode)
 
             self.tokenizer = AutoTokenizer.from_pretrained(model_name)
             self.CUDA = False