Update app.py

added spectral 2 first attempt

app.py

@@ -44,6 +44,16 @@ from pyalex import Works, Authors, Sources, Institutions, Concepts, Publishers,
 from itertools import chain
 from compress_pickle import load, dump
 
+
+
+from transformers import AutoTokenizer
+from adapters import AutoAdapterModel
+import torch
+from tqdm import tqdm
+
+
+
+
 def query_records(search_term):
     def invert_abstract(inv_index):
         if inv_index is not None:
@@ -67,6 +77,67 @@ def query_records(search_term):
 
 
 
+################# Setting up the model for specter2 embeddings ###################
+
+
+
+device = torch.device("mps" if torch.backends.mps.is_available() else "cuda")
+print(f"Using device: {device}")
+
+tokenizer = AutoTokenizer.from_pretrained('allenai/specter2_aug2023refresh_base')
+model = AutoAdapterModel.from_pretrained('allenai/specter2_aug2023refresh_base')
+
+
+
+def create_embeddings(texts_to_embedd):
+    # Set up the device
+
+
+    print(len(texts_to_embedd))
+
+    # Load the proximity adapter and activate it
+    model.load_adapter("allenai/specter2_aug2023refresh", source="hf", load_as="proximity", set_active=True)
+    model.set_active_adapters("proximity")
+ 
+    model.to(device)
+
+    def batch_generator(data, batch_size):
+        """Yield consecutive batches of data."""
+        for i in range(0, len(data), batch_size):
+            yield data[i:i + batch_size]
+
+    @spaces.GPU(duration=120)
+    def encode_texts(texts, device, batch_size=16):
+        """Process texts in batches and return their embeddings."""
+        model.eval()
+        with torch.no_grad():
+            all_embeddings = []
+            count = 0
+            for batch in tqdm(batch_generator(texts, batch_size)):
+                inputs = tokenizer(batch, padding=True, truncation=True, return_tensors="pt", max_length=512).to(device)
+                outputs = model(**inputs)
+                embeddings = outputs.last_hidden_state[:, 0, :]  # Taking the [CLS] token representation
+                
+                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()
+                    count = 0
+                
+                count +=1
+            
+            all_embeddings = torch.cat(all_embeddings, dim=0)
+        return all_embeddings
+
+    # Concatenate title and abstract
+    embeddings = encode_texts(texts_to_embedd, device, batch_size=32).cpu().numpy()  # Process texts in batches of 10
+
+    return embeddings
+
+
+
+
+
 def predict(text_input, progress=gr.Progress()):
     
     # get data.
@@ -75,6 +146,24 @@ def predict(text_input, progress=gr.Progress()):
     
     
     
+    texts_to_embedd = [title + tokenizer.sep_token + publication + tokenizer.sep_token  + abstract for title, publication, abstract in zip(records_df['title'],records_df['parsed_publication'], records_df['abstract'])]
+    
+    embeddings = create_embeddings(texts_to_embedd)
+    print(embeddings)
+    
+    
+    
+    
+    
+    
+    
+    
+    
+    
+    
+    
+    
+    
     file_name = f"{datetime.utcnow().strftime('%s')}.html"
     file_path = static_dir / file_name
     print(file_path)