Added cache directory support for Colpali

byaldi/RAGModel.py

@@ -45,6 +45,7 @@ def from_pretrained(
         index_root: str = ".byaldi",
         device: str = "cuda",
         verbose: int = 1,
+        cache_dir: Optional[str] = "/cache_dir/models",
     ):
         """Load a ColPali model from a pre-trained checkpoint.
 
@@ -61,6 +62,7 @@ def from_pretrained(
             index_root=index_root,
             device=device,
             verbose=verbose,
+            cache_dir=cache_dir,
         )
         return instance
 

byaldi/colpali.py

@@ -76,6 +76,7 @@ def __init__(
                     else None
                 ),
                 token=kwargs.get("hf_token", None) or os.environ.get("HF_TOKEN"),
+                cache_dir=kwargs.get("cache_dir", None),
             )
         elif "colqwen2" in pretrained_model_name_or_path.lower():
             self.model = ColQwen2.from_pretrained(