Support trust_remote_code in associated AutoModels, closes #1342

src/python/txtai/models/pooling/base.py

@@ -36,7 +36,9 @@ class Pooling(Module):
         super().__init__()
 
         self.model = Models.load(path, modelargs=modelargs)
-        self.tokenizer = Models.tokenizer(tokenizer if tokenizer else path)
+        self.tokenizer = Models.tokenizer(
+            tokenizer if tokenizer else path, trust_remote_code=modelargs.get("trust_remote_code", False) if modelargs else False
+        )
         self.device = Models.device(device)
 
         # Detect unbounded tokenizer typically found in older models

src/python/txtai/pipeline/hfmodel.py

@@ -64,7 +64,11 @@ class HFModel(Tensors):
         model = self.prepare(model)
 
         # Load tokenizer
-        tokenizer = transformers.AutoTokenizer.from_pretrained(tokenizer) if isinstance(tokenizer, str) else tokenizer
+        tokenizer = (
+            transformers.AutoTokenizer.from_pretrained(tokenizer, trust_remote_code=kwargs.get("trust_remote_code", False))
+            if isinstance(tokenizer, str)
+            else tokenizer
+        )
 
         return model, tokenizer
 

src/python/txtai/pipeline/image/caption.py

@@ -38,8 +38,8 @@ class Caption(HFModel):
             self.model, self.tokenizer, self.processor = path
         else:
             self.model = transformers.AutoModelForImageTextToText.from_pretrained(path, **kwargs)
-            self.tokenizer = transformers.AutoTokenizer.from_pretrained(path)
-            self.processor = transformers.AutoImageProcessor.from_pretrained(path)
+            self.tokenizer = transformers.AutoTokenizer.from_pretrained(path, trust_remote_code=kwargs.get("trust_remote_code", False))
+            self.processor = transformers.AutoImageProcessor.from_pretrained(path, trust_remote_code=kwargs.get("trust_remote_code", False))
 
         # Move model to device
         self.model = self.model.to(self.device)