wip: see w/ Manu

colpali_engine/models/paligemma/colpali_2/modeling_colpali_2.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass
-from typing import Any, ClassVar, Dict, Optional, cast
+from typing import Any, ClassVar, Dict, Optional
 
 import torch
 from torch import nn
@@ -40,9 +40,7 @@ class ColPali2(PaliGemmaPreTrainedModel):
     main_input_name: ClassVar[str] = "doc_input_ids"  # transformers-related
 
     def __init__(self, config: ColPali2Config):
-        super().__init__(config=config)
-
-        self.config = cast(ColPali2Config, self.config)
+        self.config = config
         self.vlm_backbone = PaliGemmaForConditionalGeneration(self.config.vlm_config)
 
         self.single_vector_projector = nn.Linear(

colpali_engine/models/paligemma/colpali_2/publish_base_colpali_2.py

@@ -0,0 +1,29 @@
+from typing import cast
+
+import typer
+from transformers.models.paligemma.configuration_paligemma import PaliGemmaConfig
+
+from colpali_engine.models.paligemma.colpali_2.configuration_colpali_2 import ColPali2Config
+from colpali_engine.models.paligemma.colpali_2.modeling_colpali_2 import ColPali2
+
+
+def main():
+    vlm_backbone_model_name = "google/paligemma-3b-mix-448"
+    base_colpali_2_name = "vidore/colpali-v2.0"
+
+    model_config = ColPali2Config(
+        vlm_backbone_config=cast(PaliGemmaConfig, PaliGemmaConfig.from_pretrained(vlm_backbone_model_name)),
+        single_vector_projector_dim=128,
+        single_vector_pool_strategy="mean",
+        multi_vector_projector_dim=128,
+    )
+
+    model = ColPali2(config=model_config)
+
+    model.push_to_hub(base_colpali_2_name, private=True)
+
+    return
+
+
+if __name__ == "__main__":
+    typer.run(main)