black and lint

colpali_engine/dataset/hard_neg_collator.py

@@ -1,6 +1,8 @@
+from random import randint
+
 from datasets import Dataset, DatasetDict
 from transformers import PreTrainedTokenizer, ProcessorMixin
-from random import randint
+
 from .custom_collator import CustomCollator
 
 

colpali_engine/interpretability/gen_interpretability_plots.py

@@ -11,8 +11,7 @@
 
 from colpali_engine.interpretability.plot_utils import plot_patches
 from colpali_engine.interpretability.processor import ColPaliProcessor
-from colpali_engine.interpretability.torch_utils import \
-    normalize_attention_map_per_query_token
+from colpali_engine.interpretability.torch_utils import normalize_attention_map_per_query_token
 from colpali_engine.interpretability.vit_configs import VIT_CONFIG
 from colpali_engine.models.paligemma_colbert_architecture import ColPali
 

colpali_engine/models/colbert_architectures.py

@@ -1,8 +1,14 @@
 from torch import nn
-from transformers import (BertModel, BertPreTrainedModel, CamembertModel,
-                          CamembertPreTrainedModel, LlamaModel,
-                          LlamaPreTrainedModel, XLMRobertaModel,
-                          XLMRobertaPreTrainedModel)
+from transformers import (
+    BertModel,
+    BertPreTrainedModel,
+    CamembertModel,
+    CamembertPreTrainedModel,
+    LlamaModel,
+    LlamaPreTrainedModel,
+    XLMRobertaModel,
+    XLMRobertaPreTrainedModel,
+)
 
 
 class ColCamembert(CamembertPreTrainedModel):

colpali_engine/models/paligemma_colbert_architecture.py

@@ -1,7 +1,6 @@
 import torch
 from torch import nn
-from transformers.models.paligemma.modeling_paligemma import (
-    PaliGemmaForConditionalGeneration, PaliGemmaPreTrainedModel)
+from transformers.models.paligemma.modeling_paligemma import PaliGemmaForConditionalGeneration, PaliGemmaPreTrainedModel
 
 
 class BiPaliLast(PaliGemmaPreTrainedModel):

colpali_engine/utils/train_colpali_engine_models.py

@@ -9,17 +9,18 @@
 from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
 from torch.utils.data import DataLoader
 from tqdm import tqdm
-from transformers import (AutoTokenizer, Idefics2Processor, PreTrainedModel,
-                          PreTrainedTokenizer, TrainingArguments)
+from transformers import AutoTokenizer, Idefics2Processor, PreTrainedModel, PreTrainedTokenizer, TrainingArguments
 
 from colpali_engine.dataset.custom_collator import CustomCollator
 from colpali_engine.dataset.hard_neg_collator import HardNegCollator
-from colpali_engine.loss.colbert_loss import (BiEncoderLoss, BiPairwiseCELoss,
-                                              ColbertLoss,
-                                              ColbertPairwiseCELoss,
-                                              ColbertPairwiseNegativeCELoss)
-from colpali_engine.trainer.contrastive_trainer import (
-    ContrastiveNegativeTrainer, ContrastiveTrainer)
+from colpali_engine.loss.colbert_loss import (
+    BiEncoderLoss,
+    BiPairwiseCELoss,
+    ColbertLoss,
+    ColbertPairwiseCELoss,
+    ColbertPairwiseNegativeCELoss,
+)
+from colpali_engine.trainer.contrastive_trainer import ContrastiveNegativeTrainer, ContrastiveTrainer
 from colpali_engine.trainer.retrieval_evaluator import CustomEvaluator
 from colpali_engine.utils.gpu_stats import print_gpu_utilization, print_summary
 

colpali_engine/utils/wrapper.py

@@ -1,17 +1,22 @@
 import importlib
 
 from colpali_engine.models.clip_baselines import ColSigLIP, SigLIP
-from colpali_engine.models.colbert_architectures import (BiBERT, BiXLMRoBERTa,
-                                                         ColBERT, ColCamembert,
-                                                         ColLlama,
-                                                         ColXLMRoBERTa)
-from colpali_engine.models.idefics_colbert_architecture import (BiIdefics,
-                                                                ColIdefics)
-from colpali_engine.models.paligemma_colbert_architecture import (BiNewSiglip,
-                                                                  BiPaliLast,
-                                                                  BiPaliMean,
-                                                                  ColNewSiglip,
-                                                                  ColPali)
+from colpali_engine.models.colbert_architectures import (
+    BiBERT,
+    BiXLMRoBERTa,
+    ColBERT,
+    ColCamembert,
+    ColLlama,
+    ColXLMRoBERTa,
+)
+from colpali_engine.models.idefics_colbert_architecture import BiIdefics, ColIdefics
+from colpali_engine.models.paligemma_colbert_architecture import (
+    BiNewSiglip,
+    BiPaliLast,
+    BiPaliMean,
+    ColNewSiglip,
+    ColPali,
+)
 
 if importlib.util.find_spec("transformers") is not None:
     from transformers import AutoProcessor, AutoTokenizer

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = 'colpali_engine'
 # dynamic = ["version"]
-version = '0.0.1'
+version = '0.1.0'
 description = 'This repository centralizes ressources for the ColPali project.'
 authors = [
     { name = 'Manuel Faysse', email = '[email protected]' },