addressed linting

Author: mogith-pn

dsp/modules/clarifai.py

@@ -1,23 +1,25 @@
 """Clarifai LM integration"""
 from typing import Any, Optional
+
 from dsp.modules.lm import LM
 
 try:
     from clarifai.client.model import Model
 except ImportError as err:
     raise ImportError("ClarifaiLLM requires `pip install clarifai`.") from err
 
+
 class ClarifaiLLM(LM):
     """Integration to call models hosted in clarifai platform.
-    
+
     Args:
         model (str, optional): Clarifai URL of the model. Defaults to "Mistral-7B-Instruct".
         api_key (Optional[str], optional): CLARIFAI_PAT token. Defaults to None.
         **kwargs: Additional arguments to pass to the API provider.
     Example:
         import dspy
-        dspy.configure(lm=dspy.Clarifai(model=MODEL_URL, 
-                                        api_key=CLARIFAI_PAT, 
+        dspy.configure(lm=dspy.Clarifai(model=MODEL_URL,
+                                        api_key=CLARIFAI_PAT,
                                         inference_params={"max_tokens":100,'temperature':0.6}))
     """
 
@@ -30,36 +32,36 @@ def __init__(
         super().__init__(model)
 
         self.provider = "clarifai"
-        self.pat=api_key
-        self._model= Model(url=model, pat=api_key)
-        self.kwargs = {
-            "n": 1,
-            **kwargs
-        }
-        self.history :list[dict[str, Any]] = []
-        self.kwargs['temperature'] = (
-            self.kwargs['inference_params']['temperature'] if
-            'inference_params' in self.kwargs and
-             'temperature' in self.kwargs['inference_params'] else 0.0
-        ) 
-        self.kwargs['max_tokens'] = (
-            self.kwargs['inference_params']['max_tokens'] if
-            'inference_params' in self.kwargs and
-             'max_tokens' in self.kwargs['inference_params'] else 150
+        self.pat = api_key
+        self._model = Model(url=model, pat=api_key)
+        self.kwargs = {"n": 1, **kwargs}
+        self.history: list[dict[str, Any]] = []
+        self.kwargs["temperature"] = (
+            self.kwargs["inference_params"]["temperature"]
+            if "inference_params" in self.kwargs
+            and "temperature" in self.kwargs["inference_params"]
+            else 0.0
+        )
+        self.kwargs["max_tokens"] = (
+            self.kwargs["inference_params"]["max_tokens"]
+            if "inference_params" in self.kwargs
+            and "max_tokens" in self.kwargs["inference_params"]
+            else 150
         )
-    
+
     def basic_request(self, prompt, **kwargs):
         params = (
-            self.kwargs['inference_params'] if 'inference_params' in self.kwargs
-            else {}
+            self.kwargs["inference_params"] if "inference_params" in self.kwargs else {}
         )
         response = (
-                self._model.predict_by_bytes(
-                    input_bytes= prompt.encode(encoding="utf-8"),
-                    input_type= "text",
-                    inference_params= params,
-                    ).outputs[0].data.text.raw
+            self._model.predict_by_bytes(
+                input_bytes=prompt.encode(encoding="utf-8"),
+                input_type="text",
+                inference_params=params,
             )
+            .outputs[0]
+            .data.text.raw
+        )
         kwargs = {**self.kwargs, **kwargs}
         history = {
             "prompt": prompt,
@@ -68,21 +70,22 @@ def basic_request(self, prompt, **kwargs):
         }
         self.history.append(history)
         return response
-    
+
     def request(self, prompt: str, **kwargs):
         return self.basic_request(prompt, **kwargs)
-    
-    def __call__(self,
+
+    def __call__(
+        self,
         prompt: str,
         only_completed: bool = True,
         return_sorted: bool = False,
-        **kwargs
+        **kwargs,
     ):
         assert only_completed, "for now"
         assert return_sorted is False, "for now"
 
         n = kwargs.pop("n", 1)
-        completions=[]
+        completions = []
 
         for i in range(n):
             response = self.request(prompt, **kwargs)

dspy/retrieve/clarifai_rm.py

@@ -1,11 +1,12 @@
 """Clarifai as retriver to retrieve hits"""
-from typing import List, Union
 import os
+from concurrent.futures import ThreadPoolExecutor
+from typing import List, Optional, Union
+
+import requests
+
 import dspy
 from dsp.utils import dotdict
-import requests
-from typing import Optional
-from concurrent.futures import ThreadPoolExecutor
 
 try:
     from clarifai.client.search import Search
@@ -25,51 +26,55 @@ class ClarifaiRM(dspy.Retrieve):
         clarfiai_app_id (str): Clarifai App ID, where the documents are stored.
         clarifai_pat (str): Clarifai PAT key.
         k (int): Top K documents to retrieve.
-    
+
     Examples:
         TODO
     """
 
-    def __init__(self,
-                 clarifai_user_id: str,
-                 clarfiai_app_id: str,
-                 clarifai_pat: Optional[str] = None,
-                 k: int = 3,
-                 
+    def __init__(
+        self,
+        clarifai_user_id: str,
+        clarfiai_app_id: str,
+        clarifai_pat: Optional[str] = None,
+        k: int = 3,
     ):
         self.app_id = clarfiai_app_id
         self.user_id = clarifai_user_id
-        self.pat = clarifai_pat if clarifai_pat is not None else os.environ["CLARIFAI_PAT"]
-        self.k=k
-        self.clarifai_search = Search(user_id=self.user_id, app_id=self.app_id, top_k=k, pat=self.pat)
+        self.pat = (
+            clarifai_pat if clarifai_pat is not None else os.environ["CLARIFAI_PAT"]
+        )
+        self.k = k
+        self.clarifai_search = Search(
+            user_id=self.user_id, app_id=self.app_id, top_k=k, pat=self.pat
+        )
         super().__init__(k=k)
-    
+
     def retrieve_hits(self, hits):
-                header = {"Authorization": f"Key {self.pat}"}
-                request = requests.get(hits.input.data.text.url, headers=header)
-                request.encoding = request.apparent_encoding
-                requested_text = request.text
-                return requested_text
-    
-    def forward(self, query_or_queries: Union[str, List[str]], k: Optional[int] = None
-    ) -> dspy.Prediction:
+        header = {"Authorization": f"Key {self.pat}"}
+        request = requests.get(hits.input.data.text.url, headers=header)
+        request.encoding = request.apparent_encoding
+        requested_text = request.text
+        return requested_text
 
+    def forward(
+        self, query_or_queries: Union[str, List[str]], k: Optional[int] = None
+    ) -> dspy.Prediction:
         """Uses clarifai-python SDK search function and retrieves top_k similar passages for given query,
-       Args:
-            query_or_queries : single query or list of queries
-            k : Top K relevant documents to return
-       
-       Returns:
-            passages in format of dotdict 
-       
-       Examples:
-       Below is a code snippet that shows how to use Marqo as the default retriver:
-        ```python
-        import clarifai
-        llm = dspy.Clarifai(model=MODEL_URL, api_key="YOUR CLARIFAI_PAT")
-        retriever_model = ClarifaiRM(clarifai_user_id="USER_ID", clarfiai_app_id="APP_ID", clarifai_pat="YOUR CLARIFAI_PAT")
-        dspy.settings.configure(lm=llm, rm=retriever_model)
-        ```
+        Args:
+             query_or_queries : single query or list of queries
+             k : Top K relevant documents to return
+
+        Returns:
+             passages in format of dotdict
+
+        Examples:
+        Below is a code snippet that shows how to use Marqo as the default retriver:
+         ```python
+         import clarifai
+         llm = dspy.Clarifai(model=MODEL_URL, api_key="YOUR CLARIFAI_PAT")
+         retriever_model = ClarifaiRM(clarifai_user_id="USER_ID", clarfiai_app_id="APP_ID", clarifai_pat="YOUR CLARIFAI_PAT")
+         dspy.settings.configure(lm=llm, rm=retriever_model)
+         ```
         """
         queries = (
             [query_or_queries]
@@ -81,10 +86,10 @@ def forward(self, query_or_queries: Union[str, List[str]], k: Optional[int] = No
         queries = [q for q in queries if q]
 
         for query in queries:
-            search_response= self.clarifai_search.query(ranks=[{"text_raw": query}])
+            search_response = self.clarifai_search.query(ranks=[{"text_raw": query}])
 
             # Retrieve hits
-            hits=[hit for data in search_response for hit in data.hits]
+            hits = [hit for data in search_response for hit in data.hits]
             with ThreadPoolExecutor(max_workers=10) as executor:
                 results = list(executor.map(self.retrieve_hits, hits))
             passages.extend(dotdict({"long_text": d}) for d in results)