From 452b35e7b7c624dcb4832fbf5d4daba4470bc621 Mon Sep 17 00:00:00 2001
From: Nikolai Kolodziej <7687617+kldzj@users.noreply.github.com>
Date: Mon, 24 Nov 2025 01:57:44 +0100
Subject: [PATCH] Add `trust_remote_code` configuration option (#31)

* Add `trust_remote_code` configuration option and apply it when loading models and tokenizers

* Default `trust_remote_code` to `None` and set it to `True` if previously `None` so the user wouldn't be asked multiple times

* Consistently access `trust_remote_code` from `self.settings` instead of the global `settings` object.

* Introduce `trusted_models` dictionary to manage and confirm `trust_remote_code` settings during model loading

* Assign `trust_remote_code` to `evaluate_model` in `trusted_models` instead of `model`
---
 src/heretic/config.py |  5 +++++
 src/heretic/model.py  | 17 ++++++++++++++++-
 2 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/src/heretic/config.py b/src/heretic/config.py
index efd2354..d1b0282 100644
--- a/src/heretic/config.py
+++ b/src/heretic/config.py
@@ -46,6 +46,11 @@ class Settings(BaseSettings):
         description="Device map to pass to Accelerate when loading the model.",
     )
 
+    trust_remote_code: bool | None = Field(
+        default=None,
+        description="Whether to trust remote code when loading the model.",
+    )
+
     batch_size: int = Field(
         default=0,  # auto
         description="Number of input sequences to process in parallel (0 = auto).",
diff --git a/src/heretic/model.py b/src/heretic/model.py
index c354fb9..2ec7ed5 100644
--- a/src/heretic/model.py
+++ b/src/heretic/model.py
@@ -39,7 +39,8 @@ class Model:
         print(f"Loading model [bold]{settings.model}[/]...")
 
         self.tokenizer: PreTrainedTokenizerBase = AutoTokenizer.from_pretrained(
-            settings.model
+            settings.model,
+            trust_remote_code=settings.trust_remote_code,
         )
 
         # Fallback for tokenizers that don't declare a special pad token.
@@ -48,6 +49,10 @@ class Model:
             self.tokenizer.padding_side = "left"
 
         self.model = None
+        self.trusted_models = {settings.model: settings.trust_remote_code}
+
+        if self.settings.evaluate_model is not None:
+            self.trusted_models[settings.evaluate_model] = settings.trust_remote_code
 
         for dtype in settings.dtypes:
             print(f"* Trying dtype [bold]{dtype}[/]... ", end="")
@@ -57,8 +62,14 @@ class Model:
                     settings.model,
                     dtype=dtype,
                     device_map=settings.device_map,
+                    trust_remote_code=self.trusted_models.get(settings.model),
                 )
 
+                # If we reach this point and the model requires trust_remote_code,
+                # the user must have confirmed it.
+                if self.trusted_models.get(settings.model) is None:
+                    self.trusted_models[settings.model] = True
+
                 # A test run can reveal dtype-related problems such as the infamous
                 # "RuntimeError: probability tensor contains either `inf`, `nan` or element < 0"
                 # (https://github.com/meta-llama/llama/issues/380).
@@ -93,8 +104,12 @@ class Model:
             self.settings.model,
             dtype=dtype,
             device_map=self.settings.device_map,
+            trust_remote_code=self.trusted_models.get(self.settings.model),
         )
 
+        if self.trusted_models.get(self.settings.model) is None:
+            self.trusted_models[self.settings.model] = True
+
     def get_layers(self) -> ModuleList:
         # Most multimodal models.
         with suppress(Exception):