Add Cohere, NVIDIA, Jina & HuggingFace rerankers (#37)

README.md

@@ -94,10 +94,22 @@ pip install git+https://github.com/Storia-AI/sage.git@main
 
 <br>
 <summary><strong>Optional</strong></summary>
-If you are planning on indexing GitHub issues in addition to the codebase, you will need a GitHub token:
+
+- By default, we use an <a href="https://huggingface.co/cross-encoder/ms-marco-MiniLM-L-6-v2">open-source re-ranker</a>. For higher accuracy, you can use <a href="https://cohere.com/rerank">Cohere</a>, <a href="https://developer.nvidia.com/blog/enhancing-rag-pipelines-with-re-ranking/">NVIDIA</a> or <a href="https://jina.ai/reranker/">Jina</a>:
+
+    ```
+    export COHERE_API_KEY=...
+    export NVIDIA_API_KEY=...
+    export JINA_API_KEY=...
+    ```
+
+    We are seeing significant gains in accuracy from these proprietary rerankers.
+
+- If you are planning on indexing GitHub issues in addition to the codebase, you will need a GitHub token:
 
     export GITHUB_TOKEN=...
 
+
 ## Running it
 
 <details open>

requirements.txt

@@ -3,15 +3,16 @@ Pygments==2.18.0
 cohere==5.9.2
 fastapi==0.112.2
 gradio>=4.26.0
-langchain==0.3.0
-langchain-anthropic==0.2.0
-langchain-cohere==0.3.0
-langchain-community==0.3.0
-langchain-core==0.3.0
-langchain-experimental==0.3.0
-langchain-ollama==0.2.0
-langchain-openai==0.2.0
-langchain-text-splitters==0.3.0
+langchain==0.2.16
+langchain-anthropic==0.1.23
+langchain-cohere==0.2.4
+langchain-community==0.2.17
+langchain-core==0.2.40
+langchain-experimental==0.0.65
+langchain-nvidia-ai-endpoints==0.2.2
+langchain-ollama==0.1.3
+langchain-openai==0.1.25
+langchain-text-splitters==0.2.4
 marqo==3.7.0
 nbformat==5.10.4
 openai==1.42.0

sage/chat.py

@@ -12,14 +12,12 @@ from dotenv import load_dotenv
 from langchain.chains import create_history_aware_retriever, create_retrieval_chain
 from langchain.chains.combine_documents import create_stuff_documents_chain
 from langchain.retrievers import ContextualCompressionRetriever
-from langchain.retrievers.document_compressors import CrossEncoderReranker
 from langchain.schema import AIMessage, HumanMessage
-from langchain_cohere import CohereRerank
-from langchain_community.cross_encoders import HuggingFaceCrossEncoder
 from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
 
 import sage.vector_store as vector_store
 from sage.llm import build_llm_via_langchain
+from sage.reranker import build_reranker, RerankerProvider
 
 load_dotenv()
 
@@ -30,15 +28,7 @@ def build_rag_chain(args):
 
     retriever_top_k = 5 if args.reranker_provider == "none" else 25
     retriever = vector_store.build_from_args(args).as_retriever(top_k=retriever_top_k)
-
-    if args.reranker_provider == "none":
-        compressor = None
-    if args.reranker_provider == "huggingface":
-        encoder_model = HuggingFaceCrossEncoder(model_name=args.reranker_model)
-        compressor = CrossEncoderReranker(model=encoder_model, top_n=5)
-    if args.reranker_provider == "cohere":
-        compressor = CohereRerank(model=args.reranker_model, cohere_api_key=os.environ.get("COHERE_API_KEY"), top_n=5)
-
+    compressor = build_reranker(args.reranker_provider, args.reranker_model)
     if compressor:
         retriever = ContextualCompressionRetriever(base_compressor=compressor, base_retriever=retriever)
 
@@ -94,7 +84,7 @@ def main():
         default="http://localhost:8882",
         help="URL for the Marqo server. Required if using Marqo as embedder or vector store.",
     )
-    parser.add_argument("--reranker-provider", default="huggingface", choices=["none", "huggingface", "cohere"])
+    parser.add_argument("--reranker-provider", default="huggingface", choices=[r.value for r in RerankerProvider])
     parser.add_argument(
         "--reranker-model",
         help="The reranker model name. When --reranker-provider=huggingface, we suggest choosing a model from the "
@@ -114,12 +104,6 @@ def main():
     )
     args = parser.parse_args()
 
-    if not args.index_name:
-        if args.vector_store_type == "marqo":
-            args.index_name = args.repo_id.split("/")[1]
-        elif args.vector_store_type == "pinecone":
-            parser.error("Please specify --index-name for Pinecone.")
-
     if not args.llm_model:
         if args.llm_provider == "openai":
             args.llm_model = "gpt-4"
@@ -130,12 +114,6 @@ def main():
         else:
             raise ValueError("Please specify --llm_model")
 
-    if not args.reranker_model:
-        if args.reranker_provider == "cohere":
-            args.reranker_model = "rerank-english-v3.0"
-        elif args.reranker_provider == "huggingface":
-            args.reranker_model = "cross-encoder/ms-marco-TinyBERT-L-2-v2"
-
     rag_chain = build_rag_chain(args)
 
     def source_md(file_path: str, url: str) -> str:

sage/reranker.py

@@ -0,0 +1,42 @@
+import os
+from enum import Enum
+from typing import Optional
+
+from langchain.retrievers.document_compressors import CrossEncoderReranker
+from langchain_cohere import CohereRerank
+from langchain_community.cross_encoders import HuggingFaceCrossEncoder
+from langchain_community.document_compressors import JinaRerank
+from langchain_core.documents import BaseDocumentCompressor
+from langchain_nvidia_ai_endpoints import NVIDIARerank
+
+
+class RerankerProvider(Enum):
+    NONE = "none"
+    HUGGINGFACE = "huggingface"
+    COHERE = "cohere"
+    NVIDIA = "nvidia"
+    JINA = "jina"
+
+
+def build_reranker(provider: str, model: Optional[str] = None, top_n: Optional[int] = 5) -> BaseDocumentCompressor:
+    if provider == RerankerProvider.NONE.value:
+        return None
+    if provider == RerankerProvider.HUGGINGFACE.value:
+        model = model or "cross-encoder/ms-marco-MiniLM-L-6-v2"
+        encoder_model = HuggingFaceCrossEncoder(model_name=model)
+        return CrossEncoderReranker(model=encoder_model, top_n=top_n)
+    if provider == RerankerProvider.COHERE.value:
+        if not os.environ.get("COHERE_API_KEY"):
+            raise ValueError("Please set the COHERE_API_KEY environment variable")
+        model = model or "rerank-english-v3.0"
+        return CohereRerank(model=model, cohere_api_key=os.environ.get("COHERE_API_KEY"), top_n=top_n)
+    if provider == RerankerProvider.NVIDIA.value:
+        if not os.environ.get("NVIDIA_API_KEY"):
+            raise ValueError("Please set the NVIDIA_API_KEY environment variable")
+        model = model or "nvidia/nv-rerankqa-mistral-4b-v3"
+        return NVIDIARerank(model=model, api_key=os.environ.get("NVIDIA_API_KEY"), top_n=top_n, truncate="END")
+    if provider == RerankerProvider.JINA.value:
+        if not os.environ.get("JINA_API_KEY"):
+            raise ValueError("Please set the JINA_API_KEY environment variable")
+        return JinaRerank(top_n=top_n)
+    raise ValueError(f"Invalid reranker provider: {provider}")

sage/vector_store.py

@@ -149,11 +149,15 @@ class MarqoVectorStore(VectorStore):
 def build_from_args(args: dict) -> VectorStore:
     """Builds a vector store from the given command-line arguments."""
     if args.vector_store_type == "pinecone":
+        if not args.index_name:
+            raise ValueError("Please specify --index-name for Pinecone.")
         dimension = args.embedding_size if "embedding_size" in args else None
         return PineconeVectorStore(
             index_name=args.index_name, namespace=args.repo_id, dimension=dimension, hybrid=args.hybrid_retrieval
         )
     elif args.vector_store_type == "marqo":
-        return MarqoVectorStore(url=args.marqo_url, index_name=args.index_name)
+        marqo_url = args.marqo_url or "http://localhost:8882"
+        index_name = args.index_name or args.repo_id.split("/")[1]
+        return MarqoVectorStore(url=marqo_url, index_name=index_name)
     else:
         raise ValueError(f"Unrecognized vector store type {args.vector_store_type}")