Update app.py

app.py

@@ -72,7 +72,7 @@ import umap
 
 
 
-def query_records(search_term,progress):
+def query_records(search_term):
     def invert_abstract(inv_index):
         if inv_index is not None:
             l_inv = [(w, p) for w, pos in inv_index.items() for p in pos]
@@ -95,16 +95,17 @@ def query_records(search_term,progress):
     query_length = Works().search([search_term]).count()
     
     records = []
-    total_pages = (query_length + 199) // 200  # Calculate total number of pages
+    #total_pages = (query_length + 199) // 200  # Calculate total number of pages
+    progress=gr.Progress()
     
-    for i, record in enumerate(chain(*query.paginate(per_page=200))):
+    for i, record in progress.tqdm(enumerate(chain(*query.paginate(per_page=200)))):
         records.append(record)
         
         # Calculate progress from 0 to 0.1
-        achieved_progress = min(0.1, (i + 1) / query_length * 0.1)
+        #achieved_progress = min(0.1, (i + 1) / query_length * 0.1)
         
         # Update progress bar
-        progress(achieved_progress, desc="Getting queried data...")
+        #progress(achieved_progress, desc="Getting queried data...")
 
 
     
@@ -128,7 +129,9 @@ def query_records(search_term,progress):
 
 
 
-device = torch.device("mps" if torch.backends.mps.is_available() else "cuda")
+#device = torch.device("mps" if torch.backends.mps.is_available() else "cuda")
+#print(f"Using device: {device}")
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 print(f"Using device: {device}")
 
 tokenizer = AutoTokenizer.from_pretrained('allenai/specter2_aug2023refresh_base')
@@ -168,7 +171,8 @@ def create_embeddings(texts_to_embedd):
                 all_embeddings.append(embeddings.cpu())  # Move to CPU to free GPU memory
                 #torch.mps.empty_cache()  # Clear cache to free up memory
                 if count == 100:
-                    torch.mps.empty_cache()
+                    #torch.mps.empty_cache()
+                    torch.cuda.empty_cache()
                     count = 0
                 
                 count +=1