Separate rerankers, splitters and extractors (#85)

knowledgehub/indices/rankings/cohere.py

@@ -0,0 +1,38 @@
+from __future__ import annotations
+
+import os
+
+from ...base import Document
+from .base import BaseReranking
+
+
+class CohereReranking(BaseReranking):
+    model_name: str = "rerank-multilingual-v2.0"
+    cohere_api_key: str = os.environ.get("COHERE_API_KEY", "")
+    top_k: int = 1
+
+    def run(self, documents: list[Document], query: str) -> list[Document]:
+        """Use Cohere Reranker model to re-order documents
+        with their relevance score"""
+        try:
+            import cohere
+        except ImportError:
+            raise ImportError(
+                "Please install Cohere " "`pip install cohere` to use Cohere Reranking"
+            )
+
+        cohere_client = cohere.Client(self.cohere_api_key)
+
+        # output documents
+        compressed_docs = []
+        if len(documents) > 0:  # to avoid empty api call
+            _docs = [d.content for d in documents]
+            results = cohere_client.rerank(
+                model=self.model_name, query=query, documents=_docs, top_n=self.top_k
+            )
+            for r in results:
+                doc = documents[r.index]
+                doc.metadata["relevance_score"] = r.relevance_score
+                compressed_docs.append(doc)
+
+        return compressed_docs