Merge pull request #17 from VectorInstitute/ft/test-capability-embeddings

t-SNE visualization scripts for embedding validation

Author: fatemetkl

.gitignore

@@ -4,6 +4,9 @@ __pycache__/
 *$py.class
 .vscode/
 
+# macOS system files
+.DS_Store
+
 # C extensions
 *.so
 

example_scripts/__init__.py

@@ -0,0 +1,5 @@
+"""
+The __init__.py file for example scripts.
+
+It initializes the example scripts module.
+"""

example_scripts/example_cfg/visualization_cfg.yaml

@@ -0,0 +1,30 @@
+# The capabilities directory should contain the LLM generated capability files.
+capabilities_cfg:
+  saved_capabilities_dir: /fs01/projects/aieng/public/ace/artifacts/capabilities_gpt-4o-mini_T20_R4_A5
+  domain: math
+
+embedding_cfg:
+  # The embedding model name used to generate capability embeddings used for filtering.
+  embedding_model: "text-embedding-3-small" # "text-embedding-3-small" or "text-embedding-3-large"
+  embedding_size: 512
+  # The cosine similarity threshold for filtering capabilities based on their embeddings.
+  filtering_similarity_threshold: 0.85
+
+dimensionality_reduction_cfg:
+  # Dimensionality reduction method generates the low dimensional encodings.
+  reduce_dimensionality_method: "t-sne" # "t-sne" or "cut-embedding".
+  reduced_dimensionality_size: 2
+  tsne_perplexity: 8 # Choose this hyperparameter based on the number of capabilities you have.
+
+embedding_visualization_cfg:
+  save_dir: /fs01/projects/aieng/public/ace/artifacts/visualizations
+  plot_name: "tsne_plot"
+  show_point_ids: true # Set to true when plotting a small number of capabilities.
+
+heatmap_cfg:
+  save_dir: /fs01/projects/aieng/public/ace/artifacts/visualizations
+  plot_name: "heatmap"
+  add_squares: true
+
+defaults:
+  - _self_

example_scripts/visualize_embeddings.py

@@ -0,0 +1,94 @@
+"""The visualize_embeddings.py script is used to visualize LLM generated capabilities.
+
+It generates a 2D plot of the capabilities' embeddings using the specified
+dimensionality reduction method. It also generates a heatmap of the capabilities
+similarity matrix. Similarity is computed using the original embeddings
+generated with the embedding model defined in the config file.
+"""
+
+import os  # noqa: D100
+
+import hydra
+from omegaconf import DictConfig
+
+from src.generate_capabilities import (
+    apply_dimensionality_reduction,
+    filter_capabilities,
+    generate_and_set_capabilities_embeddings,
+    generate_capability_heatmap,
+    get_previous_capabilities,
+    plot_hierarchical_capability_2d_embeddings,
+)
+
+
+@hydra.main(
+    version_base=None, config_path="example_cfg", config_name="visualization_cfg"
+)
+def main(cfg: DictConfig) -> None:
+    """
+    Run the visualization script with the specified configuration.
+
+    Args:
+        cfg (DictConfig): Configuration for the script.
+    """
+    # Load capabilities from the specified directory
+    # Set the base capability directory
+    base_capability_dir = os.path.join(
+        cfg.capabilities_cfg.saved_capabilities_dir,
+        cfg.capabilities_cfg.domain,
+    )
+    os.makedirs(base_capability_dir, exist_ok=True)
+
+    # Fetch previously generated capabilities
+    capabilities = get_previous_capabilities(capability_dir=base_capability_dir)
+    # Assert that the capabilities list is not empty
+    print(f"Loaded {len(capabilities)} capabilities from {base_capability_dir}")
+    assert capabilities, "No capabilities found in the specified directory."
+
+    # Embed capabilities using openai embedding model
+    generate_and_set_capabilities_embeddings(
+        capabilities=capabilities,
+        embedding_model_name=cfg.embedding_cfg.embedding_model,
+        embed_dimensions=cfg.embedding_cfg.embedding_size,
+    )
+    # Filter capabilities based on their embeddings
+    filtered_capabilities = filter_capabilities(
+        capabilities,
+        embedding_model_name=cfg.embedding_cfg.embedding_model,
+        similarity_threshold=cfg.embedding_cfg.filtering_similarity_threshold,
+    )
+    # Reduce the dimensionality of capability embeddings generated by the
+    # embedding model.
+    apply_dimensionality_reduction(
+        filtered_capabilities,
+        dim_reduction_method=cfg.dimensionality_reduction_cfg.reduce_dimensionality_method,
+        output_dimension_size=cfg.dimensionality_reduction_cfg.reduced_dimensionality_size,
+        embedding_model_name=cfg.embedding_cfg.embedding_model,
+        tsne_perplexity=cfg.dimensionality_reduction_cfg.tsne_perplexity,
+    )
+    # Visualize the reduced embeddings
+    print(
+        f"Visualizing {len(filtered_capabilities)} capabilities at {cfg.embedding_visualization_cfg.save_dir}"
+    )
+    plot_hierarchical_capability_2d_embeddings(
+        capabilities=filtered_capabilities,
+        dim_reduction_method=cfg.dimensionality_reduction_cfg.reduce_dimensionality_method,
+        save_dir=cfg.embedding_visualization_cfg.save_dir,
+        plot_name=cfg.embedding_visualization_cfg.plot_name,
+        show_point_ids=cfg.embedding_visualization_cfg.show_point_ids,
+    )
+    # Create and save the heatmap
+    print(
+        f"Generating heatmap for {len(filtered_capabilities)} capabilities at {cfg.heatmap_cfg.save_dir}"
+    )
+    generate_capability_heatmap(
+        capabilities=filtered_capabilities,
+        embedding_model_name=cfg.embedding_cfg.embedding_model,  # Using the original embeddings, not the reduced version.
+        save_dir=cfg.heatmap_cfg.save_dir,
+        plot_name=cfg.heatmap_cfg.plot_name,
+        add_squares=cfg.heatmap_cfg.add_squares,
+    )
+
+
+if __name__ == "__main__":
+    main()

src/capability.py

@@ -127,7 +127,7 @@ class Capability:
         Loads the capability configuration from a JSON file.
     _load_capability_repr_class() -> None
         Loads the capability representation class from a Python file.
-    _to_dict() -> Dict[str, Any]
+    to_dict() -> Dict[str, Any]
         Converts the capability attributes to a dictionary.
     to_json_str() -> str
         Converts the capability to a JSON string.
@@ -509,7 +509,20 @@ def add_and_update_tasks(
         self._load_capability_json()
         self._load_capability_repr_class()
 
-    def _to_dict(self, attribute_names: List[str] | None = None) -> Dict[str, Any]:
+    def to_dict(self, attribute_names: List[str] | None = None) -> Dict[str, Any]:
+        """
+        Return a dictionary of the capability attributes.
+
+        Args:
+            attribute_names (List[str] | None, optional): the list of attribute
+            names requested. If none, return a set of default attributes.
+            Defaults to None.
+
+        Returns
+        -------
+            Dict[str, Any]: a dictionary representation of the capability
+                based on the requested attribute names or a default set of attributes.
+        """
         if attribute_names is None:
             return {
                 "name": self.name,
@@ -521,6 +534,22 @@ def _to_dict(self, attribute_names: List[str] | None = None) -> Dict[str, Any]:
             attr: getattr(self, attr) for attr in attribute_names if hasattr(self, attr)
         }
 
+    def get_attribute(self, attribute_name: str) -> Any:
+        """
+        Get the value of a specific attribute of the capability.
+
+        Args
+        ----
+            attribute_name (str): The name of the attribute to retrieve.
+
+        Returns
+        -------
+            Any: The value of the specified attribute.
+        """
+        if not hasattr(self, attribute_name):
+            raise AttributeError(f"Attribute {attribute_name} not found in capability.")
+        return getattr(self, attribute_name)
+
     def to_json_str(self, attribute_names: List[str] | None = None) -> str:
         """
         Convert the capability to a JSON string.
@@ -538,7 +567,7 @@ def to_json_str(self, attribute_names: List[str] | None = None) -> str:
             # If only the name is requested, return the name directly
             repr_str = self.name
         else:
-            repr_str = json.dumps(self._to_dict(attribute_names), indent=4)
+            repr_str = json.dumps(self.to_dict(attribute_names), indent=4)
         return str(repr_str)
 
     def __str__(self) -> str:

src/generate_capabilities.py

@@ -6,6 +6,7 @@
 from typing import Any, Dict, List, Optional
 
 import numpy as np
+import torch
 from langsmith import tracing_context
 from tenacity import Retrying, stop_after_attempt
 
@@ -15,7 +16,9 @@
     EmbeddingGenerator,
     EmbeddingModelName,
     filter_embeddings,
+    hierarchical_2d_visualization,
     reduce_embeddings_dimensions,
+    save_embedding_heatmap,
 )
 from src.model import Model
 from src.utils import constants, prompts
@@ -145,7 +148,7 @@ def get_capability_repr_with_score(capability: Capability, model_name: str) -> s
         str: A JSON string containing the capability JSON string and score.
     """
     model_score = capability.load_scores()[model_name]
-    capability_dict = capability._to_dict()
+    capability_dict = capability.to_dict()
     capability_dict["score"] = model_score
     return json.dumps(capability_dict, indent=4)
 
@@ -340,12 +343,106 @@ def generate_capabilities_using_llm(
     }
 
 
+def plot_hierarchical_capability_2d_embeddings(
+    capabilities: List[Capability],
+    dim_reduction_method: str,
+    plot_name: str,
+    save_dir: str,
+    show_point_ids: bool,
+) -> None:
+    """Visualize the hierarchical capability embeddings.
+
+    Embeddings are retrieved based on the defined dim_reduction_method,
+    and they should be 2D.
+
+    Args
+    ----
+        capabilities (List[Capability]): The list of capabilities.
+        dim_reduction_method (str): The dimensionality reduction method to use.
+        plot_name (str): The name of the plot to save.
+        save_dir (str): The directory to save the plot.
+        show_point_ids (bool): Whether to show point IDs in the plot. Set this to
+            False for large datasets to avoid cluttering the plot.
+
+    Returns
+    -------
+        None
+    """
+    # Get the reduced embeddings.
+    reduced_embeddings = [
+        capability.get_embedding(dim_reduction_method) for capability in capabilities
+    ]
+    area_names = [capability.get_attribute("area") for capability in capabilities]
+
+    # Populate embeddings_by_area, and points_area_name_ids
+    embeddings_by_area: dict[str, List[torch.Tensor]] = {}
+    points_area_name_ids: dict[str, dict[str, int]] = {}
+    for idx in range(len(reduced_embeddings)):
+        area_name = area_names[idx]
+        if area_name not in embeddings_by_area:
+            embeddings_by_area[area_name] = []
+            points_area_name_ids[area_name] = {}
+        embeddings_by_area[area_name].append(reduced_embeddings[idx])
+        points_area_name_ids[area_name][capabilities[idx].name] = idx
+
+    hierarchical_2d_visualization(
+        embeddings_by_area=embeddings_by_area,
+        save_dir=save_dir,
+        plot_name=plot_name,
+        points_area_name_ids=points_area_name_ids if show_point_ids else None,
+    )
+
+
+def generate_capability_heatmap(
+    capabilities: List[Capability],
+    embedding_model_name: str,
+    plot_name: str,
+    save_dir: str,
+    add_squares: bool,
+) -> None:
+    """
+    Generate and save a heatmap of the capabilities based on their embeddings.
+
+    Args:
+        capabilities (List[Capability]): the list of capabilities.
+        embedding_model_name (str): name of the embedding model used
+            to generate the embeddings.
+        plot_name (str): name of the plot file to save.
+        save_dir (str): directory to save the plot.
+        add_squares (bool): whether to add squares to the heatmap.
+    """
+    # Get the embeddings based on the specified embedding model name.
+    embeddings = [
+        capability.get_embedding(embedding_model_name) for capability in capabilities
+    ]
+    # Process capabilities to populate embeddings_by_area and
+    # capability_names_by_area.
+    area_names = [capability.area for capability in capabilities]
+    embeddings_by_area: dict[str, List[torch.Tensor]] = {}
+    capability_names_by_area: dict[str, List[str]] = {}
+    for idx in range(len(capabilities)):
+        embedding_group = area_names[idx]
+        if embedding_group not in embeddings_by_area:
+            embeddings_by_area[embedding_group] = []
+            capability_names_by_area[embedding_group] = []
+        embeddings_by_area[embedding_group].append(embeddings[idx])
+        capability_names_by_area[embedding_group].append(capabilities[idx].name)
+
+    save_embedding_heatmap(
+        embeddings_by_area=embeddings_by_area,
+        capability_names_by_area=capability_names_by_area,
+        save_dir=save_dir,
+        plot_name=plot_name,
+        add_squares=add_squares,
+    )
+
+
 def apply_dimensionality_reduction(
     capabilities: List[Capability],
     dim_reduction_method: str,
     output_dimension_size: int,
     embedding_model_name: str,
-    seed: int = 42,
+    tsne_perplexity: int,
 ) -> None:  # noqa: D205
     """Apply dimensionality reduction to the capabilities.
 
@@ -391,7 +488,7 @@ def apply_dimensionality_reduction(
         embeddings,
         output_dimensions=output_dimension_size,
         dim_reduction_technique=DimensionalityReductionTechnique(dim_reduction_method),
-        seed=seed,
+        perplexity=tsne_perplexity,
     )
     # Set the reduced embeddings for each capability.
     for capability, reduced_embedding in zip(capabilities, reduced_embeddings):
@@ -425,12 +522,12 @@ def generate_and_set_capabilities_embeddings(
         embed_dimensions=embed_dimensions,
     )
     # Generate embeddings for the capabilities, all at the same time.
-    embeddings = embedding_generator.generate_embeddings(
-        texts=[
-            capability.to_json_str(attribute_names=["name", "description", "domain"])
-            for capability in capabilities
-        ]
-    )
+    # Embeddings are generated based on the capability name and description.
+    texts = []
+    for capability in capabilities:
+        capability_dict = capability.to_dict(attribute_names=["name", "description"])
+        texts.append(f"{capability_dict['name']}: {capability_dict['description']}")
+    embeddings = embedding_generator.generate_embeddings(texts)
     # Set embeddings for capabilities.
     for i, capability in enumerate(capabilities):
         capability.set_embedding(