diff --git a/src/python/txtai/models/pooling/base.py b/src/python/txtai/models/pooling/base.py index d69c96dc..24fc4e01 100644 --- a/src/python/txtai/models/pooling/base.py +++ b/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 diff --git a/src/python/txtai/pipeline/hfmodel.py b/src/python/txtai/pipeline/hfmodel.py index db5f3afe..d434d783 100644 --- a/src/python/txtai/pipeline/hfmodel.py +++ b/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 diff --git a/src/python/txtai/pipeline/image/caption.py b/src/python/txtai/pipeline/image/caption.py index cc54c62e..75d7fbba 100644 --- a/src/python/txtai/pipeline/image/caption.py +++ b/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)