add intention mapper to embedding flow , refactor into classes

app.py

@@ -18,7 +18,12 @@ from globalvars import API_BASE, intention_prompt, tasks
 from dotenv import load_dotenv
 import re 
 from utils import load_env_variables
-
+import chromadb
+from chromadb import Documents, EmbeddingFunction, Embeddings  
+from chromadb.config import Settings  
+from chromadb import HttpClient  
+from langchain_community.document_loaders import UnstructuredFileLoader  
+from utils import load_env_variables  , parse_and_route
 
 os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:30'
 os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
@@ -29,23 +34,108 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 hf_token, yi_token = load_env_variables()
 
-## use instruct embeddings
-# Load the tokenizer and model
-tokenizer = AutoTokenizer.from_pretrained('nvidia/NV-Embed-v1', token = hf_token , trust_remote_code=True)
-model = AutoModel.from_pretrained('nvidia/NV-Embed-v1' , token = hf_token ,  trust_remote_code=True).to(device)
+def clear_cuda_cache():
+    torch.cuda.empty_cache()
 
-## add chroma vector store
 
-## Make intention Mapper 
+## 01ai Yi-large Clience 
 
-intention_client = OpenAI(
+client = OpenAI(
     api_key=yi_token,
     base_url=API_BASE
 )
-intention_completion = intention_client.chat.completions.create(
-    model="yi-large",
-    messages=[{"role": "system", "content": intention_prompt},{"role": "user", "content": inputext}]
-)
+
+
+## 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")  
+        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()  
+  
+    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  
+
+## add chroma vector store
+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  
+  
+
+
 # print(completion)
 
 def respond(

requirements.txt

@@ -4,4 +4,12 @@ sentence-transformers
 torch==2.2.0
 transformers
 openai
-python-dotenv
+python-dotenv
+chromadb
+langchain-community
+unstructured[all-docs]
+libmagic
+poppler
+tesseract
+libxml2
+libxslt

utils.py

@@ -1,3 +1,5 @@
+# utils.py
+
 import re 
 from dotenv import load_dotenv
 import re 
@@ -14,18 +16,18 @@ def load_env_variables():
 
     return hf_token, yi_token
 
-def parse_and_route(example_output: str):
-    # Regex pattern to match the true task
-    pattern = r'"(\w+)":\s?true'
-    
-    # Find the true task
-    match = re.search(pattern, example_output)
-    
-    if match:
-        true_task = match.group(1)
-        if true_task in tasks:
-            return {true_task: tasks[true_task]}
-        else:
-            return {true_task: "Task description not found"}
-    else:
-        return "No true task found in the example output"
+def parse_and_route(example_output: str):  
+    # Regex pattern to match the true task  
+    pattern = r'"(\w+)":\s?true'  
+  
+    # Find the true task  
+    match = re.search(pattern, example_output)  
+  
+    if match:  
+        true_task = match.group(1)  
+        if true_task in tasks:  
+            return {true_task: tasks[true_task]}  
+        else:  
+            return {true_task: "Task description not found"}  
+    else:  
+        return "No true task found in the example output"