Fix code style in dspy/dsp (#8176)

* enable style check

* update rules

* allow wildcard imports

* fix code style for dspy/dsp

Author: chenmoneygithub

dspy/dsp/colbertv2.py

@@ -21,7 +21,10 @@ def __init__(
         self.url = f"{url}:{port}" if port else url
 
     def __call__(
-        self, query: str, k: int = 10, simplify: bool = False,
+        self,
+        query: str,
+        k: int = 10,
+        simplify: bool = False,
     ) -> Union[list[str], list[dotdict]]:
         if self.post_requests:
             topk: list[dict[str, Any]] = colbertv2_post_request(self.url, query, k)
@@ -36,9 +39,7 @@ def __call__(
 
 @request_cache()
 def colbertv2_get_request_v2(url: str, query: str, k: int):
-    assert (
-        k <= 100
-    ), "Only k <= 100 is supported for the hosted ColBERTv2 server at the moment."
+    assert k <= 100, "Only k <= 100 is supported for the hosted ColBERTv2 server at the moment."
 
     payload = {"query": query, "k": k}
     res = requests.get(url, params=payload, timeout=10)
@@ -72,88 +73,112 @@ def colbertv2_post_request_v2_wrapped(*args, **kwargs):
 
 colbertv2_post_request = colbertv2_post_request_v2_wrapped
 
+
 class ColBERTv2RetrieverLocal:
-    def __init__(self,passages:List[str],colbert_config=None,load_only:bool=False):
+    def __init__(self, passages: List[str], colbert_config=None, load_only: bool = False):
         """Colbertv2 retriever module
 
         Args:
             passages (List[str]): list of passages
             colbert_config (ColBERTConfig, optional): colbert config for building and searching. Defaults to None.
             load_only (bool, optional): whether to load the index or build and then load. Defaults to False.
         """
-        assert colbert_config is not None, "Please pass a valid colbert_config, which you can import from colbert.infra.config import ColBERTConfig and modify it"
+        assert (
+            colbert_config is not None
+        ), "Please pass a valid colbert_config, which you can import from colbert.infra.config import ColBERTConfig and modify it"
         self.colbert_config = colbert_config
 
-        assert self.colbert_config.checkpoint is not None, "Please pass a valid checkpoint like colbert-ir/colbertv2.0, which you can modify in the ColBERTConfig with attribute name checkpoint"
+        assert (
+            self.colbert_config.checkpoint is not None
+        ), "Please pass a valid checkpoint like colbert-ir/colbertv2.0, which you can modify in the ColBERTConfig with attribute name checkpoint"
         self.passages = passages
-        
-        assert self.colbert_config.index_name is not None, "Please pass a valid index_name, which you can modify in the ColBERTConfig with attribute name index_name"
+
+        assert (
+            self.colbert_config.index_name is not None
+        ), "Please pass a valid index_name, which you can modify in the ColBERTConfig with attribute name index_name"
         self.passages = passages
 
         if not load_only:
-            print(f"Building the index for experiment {self.colbert_config.experiment} with index name {self.colbert_config.index_name}")
+            print(
+                f"Building the index for experiment {self.colbert_config.experiment} with index name "
+                f"{self.colbert_config.index_name}"
+            )
             self.build_index()
-        
-        print(f"Loading the index for experiment {self.colbert_config.experiment} with index name {self.colbert_config.index_name}")
+
+        print(
+            f"Loading the index for experiment {self.colbert_config.experiment} with index name "
+            f"{self.colbert_config.index_name}"
+        )
         self.searcher = self.get_index()
 
     def build_index(self):
-
         try:
-            import colbert # noqa: F401
+            import colbert  # noqa: F401
         except ImportError:
-            print("Colbert not found. Please check your installation or install the module using pip install colbert-ai[faiss-gpu,torch].")
+            print(
+                "Colbert not found. Please check your installation or install the module using pip install "
+                "colbert-ai[faiss-gpu,torch]."
+            )
 
         from colbert import Indexer
         from colbert.infra import Run, RunConfig
-        with Run().context(RunConfig(nranks=self.colbert_config.nranks, experiment=self.colbert_config.experiment)):  
+
+        with Run().context(RunConfig(nranks=self.colbert_config.nranks, experiment=self.colbert_config.experiment)):
             indexer = Indexer(checkpoint=self.colbert_config.checkpoint, config=self.colbert_config)
             indexer.index(name=self.colbert_config.index_name, collection=self.passages, overwrite=True)
 
     def get_index(self):
         try:
-            import colbert # noqa: F401
+            import colbert  # noqa: F401
         except ImportError:
-            print("Colbert not found. Please check your installation or install the module using pip install colbert-ai[faiss-gpu,torch].")
+            print(
+                "Colbert not found. Please check your installation or install the module using pip install "
+                "colbert-ai[faiss-gpu,torch]."
+            )
 
         from colbert import Searcher
         from colbert.infra import Run, RunConfig
-        
+
         with Run().context(RunConfig(experiment=self.colbert_config.experiment)):
             searcher = Searcher(index=self.colbert_config.index_name, collection=self.passages)
         return searcher
-    
+
     def __call__(self, *args: Any, **kwargs: Any) -> Any:
         return self.forward(*args, **kwargs)
 
-    def forward(self,query:str,k:int=7,**kwargs):
+    def forward(self, query: str, k: int = 7, **kwargs):
         import torch
-        
+
         if kwargs.get("filtered_pids"):
             filtered_pids = kwargs.get("filtered_pids")
             assert type(filtered_pids) == List[int], "The filtered pids should be a list of integers"
             device = "cuda" if torch.cuda.is_available() else "cpu"
             results = self.searcher.search(
                 query,
-                #Number of passages to receive
-                k=k, 
-                #Passing the filter function of relevant 
+                # Number of passages to receive
+                k=k,
+                # Passing the filter function of relevant
                 filter_fn=lambda pids: torch.tensor(
-                    [pid for pid in pids if pid in filtered_pids],dtype=torch.int32).to(device))
+                    [pid for pid in pids if pid in filtered_pids], dtype=torch.int32
+                ).to(device),
+            )
         else:
             searcher_results = self.searcher.search(query, k=k)
         results = []
-        for pid,rank,score in zip(*searcher_results):
-            results.append(dotdict({'long_text':self.searcher.collection[pid],'score':score,'pid':pid}))
+        for pid, rank, score in zip(*searcher_results):  # noqa: B007
+            results.append(dotdict({"long_text": self.searcher.collection[pid], "score": score, "pid": pid}))
         return results
 
+
 class ColBERTv2RerankerLocal:
-    
-    def __init__(self,colbert_config=None,checkpoint:str='bert-base-uncased'):
+    def __init__(self, colbert_config=None, checkpoint: str = "bert-base-uncased"):
         try:
-            import colbert # noqa: F401
+            import colbert  # noqa: F401
         except ImportError:
-            print("Colbert not found. Please check your installation or install the module using pip install colbert-ai[faiss-gpu,torch].")
+            print(
+                "Colbert not found. Please check your installation or install the module using pip install "
+                "colbert-ai[faiss-gpu,torch]."
+            )
         """_summary_
 
         Args:
@@ -167,24 +192,25 @@ def __init__(self,colbert_config=None,checkpoint:str='bert-base-uncased'):
     def __call__(self, *args: Any, **kwargs: Any) -> Any:
         return self.forward(*args, **kwargs)
 
-    def forward(self,query:str,passages:List[str]=[]):
+    def forward(self, query: str, passages: Optional[List[str]] = None):
         assert len(passages) > 0, "Passages should not be empty"
 
         import numpy as np
         from colbert.modeling.colbert import ColBERT
         from colbert.modeling.tokenization.doc_tokenization import DocTokenizer
         from colbert.modeling.tokenization.query_tokenization import QueryTokenizer
-        
+
+        passages = passages or []
         self.colbert_config.nway = len(passages)
-        query_tokenizer = QueryTokenizer(self.colbert_config,verbose=1)
+        query_tokenizer = QueryTokenizer(self.colbert_config, verbose=1)
         doc_tokenizer = DocTokenizer(self.colbert_config)
-        query_ids,query_masks = query_tokenizer.tensorize([query])
-        doc_ids,doc_masks = doc_tokenizer.tensorize(passages)
-
-        col = ColBERT(self.checkpoint,self.colbert_config) 
-        Q = col.query(query_ids,query_masks)
-        DOC_IDS,DOC_MASKS = col.doc(doc_ids,doc_masks,keep_dims='return_mask')
-        Q_duplicated = Q.repeat_interleave(len(passages), dim=0).contiguous()
-        tensor_scores = col.score(Q_duplicated,DOC_IDS,DOC_MASKS)
+        query_ids, query_masks = query_tokenizer.tensorize([query])
+        doc_ids, doc_masks = doc_tokenizer.tensorize(passages)
+
+        col = ColBERT(self.checkpoint, self.colbert_config)
+        q = col.query(query_ids, query_masks)
+        doc_ids, doc_masks = col.doc(doc_ids, doc_masks, keep_dims="return_mask")
+        q_duplicated = q.repeat_interleave(len(passages), dim=0).contiguous()
+        tensor_scores = col.score(q_duplicated, doc_ids, doc_masks)
         passage_score_arr = np.array([score.cpu().detach().numpy().tolist() for score in tensor_scores])
-        return passage_score_arr
+        return passage_score_arr