Embedding generation and visualization README (#36)

* Added example script README

* Added example script README

* small fix

Author: fatemetkl

example_scripts/README.md

@@ -0,0 +1,141 @@
+
+## `train_test_embedding_visualization` example
+
+Here we describe the steps required for reading and selecting pre-generated capabilities and their tasks, generating capability embeddings, filtering capabilities based on those embeddings, reducing dimensionality, and visualizing capabilities. All of these steps are implemented in the `train_test_embedding_visualization.py` script, which runs the process for both `train` and `test` capabilities. The directory containing the `train` and `test` capabilities and tasks is specified in the `train_test_embedding_visualization_cfg.yaml` file.
+
+You can also find the steps for loading and visualizing LLM scores in `plot_llm_capability_scores.py`. The scores can be plotted using a spider chart or a bar chart via the `plot_capability_scores_spider_and_bar_chart()` function.
+
+
+Step 1: Read the already generated and saved train capabilities:
+
+```python
+    # Read the capabilities from the base directory
+    train_capability_dir = os.path.join(
+        cfg.capabilities_cfg.saved_capabilities_dir,
+        cfg.capabilities_cfg.domain,
+    )
+    # Fetch previously generated capabilities
+    capabilities = get_previous_capabilities(capability_dir=train_capability_dir)
+
+```
+
+Step 2: Sort and keep complete capabilities. Complete capabilities have enough verified tasks generated for them.
+
+```python
+
+    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)
+        ),
+    )
+    capabilities = sorted(capabilities, key=lambda x: x.name)
+
+  ```
+
+
+
+Step 3: Generate capability embeddings using openai model, and assign embeddings to each capability object.
+
+```python   
+    # 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,
+    )
+```
+
+Step 4: Filter capabilities based on the embeddings such that if embeddings are too similar according to a threshold, one of them (capability) should be removed.
+
+```python
+    # 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,
+    )
+```
+
+Step 5: Capability embedding dimensionality reduction. 
+
+```python
+    # Reduce the dimensionality of capability embeddings generated by the
+    # embedding model.
+    dim_reduction = apply_dimensionality_reduction(
+        filtered_capabilities,
+        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,
+    )
+```
+
+
+Step 6: Visualize the reduced embeddings.
+
+```python
+    # Plot training capabilities
+    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,
+    )
+```
+
+Step 7: Capability Heatmap
+
+```python
+    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,
+    )
+```
+
+
+Step 8: **Test capabilities** are also loaded and their embeddings are generated using openai embedding model just like the previous steps. The only difference is that we should use the already fitted PCA model for dimensionality reduction. Also, we visualize test capabilities 2D embeddings with train embeddings to see their relative distance.
+
+```python
+        # Use the fitted PCA dim reduction to transform the test capabilities
+        apply_dimensionality_reduction_to_test_capabilities(
+            test_capabilities,
+            dim_reduction_method=dim_reduction,
+            embedding_model_name=cfg.embedding_cfg.embedding_model,
+        )
+```
+
+Step 9: Visualize train and test capability embeddings together.
+
+```python
+        all_capabilities = filtered_capabilities + test_capabilities
+        logger.info(
+            f"Visualizing {len(all_capabilities)} train and test capabilities at {cfg.embedding_visualization_cfg.save_dir}"
+        )
+        plot_hierarchical_capability_2d_embeddings(
+            capabilities=all_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 and Test",
+            show_point_ids=cfg.embedding_visualization_cfg.show_point_ids,
+        )
+```
+
+
+### How are embeddings generated?
+
+The `generate_and_set_capabilities_embeddings()` function in `src/utils/embedding_utils.py` handles this process. Capability name and descriptions are extracted to form the representation string `rep_string`. Then, embeddings are generated using the OpenAI embedding model via `embedding_generator`. Finally, the embeddings are assigned to each capability object.
+The representation string was chosen based on visualization-based experiments and is defined as:
+
+```python 
+        rep_string = f"{capability_dict['name']} - {capability.area}: {capability_dict['description']}"
+```

example_scripts/example_cfg/train_test_embedding_visualization_cfg.yaml

@@ -19,7 +19,7 @@ dimensionality_reduction_cfg:
   normalize_output: False
 
 embedding_visualization_cfg:
-  save_dir: /fs01/projects/aieng/public/acecapabilities_o4-mini_C100_R5_A10_T100/visualizations
+  save_dir: /fs01/projects/aieng/public/ace/capabilities_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.