Merge branch 'main' of https://huggingface.co/spaces/Tonic/01aiYi-NvidiaEmbed

app.py

@@ -1,4 +1,4 @@
-# app.py  
+# app.py
 import spaces
 from torch.nn import DataParallel  
 from torch import Tensor  
@@ -10,19 +10,16 @@ from langchain_chroma import Chroma
 from chromadb import Documents, EmbeddingFunction, Embeddings  
 from chromadb.config import Settings  
 import chromadb #import HttpClient 
-import os  
+import os
+import tempfile
 import re 
 import uuid  
 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()
 
@@ -31,6 +28,13 @@ os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
 os.environ['CUDA_CACHE_DISABLE'] = '1'  
   
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")  
+
+# Ensure the temporary directory exists
+temp_dir = '/tmp/gradio/'
+os.makedirs(temp_dir, exist_ok=True)
+
+# Set Gradio cache directory
+gr.components.file.GRADIO_CACHE = temp_dir
   
 ### Utils  
 hf_token, yi_token = load_env_variables()  
@@ -110,7 +114,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 +144,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 +154,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 +164,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 +203,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 +230,4 @@ with gr.Blocks() as demo:
   
 if __name__ == "__main__":  
     # os.system("chroma run --host localhost --port 8000 &")
-    demo.launch()  
+    demo.launch()