Merge pull request #1 from peremartra/April-1

Add bias visualization module with comprehensive metrics and utilities

Author: peremartra

.gitignore

@@ -0,0 +1,63 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+*.egg-info/
+.installed.cfg
+*.egg
+
+# Installer logs
+pip-log.txt
+pip-delete-this.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+.hypothesis/
+.pytest_cache/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# pyenv
+.python-version
+
+# mypy
+.mypy_cache/
+.dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# mkdocs build
+site/
+
+# Local files
+howto.txt
+
+# Visualization test output
+visualization_test_output/

CONTRIBUTING.md

@@ -28,6 +28,12 @@ By participating in this project, you agree to maintain a respectful and inclusi
 3. Install development dependencies:
    ```bash
    pip install -e ".[dev]"
+   
+   # For working on bias visualization
+   pip install -e ".[viz]"
+   
+   # For working on evaluation tools
+   pip install -e ".[eval]"
    ```
 4. Create a new branch for your feature or bugfix:
    ```bash
@@ -98,6 +104,8 @@ For new features:
 - Add unit tests for each function or method
 - Add integration tests for interactions between components
 - Ensure tests cover both normal behavior and error cases
+- For bias visualization features, test both the numerical computations and visualization generation
+- Mock transformer models for unit tests to avoid requiring large model downloads
 
 ## Documentation
 
@@ -114,15 +122,18 @@ Documentation is a crucial part of the project. Please follow these guidelines:
 
 3. **README**: Update the README.md if your changes affect the installation, basic usage, or other key aspects.
 
+4. **Visualization Examples**: When adding new visualization features, include visual examples in the documentation.
+
 ## Future Roadmap
 
 OptiPFair is an evolving project with plans for several future enhancements. If you're interested in contributing to these areas, please join the discussion in the related issues:
 
 1. **Attention Layer Pruning**: Implementation of structured pruning for attention mechanisms.
-2. **Bias visualisations**: Implement visualizations of bias in pair prompts. 
+2. **Bias-aware Pruning**: Techniques that optimize for both efficiency and fairness.
 3. **Block Pruning**: Methods for pruning entire transformer blocks.
 4. **Evaluation Framework**: Comprehensive evaluation suite for pruned models.
 5. **Fine-tuning Integration**: Tools for fine-tuning after pruning.
+6. **Extended Bias Analysis**: Support for intersectional and multi-attribute bias analysis.
 
 ## Questions?
 

README.md

@@ -8,18 +8,22 @@
     </h3>
 </div>
 
-A Python library for structured pruning of large language models, with a focus on GLU architectures.
+A Python library for structured pruning, and Bias visualization,  of large language models, with a focus on GLU architectures and fairness analysis.
+
 
 ## Overview
 
 OptiPFair enables efficient pruning of large language models while maintaining their performance. It implements various structured pruning methods, starting with MLP pruning for GLU architectures (as used in models like LLaMA, Mistral, etc.).
 
+
 Key features:
 - GLU architecture-aware pruning that preserves model structure 
 - Multiple neuron importance calculation methods
 - Support for both pruning percentage and target expansion rate
 - Simple Python API and CLI interface
 - Progress tracking and detailed statistics
+- **NEW**: Bias visualization tools to analyze and understand fairness issues
+
 
 ## Installation
 
@@ -62,6 +66,37 @@ print(f"Reduction: {stats['reduction']:,} parameters ({stats['percentage_reducti
 # Save pruned model
 pruned_model.save_pretrained("./pruned-llama-model")
 ```
+### Bias Visualization
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from optipfair.bias import visualize_bias
+
+# Load model and tokenizer
+model_name = "meta-llama/Llama-3.2-1B"
+model = AutoModelForCausalLM.from_pretrained(model_name)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+# Define prompt pairs that differ in demographic attributes
+prompt_pairs = [
+    ("The white man walked through the neighborhood. The police officer thought he", 
+     "The Black man walked through the neighborhood. The police officer thought he")
+]
+
+# Generate visualizations and get bias metrics
+_, metrics = visualize_bias(
+    model, 
+    tokenizer,
+    prompt_pairs=prompt_pairs,
+    visualization_types=["mean_diff", "pca"],
+    layers="first_middle_last",
+    output_dir="./bias_analysis"
+)
+
+# Print overall bias metrics
+overall = metrics["pair_1"]["metrics"]["overall_metrics"]
+print(f"Mean activation difference: {overall['mean_difference']:.6f}")
+```
 
 ### Command-Line Interface
 

completebias_test.py

@@ -0,0 +1,169 @@
+"""
+Comprehensive test script for OptiPFair bias visualization module.
+
+This script tests all three main visualization types (mean differences, heatmap, and PCA)
+with a real model to verify functionality.
+"""
+
+import os
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from optipfair.bias import (
+    visualize_mean_differences,
+    visualize_heatmap,
+    visualize_pca,
+    visualize_bias
+)
+
+# Create output directory for saved visualizations
+OUTPUT_DIR = "./visualization_test_output"
+os.makedirs(OUTPUT_DIR, exist_ok=True)
+
+def test_bias_visualizations():
+    """Test all bias visualization functions with a real model."""
+    print("Loading model and tokenizer...")
+    try:
+        # Load a small model for testing - use any model you have access to
+        model_name = "meta-llama/Llama-3.2-1B"  # Change to any model you have access to
+        model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            torch_dtype=torch.float16,  # Use half precision for memory efficiency
+            device_map="auto"  # Let the library decide the best device allocation
+        )
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        
+        # Define test prompt pairs
+        prompt_pairs = [
+            (
+                "The white doctor examined the patient. The nurse thought",
+                "The Black doctor examined the patient. The nurse thought"
+            )
+        ]
+        prompt_pair = prompt_pairs[0]  # Use the first pair for individual tests
+        
+        print("\n" + "="*50)
+        print("TESTING INDIVIDUAL VISUALIZATION FUNCTIONS")
+        print("="*50)
+        
+        # Test 1: Mean Differences Visualization
+        print("\nTesting mean differences visualization...")
+        try:
+            visualize_mean_differences(
+                model,
+                tokenizer,
+                prompt_pair,
+                layer_type="mlp_output",
+                layers="first_middle_last",
+                output_dir=os.path.join(OUTPUT_DIR, "mean_diff"),
+                figure_format="png"
+            )
+            print("✅ Mean differences visualization successful!")
+        except Exception as e:
+            print(f"❌ Mean differences visualization failed: {e}")
+        
+        # Test 2: Heatmap Visualization
+        # For heatmap we need a specific layer, so let's get layer 8 or the middle layer
+        print("\nTesting heatmap visualization...")
+        try:
+            visualize_heatmap(
+                model,
+                tokenizer,
+                prompt_pair,
+                layer_key="mlp_output_layer_8",  # Using middle layer - adjust if your model has fewer layers
+                output_dir=os.path.join(OUTPUT_DIR, "heatmap"),
+                figure_format="png"
+            )
+            print("✅ Heatmap visualization successful!")
+        except Exception as e:
+            print(f"❌ Heatmap visualization failed: {e}")
+            # If the specific layer fails, try with layer 0 which should exist in any model
+            print("Retrying with layer 0...")
+            try:
+                visualize_heatmap(
+                    model,
+                    tokenizer,
+                    prompt_pair,
+                    layer_key="mlp_output_layer_0",
+                    output_dir=os.path.join(OUTPUT_DIR, "heatmap"),
+                    figure_format="png"
+                )
+                print("✅ Heatmap visualization with layer 0 successful!")
+            except Exception as e2:
+                print(f"❌ Heatmap visualization with layer 0 also failed: {e2}")
+        
+        # Test 3: PCA Visualization
+        print("\nTesting PCA visualization...")
+        try:
+            visualize_pca(
+                model,
+                tokenizer,
+                prompt_pair,
+                layer_key="attention_output_layer_8",  # Using middle attention layer
+                highlight_diff=True,
+                output_dir=os.path.join(OUTPUT_DIR, "pca"),
+                figure_format="png"
+            )
+            print("✅ PCA visualization successful!")
+        except Exception as e:
+            print(f"❌ PCA visualization failed: {e}")
+            # If the specific layer fails, try with layer 0
+            print("Retrying with layer 0...")
+            try:
+                visualize_pca(
+                    model,
+                    tokenizer,
+                    prompt_pair,
+                    layer_key="attention_output_layer_0",
+                    highlight_diff=True,
+                    output_dir=os.path.join(OUTPUT_DIR, "pca"),
+                    figure_format="png"
+                )
+                print("✅ PCA visualization with layer 0 successful!")
+            except Exception as e2:
+                print(f"❌ PCA visualization with layer 0 also failed: {e2}")
+        
+        # Test 4: Main visualize_bias function (combines all visualization types)
+        print("\n" + "="*50)
+        print("TESTING MAIN VISUALIZATION FUNCTION")
+        print("="*50)
+        
+        print("\nTesting visualize_bias function...")
+        try:
+            _, metrics = visualize_bias(
+                model,
+                tokenizer,
+                prompt_pairs=prompt_pairs,
+                visualization_types=["mean_diff", "heatmap", "pca"],
+                layers="first_middle_last",
+                output_dir=os.path.join(OUTPUT_DIR, "combined"),
+                figure_format="png",
+                show_progress=True
+            )
+            
+            print("✅ visualize_bias function successful!")
+            
+            # Print some metrics to verify they're being calculated correctly
+            if metrics and "pair_1" in metrics:
+                overall = metrics["pair_1"]["metrics"]["overall_metrics"]
+                print("\nMetrics sample:")
+                print(f"  Overall mean difference: {overall['mean_difference']:.6f}")
+                print(f"  Max difference: {overall['max_difference']:.6f}")
+                
+                # Check if we have component metrics
+                if "component_metrics" in metrics["pair_1"]["metrics"]:
+                    comp_metrics = metrics["pair_1"]["metrics"]["component_metrics"]
+                    for comp_name, comp_data in comp_metrics.items():
+                        if comp_name in ["mlp_output", "attention_output"]:
+                            print(f"  {comp_name} mean difference: {comp_data['mean_difference']:.6f}")
+        except Exception as e:
+            print(f"❌ visualize_bias function failed: {e}")
+        
+        print("\nTests completed. Check the output directory for visualization results:")
+        print(f"  {os.path.abspath(OUTPUT_DIR)}")
+        
+    except Exception as e:
+        print(f"Failed to load model: {e}")
+        print("Please make sure you have access to the specified model and that your environment is set up correctly.")
+
+if __name__ == "__main__":
+    test_bias_visualizations()