updated embedding model

app.py

@@ -21,6 +21,8 @@ import gradio as gr
 from datetime import datetime
 import sys
 
+from sentence_transformers import SentenceTransformer
+
 
 
 gr.set_static_paths(paths=["static/"])
@@ -183,54 +185,16 @@ print(f"Setting up language model: {time.strftime('%Y-%m-%d %H:%M:%S')}")
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 print(f"Using device: {device}")
 
-tokenizer = AutoTokenizer.from_pretrained('allenai/specter2_aug2023refresh_base')
-model = AutoAdapterModel.from_pretrained('allenai/specter2_aug2023refresh_base')
+#tokenizer = AutoTokenizer.from_pretrained('allenai/specter2_aug2023refresh_base')
+#model = AutoAdapterModel.from_pretrained('allenai/specter2_aug2023refresh_base')
+
+model = SentenceTransformer("m7n/discipline-tuned_specter_2_024")
 
 
 @spaces.GPU(duration=60)
 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]
-
-    
-    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()
-                    torch.cuda.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
+    embeddings = model.encode(texts_to_embedd,show_progress_bar=True,batch_size=32)
 
     return embeddings