Update app.py

app.py

@@ -1,4 +1,4 @@
-# app.py  
+# app.py
 import spaces
 from torch.nn import DataParallel  
 from torch import Tensor  
@@ -17,12 +17,8 @@ import gradio as gr
 import torch  
 import torch.nn.functional as F  
 from dotenv import load_dotenv
-from utils import load_env_variables, parse_and_route , escape_special_characters 
-from globalvars import API_BASE, intention_prompt, tasks, system_message, model_name , metadata_prompt 
-# import time  
-# import httpx  
-  
-
+from utils import load_env_variables, parse_and_route, escape_special_characters 
+from globalvars import API_BASE, intention_prompt, tasks, system_message, model_name, metadata_prompt 
 
 load_dotenv()
 
@@ -110,7 +106,7 @@ class EmbeddingGenerator:
         matches = pattern.findall(metadata_output)  
         metadata = {key: value for key, value in matches}  
         return metadata 
-  
+
 class MyEmbeddingFunction(EmbeddingFunction):  
     def __init__(self, model_name: str, token: str, intention_client):  
         self.model_name = model_name  
@@ -140,7 +136,7 @@ def initialize_chroma(collection_name: str, embedding_function: MyEmbeddingFunct
   
 def add_documents_to_chroma(documents: list, embedding_function: MyEmbeddingFunction):  
     for doc in documents:  
-        embeddings, metadata = embedding_function.embedding_generator.compute_embeddings(doc)  
+        embeddings, metadata = embedding_function.create_embedding_generator().compute_embeddings(doc)  
         for embedding, meta in zip(embeddings, metadata):  
             chroma_collection.add(  
                 ids=[str(uuid.uuid1())],  
@@ -150,7 +146,7 @@ def add_documents_to_chroma(documents: list, embedding_function: MyEmbeddingFunc
             )  
   
 def query_chroma(query_text: str, embedding_function: MyEmbeddingFunction):  
-    query_embeddings, query_metadata = embedding_function.embedding_generator.compute_embeddings(query_text)  
+    query_embeddings, query_metadata = embedding_function.create_embedding_generator().compute_embeddings(query_text)  
     result_docs = chroma_collection.query(  
         query_texts=[query_text],  
         n_results=2  
@@ -160,7 +156,7 @@ def query_chroma(query_text: str, embedding_function: MyEmbeddingFunction):
 # 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)  
+embedding_function = MyEmbeddingFunction(model_name=model_name, token=hf_token, intention_client=intention_client)  
 chroma_db = initialize_chroma(collection_name="Tonic-instruct", embedding_function=embedding_function)
 
 def respond(  
@@ -199,7 +195,7 @@ def upload_documents(files):
     return "Documents uploaded and processed successfully!"  
   
 def query_documents(query):  
-    results = query_chroma(query)  
+    results = query_chroma(query, embedding_function)  
     return "\n\n".join([result.content for result in results])  
   
 with gr.Blocks() as demo:  
@@ -226,4 +222,4 @@ with gr.Blocks() as demo:
   
 if __name__ == "__main__":  
     # os.system("chroma run --host localhost --port 8000 &")
-    demo.launch()  
+    demo.launch()