remove shared link

arxiv_tool/app.py

@@ -4,7 +4,12 @@ from plot import EmbeddingPlotter
 
 TITLE = "Search tool for ArXiv papers"
 DESCRIPTION = "<center>Find your most beloved ArXiv papers!</center>"
-EXAMPLES=["RoBERTa optimisation", "Permutation invariant AI models", "Gradient descent", "Black hole information theory"]
+EXAMPLES = [
+    "RoBERTa optimisation",
+    "Permutation invariant AI models",
+    "Gradient descent",
+    "Black hole information theory",
+]
 ARTICLE = r"<center>Done by dr. Gabriel Lopez<br> For more please visit: <a href='https://sites.google.com/view/dr-gabriel-lopez/home'>My Page</a></center>"
 
 # interface function
@@ -14,21 +19,27 @@ def search_and_plot(querry):
     df, result = SentenceEncoder().transform(df, querry, model, embeddings)
     # plot
     fig1, fig2 = EmbeddingPlotter().transform(df, embeddings)
-    return result[['title', 'similarity']], fig1, fig2
+    return result[["title", "similarity"]], fig1, fig2
+
 
 # gradio elements
-in_textbox = gr.Textbox(label="Search on ArXiv:", placeholder="what do you want to learn today?...", lines=1)
+in_textbox = gr.Textbox(
+    label="Search on ArXiv:", placeholder="what do you want to learn today?...", lines=1
+)
 # in_examples = gr.Examples(examples=["BERT optimization", "Gradient descent", "Black hole information theory"], inputs=in_textbox)
 out_dataframe = gr.DataFrame(label="Most similar papers on ArXiv:")
 out_plot_sphere = gr.Plot(label="Embedding projection over a unit sphere")
-out_plot_projected_sphere = gr.Plot(label="Lambert-conformal projection over a plane", visible=False)
+out_plot_projected_sphere = gr.Plot(
+    label="Lambert-conformal projection over a plane", visible=False
+)
 
 # launch interface
-gr.Interface(inputs=in_textbox, 
-             outputs=[out_dataframe,out_plot_sphere,out_plot_projected_sphere], 
-             examples=EXAMPLES,
-             fn=search_and_plot,
-             title=TITLE,
-             description=DESCRIPTION,
-             article=ARTICLE,
-             ).launch(share=True)
+gr.Interface(
+    inputs=in_textbox,
+    outputs=[out_dataframe, out_plot_sphere, out_plot_projected_sphere],
+    examples=EXAMPLES,
+    fn=search_and_plot,
+    title=TITLE,
+    description=DESCRIPTION,
+    article=ARTICLE,
+).launch()

arxiv_tool/core.py

@@ -2,73 +2,80 @@ import pandas as pd
 import numpy as np
 import nmslib
 from sentence_transformers import SentenceTransformer
+
 # TODO: Use pipe, remove embeddings
 
 
 class SentenceEncoder:
-    """ Encodes the querry and papers data set and finds elements with the lowest cosine similarity """
+    """Encodes the querry and papers data set and finds elements with the lowest cosine similarity"""
 
     def load_and_encode(self):
-      # load
-      df = self._load()
-      # encode
-      df, model, embeddings = self._encode_papers(df)
-      return df, model, embeddings
+        # load
+        df = self._load()
+        # encode
+        df, model, embeddings = self._encode_papers(df)
+        return df, model, embeddings
 
     def transform(self, df, querry, model, embeddings):
-      # create_index
-      emb_querry = self._econde_querry(querry, model)
-      # search 
-      result = self._make_search(df,emb_querry, embeddings)
-      # add_relevant_columns 
-      df = self._add_relevant_columns(df, result)
-      return df, result
+        # create_index
+        emb_querry = self._econde_querry(querry, model)
+        # search
+        result = self._make_search(df, emb_querry, embeddings)
+        # add_relevant_columns
+        df = self._add_relevant_columns(df, result)
+        return df, result
 
     def _load(self):
-      # Load data
-      df = pd.read_csv("data/arxiv.csv")
-      return df
+        # Load data
+        df = pd.read_csv("data/arxiv.csv")
+        return df
 
-    def _encode_papers(self,df):
-      # Encode the papers title
-      checkpoint = 'distilbert-base-uncased'
-      model = SentenceTransformer(checkpoint)
-      embeddings = model.encode(df['title'], convert_to_tensor=True)
-      # embeddings column
-      df['embeddings'] = np.array(embeddings).tolist()
-      return df, model, embeddings
+    def _encode_papers(self, df):
+        # Encode the papers title
+        checkpoint = "distilbert-base-uncased"
+        model = SentenceTransformer(checkpoint)
+        embeddings = model.encode(df["title"], convert_to_tensor=True)
+        # embeddings column
+        df["embeddings"] = np.array(embeddings).tolist()
+        return df, model, embeddings
 
-    def _econde_querry(self,querry, model):
-      # Encode the querry
-      emb_querry = model.encode([querry])
-      return emb_querry
+    def _econde_querry(self, querry, model):
+        # Encode the querry
+        emb_querry = model.encode([querry])
+        return emb_querry
 
     def _make_search(self, df, emb_querry, embeddings):
-      # initialize a new index, using a HNSW index on Cosine Similarity
-      index = nmslib.init(method='hnsw', space='cosinesimil')
-      index.addDataPointBatch(embeddings)
-      index.createIndex({'post': 2}, print_progress=True)
-      # search
-      result = self._extract_search_result(index, emb_querry, df, k=10)
-      return result
+        # initialize a new index, using a HNSW index on Cosine Similarity
+        index = nmslib.init(method="hnsw", space="cosinesimil")
+        index.addDataPointBatch(embeddings)
+        index.createIndex({"post": 2}, print_progress=True)
+        # search
+        result = self._extract_search_result(index, emb_querry, df, k=10)
+        return result
 
-    def _extract_search_result(self,index, emb_querry, df, k):
-      data = []
-      idx, distances = index.knnQuery(emb_querry, k=k)
-      for i, j in zip(idx, distances):
-        data.append({'index': i,
-                    'title': df.title[i], 
-                    'abstract': df.abstract[i],
-                    'similarity': 1.0 - j})
-      return pd.DataFrame(data)
+    def _extract_search_result(self, index, emb_querry, df, k):
+        data = []
+        idx, distances = index.knnQuery(emb_querry, k=k)
+        for i, j in zip(idx, distances):
+            data.append(
+                {
+                    "index": i,
+                    "title": df.title[i],
+                    "abstract": df.abstract[i],
+                    "similarity": 1.0 - j,
+                }
+            )
+        return pd.DataFrame(data)
 
     def _add_relevant_columns(self, df, result):
-      # get categories
-      df['categories_parsed'] = df.categories.str.split().apply(lambda x: x[0]).str.split('.').apply(lambda x: x[0])
-      # create columns for plotting
-      df['index_papers'] = df.index
-      df['selected'] = df.index_papers.apply(lambda x: x in list(result['index']) )
-      return df
-
-
-
+        # get categories
+        df["categories_parsed"] = (
+            df.categories.str.split()
+            .apply(lambda x: x[0])
+            .str.split(".")
+            .apply(lambda x: x[0])
+        )
+        # create columns for plotting
+        df["index_papers"] = df.index
+        df["selected"] = df.index_papers.apply(lambda x: x in list(result["index"]))
+        return df

arxiv_tool/plot.py

@@ -4,37 +4,48 @@ import plotly.express as px
 from pandas import DataFrame
 import numpy as np
 
+
 class EmbeddingPlotter:
-    """ Lower the dimensionality of the representation from 768 -> 2, over the surface of the sphere """
+    """Lower the dimensionality of the representation from 768 -> 2, over the surface of the sphere"""
 
     def transform(self, df, embeddings):
         df = self.umap_embedding(df, embeddings)
         fig1, fig2 = self.plot(df)
         return fig1, fig2
 
-
     def umap_embedding(self, df, embeddings):
         # UMAP - Spherical
-        sphere_mapper = umap.UMAP(output_metric='haversine', random_state=42).fit(np.array(embeddings))
-        df['spherical_emb_X'] = np.sin(sphere_mapper.embedding_[:,0])*np.cos(sphere_mapper.embedding_[:,1])
-        df['spherical_emb_Y'] = np.sin(sphere_mapper.embedding_[:,0])*np.sin(sphere_mapper.embedding_[:,1])
-        df['spherical_emb_Z'] = np.cos(sphere_mapper.embedding_[:,0])
+        sphere_mapper = umap.UMAP(output_metric="haversine", random_state=42).fit(
+            np.array(embeddings)
+        )
+        df["spherical_emb_X"] = np.sin(sphere_mapper.embedding_[:, 0]) * np.cos(
+            sphere_mapper.embedding_[:, 1]
+        )
+        df["spherical_emb_Y"] = np.sin(sphere_mapper.embedding_[:, 0]) * np.sin(
+            sphere_mapper.embedding_[:, 1]
+        )
+        df["spherical_emb_Z"] = np.cos(sphere_mapper.embedding_[:, 0])
         # UMAP - Lambert Conformal
-        df['lambert_conformal_emb_x'] = np.arctan2(df['spherical_emb_X'], df['spherical_emb_Y']) 
-        df['lambert_conformal_emb_y'] = -np.arccos(df['spherical_emb_Z'])
+        df["lambert_conformal_emb_x"] = np.arctan2(
+            df["spherical_emb_X"], df["spherical_emb_Y"]
+        )
+        df["lambert_conformal_emb_y"] = -np.arccos(df["spherical_emb_Z"])
         return df
 
     def plot(self, df):
         # on the 3d sphere
-        fig1 = px.scatter_3d(df, 
-                    x='spherical_emb_X', 
-                    y='spherical_emb_Y', 
-                    z='spherical_emb_Z',
-                    color="categories_parsed")
+        fig1 = px.scatter_3d(
+            df,
+            x="spherical_emb_X",
+            y="spherical_emb_Y",
+            z="spherical_emb_Z",
+            color="categories_parsed",
+        )
         # on the projected spehre
-        fig2 = px.scatter(data_frame=df ,
-                x='lambert_conformal_emb_x', 
-                y='lambert_conformal_emb_y', 
-                color="categories_parsed",
-                )
-        return fig1, fig2
+        fig2 = px.scatter(
+            data_frame=df,
+            x="lambert_conformal_emb_x",
+            y="lambert_conformal_emb_y",
+            color="categories_parsed",
+        )
+        return fig1, fig2