save

app.py

@@ -97,7 +97,7 @@ def compress_and_upload(
 
         # --- 1. Load Model and Tokenizer ---
         model = AutoModelForCausalLM.from_pretrained(
-            model_id, torch_dtype="auto", device_map=None, token=token
+            model_id, torch_dtype="auto", device_map=None, token=token, trust_remote_code=True
         )
 
         output_dir = f"{model_id.split('/')[-1]}-{quant_method}"

tests/test_app.py

@@ -112,7 +112,7 @@ def test_compress_and_upload_success(
 
     mock_whoami.assert_called_once_with(token="test_token")
     mock_auto_model_for_causal_lm.from_pretrained.assert_called_once_with(
-        model_id, torch_dtype="auto", device_map=None, token="test_token"
+        model_id, torch_dtype="auto", device_map=None, token="test_token", trust_remote_code=True
     )
     mock_oneshot.assert_called_once()
     assert mock_oneshot.call_args[1]["model"] == mock_auto_model_for_causal_lm.from_pretrained.return_value
@@ -135,6 +135,23 @@ def test_compress_and_upload_success(
     assert "✅ Success!" in result
     assert "https://huggingface.co/test_user/test_model-AWQ" in result
 
+def test_compress_and_upload_with_trust_remote_code(
+    mock_hf_api,
+    mock_whoami,
+    mock_auto_model_for_causal_lm,
+    mock_oneshot,
+    mock_model_card,
+    mock_gr_oauth_token,
+):
+    model_id = "org/test_model"
+    quant_method = "AWQ"
+    compress_and_upload(model_id, quant_method, mock_gr_oauth_token)
+
+    mock_auto_model_for_causal_lm.from_pretrained.assert_called_once_with(
+        model_id, torch_dtype="auto", device_map=None, token="test_token", trust_remote_code=True
+    )
+
+
 def test_compress_and_upload_model_no_architecture(
     mock_hf_api,
     mock_whoami,