scatterplots

app.py

@@ -5,7 +5,7 @@ import re
 from task_visualizations import TaskVisualizations
 import plotly.graph_objects as go
 from functools import partial
-from text_visualization import WordCloudExtractor
+from text_visualization import WordCloudExtractor,  EmbeddingVisualizer
 
 logging.basicConfig(level=logging.INFO)
 
@@ -108,6 +108,15 @@ def setup_repository_representations_tab(repos, representation_types):
         )
 
 
+def load_embeddings_intro_description():
+    return """
+    The following plots show embeddings obtained with MPNet sentence transformer after applying 2d UMAP algorithm for dimensionality reduction.
+
+    In the first scatterplot we display PapersWithCode tasks that are colored by area.
+    """
+
+def load_embeddings_description():
+    return
 ## main
 repos_df = load_repo_df(AppConfig.repo_representations_path)
 repos = list(repos_df["repo_name"].unique())
@@ -119,18 +128,45 @@ task_visualizations = TaskVisualizations(
     AppConfig.selected_task_counts_path,
     AppConfig.tasks_path,
 )
+display_df = pd.read_parquet("data/selected_repos_representations_umap2d.parquet")
+display_df["is_task"] = display_df["representation"] == "task"
+embedding_visualizer = EmbeddingVisualizer(display_df=display_df)
+
+
+descriptions = {
+    "intro": load_embeddings_intro_description(),
+
+    "Basic representations":  """Now we show the embeddings of tasks and repos, using various texts or representations.
+
+    The fact that selected code and/or dependency signatures (containing mostly repo's file names) are dissimilar from task names
+    should not be surprising. For our problem this illustrates the fact that these representations work poorly for retrieval.
+    """,
+    "Dependency graph based representations": """
+    Note the difference between embeddings of generated tasks and repository signatures (which contain them)
+    """,
+    "READMEs": """
+    """
+}
 
 with gr.Blocks() as demo:
+    with gr.Tab("Explore Repository Embeddings"):
+
+        tab_elems = [
+            gr.Markdown("## Tasks by area"),
+            gr.Markdown(descriptions["intro"]),
+            gr.Plot(embedding_visualizer.make_task_area_scatterplot()),
+        ]
+
+        embedding_plots = embedding_visualizer.make_embedding_plots(color_col="representation")
+        for plot_name in ["Basic representations", "Dependency graph based representations", "READMEs"]:
+            tab_elems.append(gr.Markdown(f"## {plot_name}"))
+            if descriptions.get(plot_name):
+                tab_elems.append(gr.Markdown(descriptions[plot_name]))
+            tab_elems.append(gr.Plot(embedding_plots[plot_name]))
+        gr.Column(tab_elems)
     with gr.Tab("Explore Repository Representations"):
         setup_repository_representations_tab(repos, representation_types)
     with gr.Tab("Explore PapersWithCode Tasks"):
-        task_counts_description = """
-        ## PapersWithCode Tasks Visualization
-
-        PapersWithCode tasks are grouped by area.
-
-        In addition to showing task distribution across the original dataset we display task counts in the repositories we selected.
-        """.strip()
 
         gr.Markdown(task_counts_description)
 
@@ -168,4 +204,6 @@ with gr.Blocks() as demo:
             outputs=[selected_repos_tasks_plot],
         )
 
-demo.launch()
+        gr.Plot(embedding_visualizer.make_task_area_scatterplot())
+
+demo.launch(share=True)

data/selected_repos_representations_umap2d.parquet

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

pyproject.toml

@@ -5,7 +5,15 @@ description = "Add your description here"
 readme = "README.md"
 requires-python = ">=3.10"
 dependencies = [
+    "bm25s>=0.2.3",
+    "datasets>=3.1.0",
+    "gradio>=5.5.0",
+    "llvmlite==0.41.0",
+    "nbformat>=5.10.4",
+    "plotly>=5.24.1",
     "pydantic>=2.9.2",
     "scikit-learn>=1.5.2",
+    "sentence-transformers>=3.3.1",
+    "umap-learn>=0.5.7",
     "wordcloud>=1.9.3",
 ]

text_visualization.py

@@ -4,6 +4,9 @@ import wordcloud
 from pydantic import BaseModel, Field
 import numpy as np
 import PIL
+import plotly.express as px
+import pandas as pd
+import plotly.graph_objects as go
 
 
 class WordCloudExtractor(BaseModel):
@@ -21,11 +24,11 @@ class WordCloudExtractor(BaseModel):
         """
         Extract word frequencies from a corpus using TF-IDF vectorization
         and generate word cloud frequencies.
-        
+
         Args:
             texts: List of text documents
             max_features: Maximum number of words to include
-            
+
         Returns:
             Dictionary of word frequencies suitable for WordCloud
         """
@@ -34,17 +37,79 @@ class WordCloudExtractor(BaseModel):
             max_features=max_words,
             **tfidf_params
         )
-        
+
         # Fit and transform the texts
         tfidf_matrix = tfidf.fit_transform(texts)
-        
+
         # Get feature names (words)
         feature_names = tfidf.get_feature_names_out()
-        
+
         # Calculate mean TF-IDF scores across documents
         mean_tfidf = np.array(tfidf_matrix.mean(axis=0)).flatten()
-        
+
         # Create frequency dictionary
         frequencies = dict(zip(feature_names, mean_tfidf))
-        
+
         return frequencies
+
+
+class EmbeddingVisualizer(BaseModel):
+    display_df: pd.DataFrame
+    plot_kwargs: Dict[str, Any] = Field(default_factory=lambda: dict(
+        range_x=(3, 16.5),
+        range_y=(-3, 11),
+        width=1200,
+        height=800,
+        x="x",
+        y="y",
+        template="plotly_white",
+    ))
+
+    def make_embedding_plots(self, color_col=None, hover_data=["name"], filter_df_fn=None):
+        """
+        plots Plotly scatterplot of UMAP embeddings
+        """
+        display_df = self.display_df
+        if filter_df_fn is not None:
+            display_df = filter_df_fn(display_df)
+
+        display_df = display_df.sort_values("representation", ascending=False)
+        readme_df = display_df[display_df["representation"].isin(["readme", "generated_readme", "task"])]
+        raw_df = display_df[display_df["representation"].isin(["dependency_signature", "selected_code", "task"])]
+        dependency_df = display_df[display_df["representation"].isin(["repository_signature", "dependency_signature", "generated_tasks", "task"])]
+
+        plots = [
+            self._make_task_and_repos_scatterplot(df, hover_data, color_col)
+            for df in [readme_df, raw_df, dependency_df]
+        ]
+        return dict(zip(["READMEs", "Basic representations", "Dependency graph based representations"], plots))
+
+    def _make_task_and_repos_scatterplot(self, df, hover_data, color_col):
+        # Set opacity and symbol based on is_task
+        df['size'] = df['is_task'].apply(lambda x: 0.25 if x else 0.1)
+        df['symbol'] = df['is_task'].apply(int)
+
+        combined_fig = px.scatter(
+            df,
+            hover_name="name",
+            hover_data=hover_data,
+            color=color_col,
+            color_discrete_sequence=px.colors.qualitative.Set1,
+            opacity=0.5,
+            **self.plot_kwargs
+        )
+        combined_fig.data = combined_fig.data[::-1]
+
+        return combined_fig
+
+    def make_task_area_scatterplot(self, n_areas=6):
+        display_df = self.display_df
+        displayed_tasks_df = display_df[display_df["representation"] == "task"].sort_values("representation")
+        displayed_tasks_df = displayed_tasks_df.merge(pd.read_csv("data/paperswithcode_tasks.csv"), left_on="name", right_on="task")
+        displayed_tasks_df= displayed_tasks_df[displayed_tasks_df["area"].isin(displayed_tasks_df["area"].value_counts().head(n_areas).index)]
+        tasks_fig = px.scatter(displayed_tasks_df, color="area", hover_data=["name"], opacity=0.7, **self.plot_kwargs)
+        print("N DISPLAYED TASKS", len(displayed_tasks_df))
+        return tasks_fig
+
+    class Config:
+        arbitrary_types_allowed = True