TT-3084-clarifai-dspy-integration

Author: mogith-pn

dsp/modules/__init__.py

@@ -6,6 +6,7 @@
 from .cohere import *
 from .sbert import *
 from .pyserini import *
+from .clarifai import *
 
 from .hf_client import HFClientTGI
 from .hf_client import Anyscale

dsp/modules/clarifai.py

@@ -0,0 +1,88 @@
+import math
+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`.")
+
+class ClarifaiLLM(LM):
+    """Integration to call models hosted in clarifai platform."""
+
+    def __init__(
+        self,
+        model: str = "https://clarifai.com/mistralai/completion/models/mistral-7B-Instruct", #defaults to mistral-7B-Instruct
+        api_key: Optional[str] = None,
+        stop_sequences: list[str] = [],
+        **kwargs,
+    ):
+        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
+        )
+    
+    def basic_request(self, prompt, **kwargs):
+        
+        params = (
+            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
+                
+            )
+
+        history = {
+            "prompt": prompt,
+            "response": response,
+            "kwargs": 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)
+    
+    def __call__(self,
+        prompt: str,
+        only_completed: bool = True,
+        return_sorted: bool = False,
+        **kwargs
+    ):
+        assert only_completed, "for now"
+        assert return_sorted is False, "for now"
+
+        n = kwargs.pop("n", 1)
+        completions=[]
+
+        for i in range(n):
+            response = self.request(prompt, **kwargs)
+            completions.append(response)
+
+        return completions

dsp/modules/lm.py

@@ -45,14 +45,23 @@ def inspect_history(self, n: int = 1, skip: int = 0):
             prompt = x["prompt"]
 
             if prompt != last_prompt:
-                printed.append(
-                    (
-                        prompt,
-                        x["response"].generations
-                        if provider == "cohere"
-                        else x["response"]["choices"],
+
+                if provider=="clarifai":
+                    printed.append(
+                        (
+                            prompt,
+                            x['response']
+                        ) 
+                    )
+                else:    
+                    printed.append(
+                        (
+                            prompt,
+                            x["response"].generations
+                            if provider == "cohere"
+                            else x["response"]["choices"],
+                        )
                     )
-                )
 
             last_prompt = prompt
 
@@ -71,6 +80,8 @@ def inspect_history(self, n: int = 1, skip: int = 0):
                 text = choices[0].text
             elif provider == "openai" or provider == "ollama":
                 text = ' ' + self._get_choice_text(choices[0]).strip()
+            elif provider == "clarifai":
+                text=choices
             else:
                 text = choices[0]["text"]
             self.print_green(text, end="")

dspy/__init__.py

@@ -18,6 +18,7 @@
 OpenAI = dsp.GPT3
 ColBERTv2 = dsp.ColBERTv2
 Pyserini = dsp.PyseriniRetriever
+Clarifai=dsp.ClarifaiLLM
 
 HFClientTGI = dsp.HFClientTGI
 HFClientVLLM = HFClientVLLM

dspy/retrieve/clarifai_rm.py

@@ -0,0 +1,96 @@
+from collections import defaultdict
+from typing import List, Union
+import os
+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
+    from google.protobuf import json_format
+except ImportError:
+    raise ImportError(
+        "Clarifai is not installed. Install it using `pip install clarifai`"
+    )
+
+
+class ClarifaiRM(dspy.Retrieve):
+    """
+    Retrieval module uses clarifai to return the Top K relevant pasages for the given query.
+    """
+
+    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
+
+        super().__init__(k=k)
+    
+    def forward(self, query_or_queries: Union[str, List[str]], k: Optional[int]):
+
+        """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)
+        ```
+        """
+        queries = (
+            [query_or_queries]
+            if isinstance(query_or_queries, str)
+            else query_or_queries
+        )
+        k = k if k is not None else self.k
+        passages = []
+        queries = [q for q in queries if q]
+
+        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
+            
+
+
+
+
+
+
+
+
+