bug fixes and usability improvements

app.py

@@ -57,8 +57,8 @@ class ArxivSearch:
                 outputs=self.output_md
             )
             self.embedding_dropdown.change(
-                self.search_function,
-                inputs=[self.query_box, self.embedding_dropdown],
+                self.model_switch,
+                inputs=[self.embedding_dropdown],
                 outputs=self.output_md
             )
             self.plot_button.click(
@@ -73,11 +73,12 @@ class ArxivSearch:
             )
 
         self.load_data(dataset)
-        self.load_model('tfidf')
-        self.load_model('word2vec')
-        self.load_model('bert')
-        self.load_model('scibert')
-        self.load_model('sbert')
+        self.load_model(embedding)
+        # self.load_model('tfidf')
+        # self.load_model('word2vec')
+        # self.load_model('bert')
+        # self.load_model('scibert')
+        # self.load_model('sbert')
 
         self.iface.launch()
 
@@ -114,7 +115,6 @@ class ArxivSearch:
             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)
@@ -149,7 +149,9 @@ class ArxivSearch:
             z=reduced_data[:, 2],
             mode='markers',
             marker=dict(size=3.5, color="#ffffff", opacity=0.2),
-            name='All Documents'
+            name='All Documents',   
+            text=[f"<br>: {self.arxiv_ids[i] if self.arxiv_ids[i] else self.documents[i].split()[:10]}" for i in range(len(self.documents))],
+            hoverinfo='text'
         )
         layout = go.Layout(
             margin=dict(l=0, r=0, b=0, t=0),
@@ -172,7 +174,9 @@ class ArxivSearch:
                 z=reduced_results_points[:, 2],
                 mode='markers',
                 marker=dict(size=3.5, color='orange', opacity=0.75),
-                name='Results'
+                name='Results',
+                text=[f"<br>Snippet: {self.documents[i][:200]}" for i in results_indices],
+                hoverinfo='text'
             )
             if not self.embedding == "tfidf" and self.query_encoding is not None and self.query_encoding.shape[0] > 0:
                 query_trace = go.Scatter3d(
@@ -181,7 +185,9 @@ class ArxivSearch:
                     z=query_point[:, 2],
                     mode='markers',
                     marker=dict(size=5, color='red', opacity=0.8),
-                    name='Query'
+                    name='Query',
+                    text=[f"<br>Query: {self.query}"],
+                    hoverinfo='text'
                 )
                 fig = go.Figure(data=[trace, results_trace, query_trace], layout=layout)
             else:
@@ -209,7 +215,7 @@ class ArxivSearch:
         if not tokens:
             return []
         vectors = np.array([self.wv_model[word] for word in tokens])
-        query_vec = normalize(np.mean(vectors, axis=0).reshape(1, -1))
+        query_vec = np.mean(vectors, axis=0).reshape(1, -1)
         self.query_encoding = query_vec
         sims = cosine_similarity(query_vec, self.word2vec_embeddings).flatten()
         top_indices = sims.argsort()[::-1][:top_n]
@@ -219,7 +225,6 @@ class ArxivSearch:
         with torch.no_grad():
             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 = outputs.last_hidden_state[:, 0, :].numpy()
 
         self.query_encoding = query_vec
@@ -251,6 +256,38 @@ class ArxivSearch:
         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 model_switch(self, embedding, progress=gr.Progress()):
+        if self.embedding != embedding:
+            old_embedding = self.embedding
+            print(f"Switching model to {embedding}")
+            self.load_model(embedding)
+            print(f"Loaded {embedding} model")
+            self.embedding = embedding
+            if old_embedding == "tfidf":
+                del self.tfidf_matrix
+                del self.feature_names
+            if old_embedding == "word2vec":
+                del self.word2vec_embeddings
+                del self.wv_model
+            if old_embedding == "bert":
+                del self.bert_embeddings
+                del self.tokenizer
+                del self.model
+            if old_embedding == "scibert":
+                del self.scibert_embeddings
+                del self.sci_tokenizer
+                del self.sci_model
+            if old_embedding == "sbert":
+                del self.sbert_model
+                del self.sbert_embedding
+                del self.cross_encoder
+            print(f"old embedding removed")
+            if hasattr(self, "query") and self.query:
+                return self.search_function(self.query, self.embedding)
+            else:
+                return ""  # Or a message like "Model switched. Please enter a query."
+        return gr.update()  # No change if embedding is the same
 
     def load_model(self, embedding):
         self.embedding = embedding
@@ -291,8 +328,9 @@ class ArxivSearch:
     def set_embedding(self, embedding):
         self.embedding = embedding
 
-    def search_function(self, query, embedding):
+    def search_function(self, query, embedding, progress=gr.Progress()):
         self.set_embedding(embedding)
+        self.query = query
         query = query.encode().decode('unicode_escape')  # Interpret escape sequences
 
         # Load or switch embedding model here if needed