Merge pull request #31 from VectorInstitute/ft/cap_embeddings_and_scores

Capability score visualization

Author: fatemetkl

example_scripts/example_cfg/capability_score_visualization.yaml

@@ -0,0 +1,42 @@
+scientist_llm:
+  name: o4-mini
+  provider: openai
+exp_cfg:
+  exp_id: o4-mini_C100_R5_A10_T100
+# The capabilities directory should contain the LLM generated capability files.
+capabilities_cfg:
+  saved_capabilities_dir: /fs01/projects/aieng/public/ace/artifacts/gcp_artifacts/capabilities_o4-mini_C100_R5_A10_T100
+  domain: math
+  # Method used to generate capabilities
+  method: "hierarchical"
+  # Number of seed capabilities to use for initial capability generation
+  # Set to -1 to use all seed capabilities
+  num_seed_capabilities: 1
+  # Number of initial capabilities to generate using the scientist LLM
+  num_gen_capabilities: 100
+  # Buffer for capability generation
+  num_gen_capabilities_buffer: 0.2
+  # Number of capability areas to generate
+  num_capability_areas: 10
+  # Number of initial capabilities to generate per run
+  num_gen_capabilities_per_run: 5
+  # Number of tasks to generate for each capability
+  num_gen_tasks_per_capability: 100
+  # Buffer for task generation
+  num_gen_tasks_buffer: 0.2
+
+score_cfg:
+  subject_llm_names:
+    - claude-3-7-sonnet-20250219
+    - o3-mini
+    - gemini-2.0-flash
+    - o1-mini
+    - Meta-Llama-3.1-70B-Instruct
+
+  # The scores directory should contain the LLM capability scores.
+  read_score_dir: /fs01/projects/aieng/public/ace/artifacts/gcp_artifacts/scores
+  # The directory to save the capability scores plots.
+  plot_capabilities_score_dir: /fs01/projects/aieng/public/ace/artifacts/gcp_artifacts/plots
+
+defaults:
+  - _self_

example_scripts/example_cfg/train_test_embedding_visualization_cfg.yaml

@@ -1,14 +1,14 @@
 # The capabilities directory should contain the LLM generated capability files.
 capabilities_cfg:
-  saved_capabilities_dir: /fs01/projects/aieng/public/ace/artifacts/capabilities_o3-mini_C20_R3_A5_T5
+  saved_capabilities_dir: /fs01/projects/aieng/public/ace/capabilities_o4-mini_C100_R5_A10_T100
   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.9
+  filtering_similarity_threshold: 0.90
 
 dimensionality_reduction_cfg:
   # Dimensionality reduction method generates the low dimensional encodings.
@@ -19,12 +19,12 @@ dimensionality_reduction_cfg:
   normalize_output: False
 
 embedding_visualization_cfg:
-  save_dir: /fs01/projects/aieng/public/ace/artifacts/visualizations
-  plot_name: "Non-normalized PCA Embeddings"
-  show_point_ids: true # Set to true when plotting a small number of capabilities.
+  save_dir: /fs01/projects/aieng/public/acecapabilities_o4-mini_C100_R5_A10_T100/visualizations
+  plot_name: "PCA Embeddings"
+  show_point_ids: False # Set to true when plotting a small number of capabilities.
 
 heatmap_cfg:
-  save_dir: /fs01/projects/aieng/public/ace/artifacts/visualizations
+  save_dir: /fs01/projects/aieng/public/ace/capabilities_o4-mini_C100_R5_A10_T100/visualizations
   plot_name: "embedding_heatmap"
   add_squares: true
 

example_scripts/plot_llm_capability_scores.py

@@ -0,0 +1,71 @@
+import logging  # noqa: D100
+import os  # noqa: D100
+
+import hydra
+from omegaconf import DictConfig
+from tqdm import tqdm
+
+from src.generate_capabilities import (
+    get_previous_capabilities,
+    plot_capability_scores_spider_and_bar_chart,
+    select_complete_capabilities,
+)
+from src.utils.data_utils import get_run_id
+
+
+logger = logging.getLogger(__name__)
+
+
+@hydra.main(
+    version_base=None,
+    config_path="example_cfg",
+    config_name="capability_score_visualization",
+)
+def main(cfg: DictConfig) -> None:
+    """Plot capability scores across areas for each subject LLM."""
+    run_id = get_run_id(cfg)
+    # Set the base capability directory
+    capability_dir = os.path.join(
+        cfg.capabilities_cfg.saved_capabilities_dir,
+        cfg.capabilities_cfg.domain,
+    )
+
+    # Fetch previously generated capabilities
+    # Read the capabilities from the base directory
+    capabilities = get_previous_capabilities(
+        capability_dir=capability_dir,
+        score_dir_suffix=run_id,
+    )
+    capabilities = sorted(capabilities, key=lambda x: x.name)
+    logger.info(f"All capability names:\n{capabilities}")
+    # Select complete capabilities (same set of capabilities were evaluated)
+    capabilities = select_complete_capabilities(
+        capabilities=capabilities,
+        strict=False,
+        num_tasks_lower_bound=int(
+            cfg.capabilities_cfg.num_gen_tasks_per_capability
+            * (1 - cfg.capabilities_cfg.num_gen_tasks_buffer)
+        ),
+    )
+    # Sort capabilities by name
+    capabilities = sorted(capabilities, key=lambda x: x.name)
+    # Pre-load capability scores
+    for subject_llm_name in cfg.score_cfg.subject_llm_names:
+        for capability in tqdm(capabilities, desc="Loading capability scores"):
+            capability.load_scores(
+                subject_llm_name=subject_llm_name,
+            )
+
+    # Plot capability scores based on area --> spider and bar charts.
+    plot_capability_scores_spider_and_bar_chart(
+        capabilities,
+        cfg.score_cfg.subject_llm_names,
+        cfg.score_cfg.plot_capabilities_score_dir,
+        plot_name="llm_scores",
+        plot_spider_chart=True,
+        plot_grouped_bars=True,
+    )
+
+
+if __name__ == "__main__":
+    main()

example_scripts/train_test_embedding_visualization.py

@@ -37,7 +37,6 @@ def main(cfg: DictConfig) -> None:
     train_capability_dir = os.path.join(
         cfg.capabilities_cfg.saved_capabilities_dir,
         cfg.capabilities_cfg.domain,
-        "train",
     )
 
     # Fetch previously generated capabilities
@@ -65,8 +64,10 @@ def main(cfg: DictConfig) -> None:
         dim_reduction_method_name=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,
         normalize_output=cfg.dimensionality_reduction_cfg.normalize_output,
     )
+
     # Visualize the reduced embeddings
     logger.info(
         f"Visualizing {len(filtered_capabilities)} train capabilities at {cfg.embedding_visualization_cfg.save_dir}"
@@ -76,12 +77,13 @@ def main(cfg: DictConfig) -> None:
         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 + " Train",
+        plot_name=cfg.embedding_visualization_cfg.plot_name,
         show_point_ids=cfg.embedding_visualization_cfg.show_point_ids,
     )
     # Create and save the heatmap
     logger.info(
-        f"Generating heatmap for {len(filtered_capabilities)} train capabilities at {cfg.heatmap_cfg.save_dir}"
+        f"Generating heatmap for {len(filtered_capabilities)} train capabilities\
+            at {cfg.heatmap_cfg.save_dir}"
     )
     generate_capability_heatmap(
         capabilities=filtered_capabilities,
@@ -90,6 +92,22 @@ def main(cfg: DictConfig) -> None:
         plot_name=cfg.heatmap_cfg.plot_name,
         add_squares=cfg.heatmap_cfg.add_squares,
     )
+
+    _ = apply_dimensionality_reduction(
+        filtered_capabilities,
+        dim_reduction_method_name="t-sne",
+        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,
+        normalize_output=cfg.dimensionality_reduction_cfg.normalize_output,
+    )
+    plot_hierarchical_capability_2d_embeddings(
+        capabilities=filtered_capabilities,
+        dim_reduction_method="t-sne",
+        save_dir=cfg.embedding_visualization_cfg.save_dir,
+        plot_name="t-SNE Embedding",
+        show_point_ids=cfg.embedding_visualization_cfg.show_point_ids,
+    )
     # Test capabilities
     # Only PCA can be used for test capabilities.
     if cfg.dimensionality_reduction_cfg.reduce_dimensionality_method == "pca":

src/generate_capabilities.py

@@ -21,6 +21,8 @@
     filter_embeddings,
     hierarchical_2d_visualization,
     save_embedding_heatmap,
+    visualize_llm_scores_area_grouped_bar_chart,
+    visualize_llm_scores_spider_chart,
 )
 from src.model import Model
 from src.utils import constants, prompts
@@ -354,12 +356,64 @@ def generate_capabilities_using_llm(
     }
 
 
+def plot_capability_scores_spider_and_bar_chart(
+    capabilities: List[Capability],
+    subject_llm_names: List[str],
+    save_dir: str,
+    plot_name: str,
+    plot_spider_chart: bool = True,
+    plot_grouped_bars: bool = True,
+) -> None:
+    """Plot capability scores using a spider chart.
+
+    Args
+    ----
+        capabilities (List[Capability]): The list of capabilities.
+        subject_llm_names (List[str]): The names of the subject LLMs.
+        save_dir (str): The directory to save the plot.
+        plot_name (str): The name of the plot to save.
+        plot_spider_chart (bool): Whether to plot a spider chart.
+        plot_grouped_bars (bool): Whether to plot grouped bars.
+
+    """
+    # Group capabilities by area
+    llm_scores_by_area: Dict[str, Dict[str, List[float]]] = {}
+    # example: {"area1": {"llm1": [score1, score2], "llm2": [score3, score4]}} # noqa
+    for capability in capabilities:
+        if capability.area not in llm_scores_by_area:
+            llm_scores_by_area[capability.area] = {}
+        for llm_name in subject_llm_names:
+            if llm_name not in llm_scores_by_area[capability.area]:
+                llm_scores_by_area[capability.area][llm_name] = []
+            # Append the score for the capability
+            llm_scores_by_area[capability.area][llm_name].append(
+                capability.scores[llm_name]["mean"]
+            )
+    # Take the average of the scores for each area
+    # Example: {"area1": {"llm1": (mean1,std1), "llm2": (mean2,std2)}} # noqa
+    avg_llm_scores_by_area: Dict[str, Dict[str, Any]] = {}
+    for area, llm_scores in llm_scores_by_area.items():
+        avg_llm_scores_by_area[area] = {}
+        for llm_name, scores in llm_scores.items():
+            avg_llm_scores_by_area[area][llm_name] = (np.mean(scores), np.std(scores))
+
+    if plot_spider_chart:
+        visualize_llm_scores_spider_chart(
+            avg_llm_scores_by_area, save_dir, f"{plot_name}_spider_chart"
+        )
+    if plot_grouped_bars:
+        visualize_llm_scores_area_grouped_bar_chart(
+            avg_llm_scores_by_area, save_dir, f"{plot_name}_bar_chart"
+        )
+
+
 def plot_hierarchical_capability_2d_embeddings(
     capabilities: List[Capability],
     dim_reduction_method: str,
     plot_name: str,
     save_dir: str,
     show_point_ids: bool,
+    save_area_legend: bool = True,
 ) -> None:
     """Visualize the hierarchical capability embeddings.
 
@@ -374,6 +428,11 @@ def plot_hierarchical_capability_2d_embeddings(
         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.
+        save_area_legend (bool): Whether to save the area legend as a separate plot.
+
+    Returns
+    -------
+        None
     """
     # Get the reduced embeddings.
     reduced_embeddings = [
@@ -397,6 +456,7 @@ def plot_hierarchical_capability_2d_embeddings(
         save_dir=save_dir,
         plot_name=plot_name,
         points_area_name_ids=points_area_name_ids if show_point_ids else None,
+        save_area_legend=save_area_legend,
     )
 
 
@@ -568,7 +628,8 @@ def generate_and_set_capabilities_embeddings(
     texts = []
     for capability in capabilities:
         capability_dict = capability.to_dict(attribute_names=["name", "description"])
-        texts.append(f"{capability_dict['name']}: {capability_dict['description']}")
+        rep_string = f"{capability_dict['name']} - {capability.area}: {capability_dict['description']}"
+        texts.append(rep_string)
     embeddings = embedding_generator.generate_embeddings(texts)
     # Set embeddings for capabilities.
     for i, capability in enumerate(capabilities):