added sbert and simplified plotting

BERT embeddings/sbert_embedding.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c5285b4f848c6dd13e7141a2684857daf6d8d02fdb18fb4812182fe31780c717
+size 40407781

app.py

@@ -11,8 +11,8 @@ from gensim.models import KeyedVectors
 import plotly.graph_objects as go
 from sklearn.decomposition import PCA
 from transformers import AutoTokenizer, AutoModel
-
-
+from sentence_transformers import CrossEncoder
+from sentence_transformers import SentenceTransformer
 
 class ArxivSearch:
     def __init__(self, dataset, embedding="bert"):
@@ -28,7 +28,7 @@ class ArxivSearch:
 
         # model selection
         self.embedding_dropdown = gr.Dropdown(
-            choices=["tfidf", "word2vec", "bert", "scibert"],
+            choices=["tfidf", "word2vec", "bert", "scibert", "sbert"],
             value="bert",
             label="Model"
             )
@@ -56,7 +56,6 @@ class ArxivSearch:
                 inputs=[self.query_box, self.embedding_dropdown],
                 outputs=self.output_md
             )
-
             self.embedding_dropdown.change(
                 self.search_function,
                 inputs=[self.query_box, self.embedding_dropdown],
@@ -78,6 +77,7 @@ class ArxivSearch:
         self.load_model('word2vec')
         self.load_model('bert')
         self.load_model('scibert')
+        self.load_model('sbert')
 
         self.iface.launch()
 
@@ -113,6 +113,16 @@ class ArxivSearch:
             self.documents.append(text.strip())
             self.arxiv_ids.append(arxiv_id)
 
+    def plot_dense(self, embedding, pca, results_indices):
+        print(self.query_encoding.shape[0])
+        all_indices = list(set(results_indices) | set(range(min(5000, embedding.shape[0]))))
+        all_data = embedding[all_indices]
+        pca.fit(all_data)
+        reduced_data = pca.transform(embedding[:5000])
+        reduced_results_points = pca.transform(embedding[results_indices]) if len(results_indices) > 0 else np.empty((0, 3))
+        query_point = pca.transform(self.query_encoding) if self.query_encoding is not None and self.query_encoding.shape[0] > 0 else np.empty((0, 3))
+        return reduced_data, reduced_results_points, query_point
+
     def plot_3d_embeddings(self):
         # Example: plot random points, replace with your embeddings
         pca = PCA(n_components=3)
@@ -123,29 +133,14 @@ class ArxivSearch:
             pca.fit(all_data)
             reduced_data = pca.transform(self.tfidf_matrix[:5000].toarray())
             reduced_results_points = pca.transform(self.tfidf_matrix[results_indices].toarray()) if len(results_indices) > 0 else np.empty((0, 3))
-
         elif self.embedding == "word2vec":
-            all_indices = list(set(results_indices) | set(range(min(5000, self.word2vec_embeddings.shape[0]))))
-            all_data = self.word2vec_embeddings[all_indices]
-            pca.fit(all_data)
-            reduced_data = pca.transform(self.word2vec_embeddings[:5000])
-            reduced_results_points = pca.transform(self.word2vec_embeddings[results_indices]) if len(results_indices) > 0 else np.empty((0, 3))
-            query_point = pca.transform(self.query_encoding) if self.query_encoding is not None and self.query_encoding.shape[0] > 0 else np.empty((0, 3))
-
+            reduced_data, reduced_results_points, query_point = self.plot_dense(self.word2vec_embeddings, pca, results_indices)
         elif self.embedding == "bert":
-            all_indices = list(set(results_indices) | set(range(min(5000, self.bert_embeddings.shape[0]))))
-            all_data = self.bert_embeddings[all_indices]
-            pca.fit(all_data)
-            reduced_data = pca.transform(self.bert_embeddings[:5000])
-            reduced_results_points = pca.transform(self.bert_embeddings[results_indices]) if len(results_indices) > 0 else np.empty((0, 3))
-            query_point = pca.transform(self.query_encoding) if self.query_encoding is not None and self.query_encoding.shape[0] > 0 else np.empty((0, 3))
+            reduced_data, reduced_results_points, query_point = self.plot_dense(self.bert_embeddings, pca, results_indices)
+        elif self.embedding == "sbert":
+            reduced_data, reduced_results_points, query_point = self.plot_dense(self.sbert_embedding, pca, results_indices)
         elif self.embedding == "scibert":
-            all_indices = list(set(results_indices) | set(range(min(5000, self.scibert_embeddings.shape[0]))))
-            all_data = self.scibert_embeddings[all_indices]
-            pca.fit(all_data)
-            reduced_data = pca.transform(self.scibert_embeddings[:5000])
-            reduced_results_points = pca.transform(self.scibert_embeddings[results_indices]) if len(results_indices) > 0 else np.empty((0, 3))
-            query_point = pca.transform(self.query_encoding) if self.query_encoding is not None and self.query_encoding.shape[0] > 0 else np.empty((0, 3))
+            reduced_data, reduced_results_points, query_point = self.plot_dense(self.scibert_embeddings, pca, results_indices)
         else:
             raise ValueError(f"Unsupported embedding type: {self.embedding}")
         trace = go.Scatter3d(
@@ -159,9 +154,9 @@ class ArxivSearch:
         layout = go.Layout(
             margin=dict(l=0, r=0, b=0, t=0),
             scene=dict(
-                xaxis_title='X',
-                yaxis_title='Y',
-                zaxis_title='Z',
+                xaxis_title='PCA 1',
+                yaxis_title='PCA 2',
+                zaxis_title='PCA 3',
                 xaxis=dict(backgroundcolor='black', color='white', gridcolor='gray', zerolinecolor='gray'),
                 yaxis=dict(backgroundcolor='black', color='white', gridcolor='gray', zerolinecolor='gray'),
                 zaxis=dict(backgroundcolor='black', color='white', gridcolor='gray', zerolinecolor='gray'),
@@ -222,40 +217,40 @@ class ArxivSearch:
 
     def bert_search(self, query, top_n=10):
         with torch.no_grad():
-            inputs = self.tokenizer(query, return_tensors="pt", truncation=True, padding=True)
+            inputs = self.tokenizer((query+' ')*2, return_tensors="pt", truncation=True, max_length=512, padding='max_length')
             outputs = self.model(**inputs)
-            query_vec = normalize(outputs.last_hidden_state[:, 0, :].numpy())
+            # query_vec = normalize(outputs.last_hidden_state[:, 0, :].numpy())
+            query_vec = outputs.last_hidden_state[:, 0, :].numpy()
 
         self.query_encoding = query_vec
         sims = cosine_similarity(query_vec, self.bert_embeddings).flatten()
         top_indices = sims.argsort()[::-1][:top_n]
+        print(f"sim, top_indices: {sims}, {top_indices}")
         return [(i, sims[i]) for i in top_indices]
 
     def scibert_search(self, query, top_n=10):
         with torch.no_grad():
             inputs = self.sci_tokenizer(query, return_tensors="pt", truncation=True, padding=True, max_length=512)
             outputs = self.sci_model(**inputs)
-            query_vec = normalize(outputs.last_hidden_state[:, 0, :].numpy())
+            query_vec = outputs.last_hidden_state[:, 0, :].numpy()
 
         self.query_encoding = query_vec
         sims = cosine_similarity(query_vec, self.scibert_embeddings).flatten()
         top_indices = sims.argsort()[::-1][:top_n]
+        print(f"sim, top_indices: {sims}, {top_indices}")
         return [(i, sims[i]) for i in top_indices]
 
-    def bert_search_2(self, query, top_n=10):
-        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
+    def sbert_search(self, query, top_n=10):
+        query_vec = self.sbert_model.encode([query])
         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]
+        cos_scores = cosine_similarity(query_vec, self.sbert_embedding)[0]
+        top_k_indices = np.argsort(cos_scores)[-50:][::-1]
+        candidates = [dataset['train'][int(i)]['text'] for i in top_k_indices]
+        scores = self.cross_encoder.predict([(query, doc) for doc in candidates])
+        final_scores = 0.7 * scores + 0.3 * cos_scores[top_k_indices]
+        top_indices = top_k_indices[final_scores.argsort()[::-1][:top_n]]
+        print(f"sim, top_indices: {final_scores}, {top_indices}")
+        return [(top_k_indices[i], final_scores[i]) for i in final_scores.argsort()[::-1][:top_n]]
 
     def load_model(self, embedding):
         self.embedding = embedding
@@ -277,6 +272,10 @@ class ArxivSearch:
             self.sci_tokenizer = AutoTokenizer.from_pretrained('allenai/scibert_scivocab_uncased')
             self.sci_model = AutoModel.from_pretrained('allenai/scibert_scivocab_uncased')
             self.sci_model.eval()
+        elif self.embedding == "sbert":
+            self.sbert_model = SentenceTransformer("all-MiniLM-L6-v2")
+            self.sbert_embedding = np.load("BERT embeddings/sbert_embedding.npz")["sbert_embedding"]
+            self.cross_encoder = CrossEncoder("cross-encoder/ms-marco-MiniLM-L6-v2")
         else:
             raise ValueError(f"Unsupported embedding type: {self.embedding}")
         
@@ -285,13 +284,15 @@ class ArxivSearch:
         pattern = re.compile(r'a\s*b\s*s\s*t\s*r\s*a\s*c\s*t|i\s*n\s*t\s*r\s*o\s*d\s*u\s*c\s*t\s*i\s*o\s*n', re.IGNORECASE)
         match = pattern.search(text)
         if match:
-            return text[:match.start()].strip()
+            return text[:match.start()].strip() if match.start() < 1000 else text[:100].strip()
         else:
-            return text[:100].strip()
+            return text[:300].strip()
 
+    def set_embedding(self, embedding):
+        self.embedding = embedding
 
     def search_function(self, query, embedding):
-        self.embedding = embedding
+        self.set_embedding(embedding)
         query = query.encode().decode('unicode_escape')  # Interpret escape sequences
 
         # Load or switch embedding model here if needed
@@ -303,6 +304,8 @@ class ArxivSearch:
             results = self.bert_search(query)
         elif self.embedding == "scibert":
             results = self.scibert_search(query)
+        elif self.embedding == "sbert":
+            results = self.sbert_search(query)
         else:
             return "No results found."
 
@@ -317,13 +320,14 @@ class ArxivSearch:
         display_rank = 1
         for idx, score in results:
             if not self.arxiv_ids[idx]:
-                continue
-
-            link = f"https://arxiv.org/abs/{self.arxiv_ids[idx].replace('arxiv:', '')}"
-            snippet = self.snippet_before_abstract(self.documents[idx]).replace('\n', '<br>')
-            output += f"### Document {display_rank}\n"
-            output += f"[arXiv Link]({link})\n\n"
-            output += f"<pre>{snippet}</pre>\n\n---\n"
+                output += f"### Document {display_rank}\n"
+                output += f"<pre>{self.documents[idx][:200]}</pre>\n\n"
+            else:
+                link = f"https://arxiv.org/abs/{self.arxiv_ids[idx].replace('arxiv:', '')}"
+                snippet = self.snippet_before_abstract(self.documents[idx]).replace('\n', '<br>')
+                output += f"### Document {display_rank}\n"
+                output += f"[arXiv Link]({link})\n\n"
+                output += f"<pre>{snippet}</pre>\n\n---\n"
             display_rank += 1
 
         return output