addressed comments

Author: mogith-pn

dsp/modules/clarifai.py

@@ -1,22 +1,30 @@
-import math
+"""Clarifai LM integration"""
 from typing import Any, Optional
-import backoff
-
 from dsp.modules.lm import LM
 
 try:
     from clarifai.client.model import Model
-except ImportError:
-    raise ImportError("ClarifaiLLM requires `pip install clarifai`.")
+except ImportError as err:
+    raise ImportError("ClarifaiLLM requires `pip install clarifai`.") from err
 
 class ClarifaiLLM(LM):
-    """Integration to call models hosted in clarifai platform."""
+    """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, 
+                                        inference_params={"max_tokens":100,'temperature':0.6}))
+    """
 
     def __init__(
         self,
-        model: str = "https://clarifai.com/mistralai/completion/models/mistral-7B-Instruct", #defaults to mistral-7B-Instruct
+        model: str = "https://clarifai.com/mistralai/completion/models/mistral-7B-Instruct",
         api_key: Optional[str] = None,
-        stop_sequences: list[str] = [],
         **kwargs,
     ):
         super().__init__(model)
@@ -41,7 +49,6 @@ def __init__(
         )
 
     def basic_request(self, prompt, **kwargs):
-        
         params = (
             self.kwargs['inference_params'] if 'inference_params' in self.kwargs
             else {}
@@ -50,11 +57,10 @@ def basic_request(self, prompt, **kwargs):
                 self._model.predict_by_bytes(
                     input_bytes= prompt.encode(encoding="utf-8"),
                     input_type= "text",
-                    inference_params= params
-                ).outputs[0].data.text.raw
-                
+                    inference_params= params,
+                    ).outputs[0].data.text.raw
             )
-
+        kwargs = {**self.kwargs, **kwargs}
         history = {
             "prompt": prompt,
             "response": response,
@@ -63,9 +69,6 @@ def basic_request(self, prompt, **kwargs):
         self.history.append(history)
         return response
 
-    def _get_choice_text(self, choice: dict[str, Any]) -> str:
-        return choice
-    
     def request(self, prompt: str, **kwargs):
         return self.basic_request(prompt, **kwargs)
 

dspy/retrieve/clarifai_rm.py

@@ -1,4 +1,4 @@
-from collections import defaultdict
+"""Clarifai as retriver to retrieve hits"""
 from typing import List, Union
 import os
 import dspy
@@ -9,16 +9,25 @@
 
 try:
     from clarifai.client.search import Search
-    from google.protobuf import json_format
-except ImportError:
+except ImportError as err:
     raise ImportError(
         "Clarifai is not installed. Install it using `pip install clarifai`"
-    )
+    ) from err
 
 
 class ClarifaiRM(dspy.Retrieve):
     """
     Retrieval module uses clarifai to return the Top K relevant pasages for the given query.
+    Assuming that you have ingested the source documents into clarifai App, where it is indexed and stored.
+
+    Args:
+        clarifai_user_id (str): Clarifai unique user_id.
+        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,
@@ -32,10 +41,18 @@ def __init__(self,
         self.user_id = clarifai_user_id
         self.pat = clarifai_pat if clarifai_pat is not None else os.environ["CLARIFAI_PAT"]
         self.k=k
-
+        
         super().__init__(k=k)
 
-    def forward(self, query_or_queries: Union[str, List[str]], k: Optional[int]):
+    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:
 
         """Uses clarifai-python SDK search function and retrieves top_k similar passages for given query,
        Args:
@@ -49,7 +66,6 @@ def forward(self, query_or_queries: Union[str, List[str]], k: Optional[int]):
        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)
@@ -60,37 +76,19 @@ def forward(self, query_or_queries: Union[str, List[str]], k: Optional[int]):
             if isinstance(query_or_queries, str)
             else query_or_queries
         )
-        k = k if k is not None else self.k
+
+        k = self.k if self.k is not None else k
         passages = []
         queries = [q for q in queries if q]
+        clarifai_search = Search(user_id=self.user_id, app_id=self.app_id, top_k=k, pat=self.pat)
 
         for query in queries:
-            clarifai_search = Search(user_id=self.user_id, app_id=self.app_id, top_k=k, pat=self.pat)
             search_response= clarifai_search.query(ranks=[{"text_raw": query}])
 
             # Retrieve hits
             hits=[hit for data in search_response for hit in data.hits]
-            executor = ThreadPoolExecutor(max_workers=10)
-
-            def retrieve_hits(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
-            
-            futures = [executor.submit(retrieve_hits, hit) for hit in hits]
-            results = [future.result() for future in futures]
-            passages=[dotdict({"long_text": d}) for d in results]
-
-            return passages
-            
-
-
+            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)
 
-
-
-
-
-
-    
+        return passages