Switch to cohere command r model

app.py

@@ -11,8 +11,7 @@ model_config = ConfigDict(protected_namespaces=())
 setting_keys = gr.Interface(
     fn=set_keys,
     inputs=[
-        #gr.Textbox(label="Enter your CO_API_KEY"),
-        gr.Textbox(label="Enter your OPENAI_API_KEY"),
+        gr.Textbox(label="Enter your CO_API_KEY"),
         gr.Textbox(label="Enter your LLAMA_CLOUD_API_KEY"),
     ],
     outputs=gr.Textbox(label="Status")
@@ -31,7 +30,10 @@ uploading_files = gr.Interface(
 qa = gr.Interface(
     fn=answer_query,
     inputs=gr.Textbox(label="Enter your question"),
-    outputs=gr.Textbox(label="Answer"), 
+    outputs=[
+        gr.Textbox(label="Answer"), 
+        gr.Textbox(label="Relevant Nodes"),
+    ],
     title="Document Q&A System"
 )
 

main.py

@@ -1,13 +1,13 @@
-from utils.retriever import get_query_engine
+from utils.retriever import get_engine
 from utils.index import create_index
 from utils.constant import INDEX_PATH, DATA_PATH
 import os
 import shutil
 
-def set_keys(openai_api_key: str, llama_cloud_api_key: str) -> str:
+
+def set_keys(co_api_key: str, llama_cloud_api_key: str) -> str:
     try:
-        #os.environ["CO_API_KEY"] = co_api_key
-        os.environ["OPENAI_API_KEY"] = openai_api_key
+        os.environ["CO_API_KEY"] = co_api_key
         os.environ["LLAMA_CLOUD_API_KEY"] = llama_cloud_api_key
         return "Keys are set successfully"
     
@@ -23,6 +23,10 @@ def handle_file(uploaded_file):
         shutil.move(temp_file_path, save_file_path)
         
         create_index(DATA_PATH, INDEX_PATH)
+        
+        engine_manager = get_engine()[1]
+        engine_manager.initialize_index()
+        
         return "File uploaded and indexed"
     
     except Exception as e:
@@ -30,19 +34,20 @@ def handle_file(uploaded_file):
 
 def answer_query(query: str) -> str:
     
-    query_engine = get_query_engine(semantic=True)
+    query_engine = get_engine(semantic=True)[0]
     response = query_engine.query(query)
     
     nodes = query_engine.retriever.retrieve(query)
+    nodes_str = ""
     
     for node in nodes:
 
         score = node.get_score()
         text = node.text
         
-        response += f"\nNode: {node.node_id}\nScore: {score:0.3f}\nText: {text[:1000]}\n"
+        nodes_str += f"\nNode: {node.node_id}\nScore: {score:0.3f}\nText: {text[:1000]}\n"
     
-    return response
+    return response, nodes_str
 
 if __name__ == "__main__":
     

utils/retriever.py

@@ -25,14 +25,18 @@ class QueryEngineManager:
             self.reranker = None
             self.query_engine = None
             self._configure()
-    
-    def _configure(self):
-        configure_settings()
+        
+    def initialize_index(self):
         self.index = load_index(path=INDEX_PATH)
         self.nodes = list(self.index.docstore.docs.values())
+        
+    def _configure(self):
+        configure_settings()
+        self.initialize_index()
         self.reranker = LLMRerank(top_n=TOP_N_RERANKER)
         
     def get_engine(self, bm25: bool = False, semantic: bool = False):
+        
         if bm25:
             self.retriever = BM25Retriever.from_defaults(
                 nodes=self.nodes,
@@ -59,6 +63,6 @@ class QueryEngineManager:
         
         return self.query_engine
 
-def get_query_engine(bm25: bool = False, semantic: bool = False):
+def get_engine(bm25: bool = False, semantic: bool = False):
     engine_manager = QueryEngineManager()
-    return engine_manager.get_engine(bm25, semantic)
+    return engine_manager.get_engine(bm25, semantic), engine_manager

utils/settings.py

@@ -1,6 +1,7 @@
 from llama_index.core import Settings
-#from llama_index.llms.cohere import Cohere
-from llama_index.llms.openai import OpenAI
+from llama_index.llms.cohere import Cohere
+#from llama_index.llms.anthropic import Anthropic
+#from llama_index.llms.openai import OpenAI
 from llama_index.embeddings.huggingface import HuggingFaceEmbedding
 from llama_index.core.node_parser import SemanticSplitterNodeParser
 
@@ -16,15 +17,20 @@ def load_llm():
     - Do not include information from external sources not provided by the user.
     '''
     
-    # llm = Cohere(
-    #     system_prompt=system_prompt,
+    llm = Cohere(
+        system_prompt=system_prompt,
         
+    )
+    
+    # llm = OpenAI(
+    #     model = "gpt-3.5-turbo",
+    #     system_prompt=system_prompt
+    # )
+    # llm = Anthropic(
+    #     model="claude-3-opus-20240229",
+    #     system_prompt=system_prompt
     # )
     
-    llm = OpenAI(
-        model = "gpt-4o-mini",
-        system_prompt=system_prompt
-    )
     
     return llm 
 
@@ -43,6 +49,7 @@ def configure_settings():
     llm_replicate = load_llm()
     embed_model = load_embed_model()
     
+    #Settings.tokenizer = Anthropic().tokenizer
     Settings.llm = llm_replicate
     Settings.embed_model = embed_model
     Settings.node_parser = SemanticSplitterNodeParser(