refactor vectordb

app.py

@@ -98,40 +98,121 @@ class MyEmbeddingFunction(EmbeddingFunction):
         embeddings = [self.embedding_generator.compute_embeddings(doc) for doc in input]  
         embeddings = [item for sublist in embeddings for item in sublist]  
         return embeddings  
-
-class DocumentLoader:  
-    def __init__(self, file_path: str, mode: str = "elements"):  
-        self.file_path = file_path  
-        self.mode = mode  
-  
-    def load_documents(self):  
-        loader = UnstructuredFileLoader(self.file_path, mode=self.mode)  
-        docs = loader.load()  
-        return [doc.page_content for doc in docs]  
-
-class ChromaManager:  
-    def __init__(self, collection_name: str, embedding_function: MyEmbeddingFunction):  
-        self.client = HttpClient(settings=Settings(allow_reset=True))  
-        self.client.reset()  # resets the database  
-        self.collection = self.client.create_collection(collection_name)  
-        self.embedding_function = embedding_function  
-  
-    def add_documents(self, documents: list):  
-        for doc in documents:  
-            self.collection.add(ids=[str(uuid.uuid1())], documents=[doc], embeddings=self.embedding_function([doc]))  
-  
-    def query(self, query_text: str):  
-        db = Chroma(client=self.client, collection_name=self.collection.name, embedding_function=self.embedding_function)  
-        result_docs = db.similarity_search(query_text)  
-        return result_docs  
+# main.py  
+import os  
+import uuid  
+import gradio as gr  
+import torch  
+import torch.nn.functional as F  
+from torch.nn import DataParallel  
+from torch import Tensor  
+from transformers import AutoTokenizer, AutoModel  
+from huggingface_hub import InferenceClient  
+from openai import OpenAI  
+from langchain_community.document_loaders import UnstructuredFileLoader  
+from chromadb import Documents, EmbeddingFunction, Embeddings  
+from chromadb.config import Settings  
+from chromadb import HttpClient  
+from utils import load_env_variables, parse_and_route  
+from globalvars import API_BASE, intention_prompt, tasks, system_message, model_name  
+  
+os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:30'  
+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()  
+  
+client = OpenAI(api_key=yi_token, base_url=API_BASE)  
+  
+class EmbeddingGenerator:  
+    def __init__(self, model_name: str, token: str, intention_client):  
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")  
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name, token=token, trust_remote_code=True)  
+        self.model = AutoModel.from_pretrained(model_name, token=token, trust_remote_code=True).to(self.device)  
+        self.intention_client = intention_client  
+  
+    def clear_cuda_cache(self):  
+        torch.cuda.empty_cache()  
+  
+    @spaces.GPU  
+    def compute_embeddings(self, input_text: str):  
+        # Get the intention  
+        intention_completion = self.intention_client.chat.completions.create(  
+            model="yi-large",  
+            messages=[  
+                {"role": "system", "content": intention_prompt},  
+                {"role": "user", "content": input_text}  
+            ]  
+        )  
+        intention_output = intention_completion.choices[0].message['content']  
+  
+        # Parse and route the intention  
+        parsed_task = parse_and_route(intention_output)  
+        selected_task = list(parsed_task.keys())[0]  
+  
+        # Construct the prompt  
+        try:  
+            task_description = tasks[selected_task]  
+        except KeyError:  
+            print(f"Selected task not found: {selected_task}")  
+            return f"Error: Task '{selected_task}' not found. Please select a valid task."  
+  
+        query_prefix = f"Instruct: {task_description}\nQuery: "  
+        queries = [input_text]  
+  
+        # Get the embeddings  
+        with torch.no_grad():  
+            inputs = self.tokenizer(queries, return_tensors='pt', padding=True, truncation=True, max_length=4096).to(self.device)  
+            outputs = self.model(**inputs)  
+            query_embeddings = outputs.last_hidden_state.mean(dim=1)  
+  
+            # Normalize embeddings  
+            query_embeddings = F.normalize(query_embeddings, p=2, dim=1)  
+            embeddings_list = query_embeddings.detach().cpu().numpy().tolist()  
+            self.clear_cuda_cache()  
+            return embeddings_list  
+  
+class MyEmbeddingFunction(EmbeddingFunction):  
+    def __init__(self, embedding_generator: EmbeddingGenerator):  
+        self.embedding_generator = embedding_generator  
+  
+    def __call__(self, input: Documents) -> Embeddings:  
+        embeddings = [self.embedding_generator.compute_embeddings(doc) for doc in input]  
+        embeddings = [item for sublist in embeddings for item in sublist]  
+        return embeddings  
+  
+def load_documents(file_path: str, mode: str = "elements"):  
+    loader = UnstructuredFileLoader(file_path, mode=mode)  
+    docs = loader.load()  
+    return [doc.page_content for doc in docs]  
+  
+def initialize_chroma(collection_name: str, embedding_function: MyEmbeddingFunction):  
+    client = HttpClient(settings=Settings(allow_reset=True))  
+    client.reset()  # resets the database  
+    collection = client.create_collection(collection_name)  
+    return client, collection  
+  
+def add_documents_to_chroma(client, collection, documents: list, embedding_function: MyEmbeddingFunction):  
+    for doc in documents:  
+        collection.add(ids=[str(uuid.uuid1())], documents=[doc], embeddings=embedding_function([doc]))  
+  
+def query_chroma(client, collection_name: str, query_text: str, embedding_function: MyEmbeddingFunction):  
+    db = Chroma(client=client, collection_name=collection_name, embedding_function=embedding_function)  
+    result_docs = db.similarity_search(query_text)  
+    return result_docs  
   
-
 # 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(collection_name="Tonic-instruct" , embedding_function=embedding_function)  
-  
+chroma_client, chroma_collection = initialize_chroma(collection_name="Tonic-instruct", embedding_function=embedding_function)
+
 def respond(  
     message,  
     history: list[tuple[str, str]],