add automatic metadata

app.py

@@ -1,10 +1,5 @@
 # main.py  
 import spaces
-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  
@@ -15,9 +10,16 @@ from langchain_chroma import Chroma
 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  
+import os  
+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  
+from globalvars import API_BASE, intention_prompt, tasks, system_message, model_name , metadata_prompt 
+
 
 load_dotenv()
 
@@ -34,7 +36,7 @@ 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")  
@@ -44,7 +46,7 @@ class EmbeddingGenerator:
   
     def clear_cuda_cache(self):  
         torch.cuda.empty_cache()  
-  
+
     @spaces.GPU  
     def compute_embeddings(self, input_text: str):  
         # Get the intention  
@@ -71,6 +73,17 @@ class EmbeddingGenerator:
         query_prefix = f"Instruct: {task_description}\nQuery: "  
         queries = [input_text]  
   
+        # Get the metadata  
+        metadata_completion = self.intention_client.chat.completions.create(  
+            model="yi-large",  
+            messages=[  
+                {"role": "system", "content": metadata_prompt},  
+                {"role": "user", "content": input_text}  
+            ]  
+        )  
+        metadata_output = metadata_completion.choices[0].message['content']  
+        metadata = self.extract_metadata(metadata_output)  
+  
         # Get the embeddings  
         with torch.no_grad():  
             inputs = self.tokenizer(queries, return_tensors='pt', padding=True, truncation=True, max_length=4096).to(self.device)  
@@ -80,8 +93,19 @@ class EmbeddingGenerator:
             # Normalize embeddings  
             query_embeddings = F.normalize(query_embeddings, p=2, dim=1)  
             embeddings_list = query_embeddings.detach().cpu().numpy().tolist()  
+  
+            # Include metadata in the embeddings  
+            embeddings_with_metadata = [{"embedding": emb, "metadata": metadata} for emb in embeddings_list]  
+  
             self.clear_cuda_cache()  
-            return embeddings_list  
+            return embeddings_with_metadata  
+  
+    def extract_metadata(self, metadata_output: str):  
+        # Regex pattern to extract key-value pairs  
+        pattern = re.compile(r'\"(\w+)\": \"([^\"]+)\"')  
+        matches = pattern.findall(metadata_output)  
+        metadata = {key: value for key, value in matches}  
+        return metadata  
   
 class MyEmbeddingFunction(EmbeddingFunction):  
     def __init__(self, embedding_generator: EmbeddingGenerator):