small updates

app.py

@@ -149,7 +149,7 @@ class ArxivSearch:
             y=reduced_data[:, 1],
             z=reduced_data[:, 2],
             mode='markers',
-            marker=dict(size=3.5, color='#cccccc', opacity=0.35),
+            marker=dict(size=3.5, color="#ffffff", opacity=0.2),
             name='All Documents'
         )
         layout = go.Layout(
@@ -223,6 +223,21 @@ class ArxivSearch:
         top_indices = sims.argsort()[::-1][:top_n]
         return [(i, sims[i]) for i in top_indices]
 
+    def bert_search_2(self, query, top_n=5):
+        with torch.no_grad():
+            inputs = self.tokenizer(query, return_tensors="pt", truncation=True, padding=True)
+            outputs = self.model(**inputs)
+        token_embeddings = outputs.last_hidden_state
+        attention_mask = inputs['attention_mask']
+        mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
+        sentence_embeddings = torch.sum(token_embeddings * mask_expanded, dim=1)
+        sum_mask = torch.clamp(mask_expanded.sum(1), min=1e-9)
+        query_vec = sentence_embeddings / sum_mask
+        self.query_encoding = query_vec
+        sims = cosine_similarity(query_vec, self.bert_embeddings).flatten()
+        top_indices = sims.argsort()[::-1][:top_n]
+        return [(i, sims[i]) for i in top_indices]
+
     def load_model(self, embedding):
         if embedding == "tfidf":
             self.tfidf_matrix = load_npz("TF-IDF embeddings/tfidf_matrix_train.npz")