query is also shown now

app.py

@@ -23,6 +23,7 @@ class ArxivSearch:
         self.raw_texts = []
         self.arxiv_ids = []
         self.last_results = []
+        self.query_encoding = None
 
         self.embedding_dropdown = gr.Dropdown(
             choices=["tfidf", "word2vec", "bert"],
@@ -113,20 +114,7 @@ class ArxivSearch:
             self.documents.append(text.strip())
             self.arxiv_ids.append(arxiv_id)
 
-    def keyword_match_ranking(self, query, top_n=5):
-        query_terms = query.lower().split()
-        query_indices = [i for i, term in enumerate(self.feature_names) if term in query_terms]
-        if not query_indices:
-            return []
-        scores = []
-        for doc_idx in range(self.tfidf_matrix.shape[0]):
-            doc_vector = self.tfidf_matrix[doc_idx]
-            doc_score = sum(doc_vector[0, i] for i in query_indices)
-            if doc_score > 0:
-                scores.append((doc_idx, doc_score))
-        scores.sort(key=lambda x: x[1], reverse=True)
-        return scores[:top_n]
-    
+
     def plot_3d_embeddings(self, embedding):
         # Example: plot random points, replace with your embeddings
         pca = PCA(n_components=3)
@@ -144,6 +132,7 @@ class ArxivSearch:
             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))
 
         elif embedding == "bert":
             all_indices = list(set(results_indices) | set(range(min(5000, self.bert_embeddings.shape[0]))))
@@ -151,6 +140,7 @@ class ArxivSearch:
             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))
 
         else:
             raise ValueError(f"Unsupported embedding type: {embedding}")
@@ -159,7 +149,8 @@ class ArxivSearch:
             y=reduced_data[:, 1],
             z=reduced_data[:, 2],
             mode='markers',
-            marker=dict(size=3.5, color='white', opacity=0.4),
+            marker=dict(size=3.5, color='#cccccc', opacity=0.35),
+            name='All Documents'
         )
         layout = go.Layout(
             margin=dict(l=0, r=0, b=0, t=0),
@@ -182,18 +173,42 @@ class ArxivSearch:
                 z=reduced_results_points[:, 2],
                 mode='markers',
                 marker=dict(size=3.5, color='orange', opacity=0.75),
+                name='Results'
+            )
+            query_trace = go.Scatter3d(
+                x=query_point[:, 0],
+                y=query_point[:, 1],
+                z=query_point[:, 2],
+                mode='markers',
+                marker=dict(size=5, color='red', opacity=0.8),
+                name='Query'
             )
-            fig = go.Figure(data=[trace, results_trace], layout=layout)
+            fig = go.Figure(data=[trace, results_trace, query_trace], layout=layout)
         else:
             fig = go.Figure(data=[trace], layout=layout)
         return fig
-
+    
+    def keyword_match_ranking(self, query, top_n=5):
+        query_terms = query.lower().split()
+        query_indices = [i for i, term in enumerate(self.feature_names) if term in query_terms]
+        if not query_indices:
+            return []
+        scores = []
+        for doc_idx in range(self.tfidf_matrix.shape[0]):
+            doc_vector = self.tfidf_matrix[doc_idx]
+            doc_score = sum(doc_vector[0, i] for i in query_indices)
+            if doc_score > 0:
+                scores.append((doc_idx, doc_score))
+        scores.sort(key=lambda x: x[1], reverse=True)
+        return scores[:top_n]
+    
     def word2vec_search(self, query, top_n=5):
         tokens = [word for word in query.split() if word in self.wv_model.key_to_index]
         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))
+        self.query_encoding = query_vec
         sims = cosine_similarity(query_vec, self.word2vec_embeddings).flatten()
         top_indices = sims.argsort()[::-1][:top_n]
         return [(i, sims[i]) for i in top_indices]
@@ -203,6 +218,7 @@ class ArxivSearch:
             inputs = self.tokenizer(query, return_tensors="pt", truncation=True, padding=True)
             outputs = self.model(**inputs)
             query_vec = normalize(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]
         return [(i, sims[i]) for i in top_indices]