Upload app.py

app.py

@@ -18,6 +18,7 @@ from huggingface_hub import hf_hub_download
 from contextual import ContextualAI
 from openai import AzureOpenAI
 from datetime import datetime
+import sys
 
 """
 # to switch:
@@ -328,6 +329,7 @@ def run_retrieval(query):
     """
     start_time = time.time()
     query_embeddings = run_dense_retrieval(query)
+    np.save("test_query_embeddings", query_embeddings)
     print("--- Nvidia Embedding: %s seconds ---" % (time.time() - start_time))
     D, I = faiss_index.search(query_embeddings, 45)
     print("--- Faiss retrieval: %s seconds ---" % (time.time() - start_time))
@@ -362,7 +364,7 @@ def run_retrieval(query):
         out_dict.append(tmp)
     print (out_dict)
     # and now, rerank
-    out_dict = rerank_with_contextual_AI(out_dict)
+    #out_dict = rerank_with_contextual_AI(out_dict)
     return out_dict
 
 
@@ -376,10 +378,20 @@ device = torch.device("cuda")
 extractive_qa = pipeline("question-answering", model="ai-law-society-lab/extractive-qa-model", tokenizer="FacebookAI/roberta-large", device_map="auto", token=os.getenv('hf_token'))
 ids, titles, chunks = load_NJ_caselaw()
 
-ds = load_dataset("ai-law-society-lab/NJ_embeddings", token=os.getenv('hf_token'))["train"]
-ds = ds.with_format("np")
-print (ds)
-faiss_index = load_faiss_index(ds["embeddings"])
+#@profile
+def profiling_faiss_index():
+    ds = load_dataset("ai-law-society-lab/NJ_embeddings", token=os.getenv('hf_token'))["train"]
+    print (sys.getsizeof(ds))
+    ds = ds.with_format("np")
+    print (sys.getsizeof(ds))
+    print (ds)
+    #faiss_index = load_faiss_index(ds["embeddings"])
+    ds.add_faiss_index(column='embeddings')
+    #print (sys.getsizeof(faiss_index))
+
+    return ds
+faiss_index = profiling_faiss_index()
+
 
 with open("NJ_caselaw_metadata.json") as f:
     metadata = json.load(f)
@@ -389,6 +401,7 @@ with open("NJ_caselaw_metadata.json") as f:
 def load_embeddings_model(model_name = "intfloat/e5-large-v2"):
     if "NV" in model_name:
         model = AutoModel.from_pretrained('nvidia/NV-Embed-v2', trust_remote_code=True, torch_dtype=torch.bfloat16, device_map="auto")
+        #model = AutoModel.from_pretrained('nvidia/NV-Embed-v2', trust_remote_code=True, torch_dtype=torch.float16, device_map="auto")        
         model.eval()
         return model