improve retrieval and application logic

app.py

@@ -14,7 +14,7 @@ import gradio as gr
 from huggingface_hub import InferenceClient
 import openai
 from openai import OpenAI
-from globalvars import API_BASE, intention_prompt, tasks
+from globalvars import API_BASE, intention_prompt, tasks , system_message, model_name
 from dotenv import load_dotenv
 import re 
 from utils import load_env_variables
@@ -30,26 +30,17 @@ os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
 os.environ['CUDA_CACHE_DISABLE'] = '1'
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-### Utils
-
 hf_token, yi_token = load_env_variables()
 
 def clear_cuda_cache():
     torch.cuda.empty_cache()
 
-
-## 01ai Yi-large Clience 
-
 client = OpenAI(
     api_key=yi_token,
     base_url=API_BASE
 )
 
 
-## use instruct embeddings
-
-# Load the tokenizer and model
-
 class EmbeddingGenerator:  
     def __init__(self, model_name: str, token: str, intention_client):  
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")  
@@ -108,7 +99,6 @@ class MyEmbeddingFunction(EmbeddingFunction):
         embeddings = [item for sublist in embeddings for item in sublist]  
         return embeddings  
 
-## add chroma vector store
 class DocumentLoader:  
     def __init__(self, file_path: str, mode: str = "elements"):  
         self.file_path = file_path  
@@ -136,60 +126,73 @@ class ChromaManager:
         return result_docs  
   
 
-
-# print(completion)
-
-def respond(
-    message,
-    history: list[tuple[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-):
-    messages = [{"role": "system", "content": system_message}]
-
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
-
-    messages.append({"role": "user", "content": message})
-
-    response = ""
-
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        token = message.choices[0].delta.content
-
-        response += token
-        yield response
-
-"""
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
-"""
-demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
-    ],
-)
-
-
-if __name__ == "__main__":
-    demo.launch()
+# Initialize clients  
+intention_client = OpenAI(api_key=yi_token, base_url=API_BASE)  
+embedding_generator = EmbeddingGenerator(model_name=model_name, token=hf_token, intention_client=intention_client)  
+embedding_function = MyEmbeddingFunction(embedding_generator=embedding_generator)  
+chroma_manager = ChromaManager(embedding_function=embedding_function)  
+  
+def respond(  
+    message,  
+    history: list[tuple[str, str]],  
+    system_message,  
+    max_tokens,  
+    temperature,  
+    top_p,  
+):  
+    retrieved_text = query_documents(message)  
+    messages = [{"role": "system", "content": system_message}]  
+    for val in history:  
+        if val[0]:  
+            messages.append({"role": "user", "content": val[0]})  
+        if val[1]:  
+            messages.append({"role": "assistant", "content": val[1]})  
+    messages.append({"role": "user", "content": f"{retrieved_text}\n\n{message}"})  
+    response = ""  
+    for message in intention_client.chat_completion(  
+        messages,  
+        max_tokens=max_tokens,  
+        stream=True,  
+        temperature=temperature,  
+        top_p=top_p,  
+    ):  
+        token = message.choices[0].delta.content  
+        response += token  
+        yield response  
+  
+def upload_documents(files):  
+    for file in files:  
+        loader = DocumentLoader(file.name)  
+        documents = loader.load_documents()  
+        chroma_manager.add_documents(documents)  
+    return "Documents uploaded and processed successfully!"  
+  
+def query_documents(query):  
+    results = chroma_manager.query(query)  
+    return "\n\n".join([result.content for result in results])  
+  
+with gr.Blocks() as demo:  
+    with gr.Tab("Upload Documents"):  
+        with gr.Row():  
+            document_upload = gr.File(file_count="multiple", file_types=["document"])  
+            upload_button = gr.Button("Upload and Process")  
+            upload_button.click(upload_documents, inputs=document_upload, outputs=gr.Text())  
+  
+    with gr.Tab("Ask Questions"):  
+        with gr.Row():  
+            chat_interface = gr.ChatInterface(  
+                respond,  
+                additional_inputs=[  
+                    gr.Textbox(value="You are a friendly Chatbot.", label="System message"),  
+                    gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),  
+                    gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),  
+                    gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)"),  
+                ],  
+            )  
+            query_input = gr.Textbox(label="Query")  
+            query_button = gr.Button("Query")  
+            query_output = gr.Textbox()  
+            query_button.click(query_documents, inputs=query_input, outputs=query_output)  
+  
+if __name__ == "__main__":  
+    demo.launch()  

globalvars.py

@@ -3,6 +3,8 @@
 API_BASE = "https://api.01.ai/v1"
 API_KEY = "your key"
 
+model_name = 'nvidia/NV-Embed-v1'
+
 title = """
 # 👋🏻Welcome to 🙋🏻‍♂️Tonic's 📽️Nvidia 🛌🏻Embed V-1 !"""
 
@@ -84,3 +86,5 @@ intention_prompt= """
 produce a complete json schema."
 
 you will recieve a text , classify the text according to the schema above. ONLY PROVIDE THE FINAL JSON , DO NOT PRODUCE ANY ADDITION INSTRUCTION :"""
+
+system_message = """ You are a helpful assistant named YiTonic . answer the question provided based on the context above. Produce a complete answer:"""