added plot

app.py

@@ -8,6 +8,9 @@ from transformers import BertTokenizer, BertModel
 import numpy as np
 from datasets import load_dataset
 from gensim.models import KeyedVectors
+import plotly.graph_objects as go
+from sklearn.decomposition import PCA
+
 
 
 
@@ -19,6 +22,7 @@ class ArxivSearch:
         self.titles = []
         self.raw_texts = []
         self.arxiv_ids = []
+        self.last_results = []
 
         self.embedding_dropdown = gr.Dropdown(
             choices=["tfidf", "word2vec", "bert"],
@@ -26,16 +30,48 @@ class ArxivSearch:
             label="Model"
             )
         
-        self.iface = gr.Interface(
-            fn=self.search_function,
-            inputs=[
-                gr.Textbox(lines=1, placeholder="Enter your search query"),
-                self.embedding_dropdown
-            ],
-            outputs=gr.Markdown(),
-            title="arXiv Search Engine",
-            description="Search arXiv papers by keyword and embedding model.",
-        )
+
+                # Add a button to show the 3D plot
+        self.plot_button = gr.Button("Show 3D Plot")
+
+        # Define the interface using Blocks for more flexibility
+        with gr.Blocks() as self.iface:
+            gr.Markdown("# arXiv Search Engine")
+            gr.Markdown("Search arXiv papers by keyword and embedding model.")
+            with gr.Row():
+                self.query_box = gr.Textbox(lines=1, placeholder="Enter your search query", label="Query")
+                self.embedding_dropdown.render()
+                self.plot_button.render()
+            with gr.Row():
+                self.plot_output = gr.Plot()
+                self.output_md = gr.Markdown()
+
+            self.query_box.submit(
+                self.search_function,
+                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],
+                outputs=self.output_md
+            )
+            self.plot_button.click(
+                self.plot_3d_embeddings,
+                inputs=[self.embedding_dropdown],
+                outputs=self.plot_output
+            )
+
+        # self.iface = gr.Interface(
+        #     fn=self.search_function,
+        #     inputs=[
+        #         gr.Textbox(lines=1, placeholder="Enter your search query"),
+        #         self.embedding_dropdown
+        #     ],
+        #     outputs=gr.Markdown(),
+        #     title="arXiv Search Engine",
+        #     description="Search arXiv papers by keyword and embedding model.",
+        # )
 
         self.load_data(dataset)
         # self.load_model(embedding)
@@ -45,15 +81,6 @@ class ArxivSearch:
 
         self.iface.launch()
 
-            
-    # # --- Load data and embeddings ---
-    # with open("feature_names.txt", "r") as f:
-    #     feature_names = [line.strip() for line in f]
-
-    # tfidf_matrix = load_npz("tfidf_matrix_train.npz")
-
-    # Load dataset and initialize search engine
-
     def load_data(self, dataset):
         train_data = dataset["train"]
         for item in train_data.select(range(len(train_data))):
@@ -99,6 +126,57 @@ class ArxivSearch:
                 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)
+        results_indices = [i[0] for i in self.last_results]
+        if embedding == "tfidf":
+            reduced_data = pca.fit_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 embedding == "word2vec":
+            reduced_data = pca.fit_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))
+
+        elif embedding == "bert":
+            reduced_data = pca.fit_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))
+        else:
+            raise ValueError(f"Unsupported embedding type: {embedding}")
+        trace = go.Scatter3d(
+            x=reduced_data[:, 0],
+            y=reduced_data[:, 1],
+            z=reduced_data[:, 2],
+            mode='markers',
+            marker=dict(size=3.5, color='white', opacity=0.4),
+        )
+        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=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'),
+            ),
+            paper_bgcolor='black',   # Outside the plotting area
+            plot_bgcolor='black',    # Plotting area
+            font=dict(color='white') # Axis and legend text
+        )
+        if len(reduced_results_points) > 0:
+            results_trace = go.Scatter3d(
+                x=reduced_results_points[:, 0],
+                y=reduced_results_points[:, 1],
+                z=reduced_results_points[:, 2],
+                mode='markers',
+                marker=dict(size=3.5, color='orange', opacity=0.9),
+            )
+            fig = go.Figure(data=[trace, results_trace], layout=layout)
+        else:
+            fig = go.Figure(data=[trace], layout=layout)
+        return fig
 
     def word2vec_search(self, query, top_n=5):
         tokens = [word for word in query.split() if word in self.wv_model.key_to_index]
@@ -163,7 +241,12 @@ class ArxivSearch:
             return "No results found."
 
         if not results:
+            self.last_results = []
             return "No results found."
+        
+
+        if results:
+            self.last_results = results
 
         output = ""
         display_rank = 1