# OptiPFair

<div align="center">
  <img src="images/optiPfair.png" alt="OptiPFair Logo" width="600"/>
</div>

<div align="center">
  <h1>optiPfair</h1>
  <strong>The Python library for making LLMs both efficient (via pruning) and fair (via bias analysis).</strong>
</div>

<p align="center">
  <a href="https://pypi.org/project/optipfair/"><img alt="PyPI Version" src="https://img.shields.io/pypi/v/optipfair?color=blue"></a>
  <a href="https://github.com/peremartra/optipfair/blob/main/LICENSE"><img alt="License" src="https://img.shields.io/github/license/peremartra/optipfair?color=green"></a>
  <a href="https://github.com/peremartra/optipfair/stargazers"><img alt="GitHub Stars" src="https://img.shields.io/github/stars/peremartra/optipfair?style=social"></a>
</p>

<div align="center">
    <h3>
        <a href="https://peremartra.github.io/optipfair/" target="_blank">Documentation</a>
        ·
        <a href="https://github.com/peremartra/optipfair/issues" target="_blank">Report Bug</a>
        ·
        <a href="https://github.com/peremartra/optipfair/issues" target="_blank">Request Feature</a>
    </h3>
</div>

---

### 🚀 Interactive Demos: Try OptiPFair NOW

Experience the power of OptiPFair directly in your browser. No installation, no setup.

| Live Bias Visualization (Recommended Demo) | Pruning Compatibility Check | Bias Compatibility Check |
| :---: | :---: | :---: |
| Analyze any compatible model from Hugging Face with a full UI. | Check if your model's architecture can be pruned by OptiPFair. | The coder's alternative to our live demo. |
| **[🚀 Launch the Live Demo on HF Spaces](https://huggingface.co/spaces/oopere/optipfair-bias-analyzer)** | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/peremartra/optipfair/blob/main/examples/pruning_compatibility_check.ipynb) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/peremartra/optipfair/blob/main/examples/bias_compatibility_check.ipynb) |

---
### ✅ Why OptiPFair?

OptiPFair is more than just another pruning library. It's a toolkit designed for the modern AI developer who cares about both performance and responsibility.

* **Efficiency & Fairness in One Place**: Stop juggling tools. OptiPFair is the only library designed to integrate structured pruning with powerful, intuitive bias visualization and analysis.

* **Optimized for Modern Architectures**: We focus on what works now. The library is specialized for GLU-based models like LLaMA, Mistral, Gemma, and Qwen, ensuring relevant and effective pruning.

* **Go Beyond Numbers with Bias Visualization**: Don't just get a bias score. Our visualization tools (PCA, heatmaps, mean differences) help you *understand* how and where your model encodes bias, enabling more effective mitigation.

* **🤖 AI-Assisted Development**: Accelerate your workflow using the included [`LLM Reference Manual`](https://github.com/peremartra/optipfair/blob/main/optipfair_llm_reference_manual.txt). Provide it to your favorite LLM (ChatGPT, Claude) to get expert-level help and generate integration code instantly.
<p align="center">
      <img src="images/optipFair_llmflow.gif" alt="AI Pair Programming with OptiPFair" width="500"/>
</p>

* **🔬 Backed by Research**: Our methods aren't arbitrary. They are built upon and validated by ongoing applied research in model optimization and fairness analysis.

---
### ⚙️ Installation
Choose the installation method that best suits your needs. For bias visualization features, you'll need the [viz] extra.
**Standard Installation**
For core pruning functionality:
```python
pip install optipfair
```

**Full Installation (with Bias Visualization)**
To use the bias analysis and visualization tools, install with the [viz] extra dependencies:
```python
pip install "optipfair[viz]"
```

**Developer Installation**
To install from the source for contributing or development:
```bash
git clone https://github.com/peremartra/optipfair.git
cd optipfair
pip install -e .
```
---
## ⚡ Quick Start

See how to use OptiPFair's core features in just a few lines of code.

### Pruning with the Python API

Prune 20% of the MLP neurons from a model using the Maximum Absolute Weight (MAW) method.

```python
from transformers import AutoModelForCausalLM
from optipfair import prune_model

# Load a pre-trained model
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B")

# Prune 20% of neurons from MLP layers
pruned_model, stats = prune_model(
    model=model,
    pruning_type="MLP_GLU",
    neuron_selection_method="MAW",
    pruning_percentage=20,
    show_progress=True,
    return_stats=True
)

# Print pruning statistics
print(f"Original parameters: {stats['original_parameters']:,}")
print(f"Pruned parameters: {stats['pruned_parameters']:,}")
print(f"Reduction: {stats['reduction']:,} parameters ({stats['percentage_reduction']:.2f}%)")

# Save the pruned model
pruned_model.save_pretrained("./pruned-llama-model")
```
The pruning process yields tangible results in model size and performance. Here's a sample comparison for **Llama-3.2-1B** after pruning 20% of its MLP neurons:

| Metric | Original Model | Pruned Model | Improvement |
| :--- | :---: | :---: | :---: |
| **Total Parameters** | 1.24B | 1.07B | **-13.03%** |
| **Inference Speed** | *Benchmark in progress* | *Benchmark in progress* | *Coming soon* |
| **MMLU Score** | *Benchmark in progress* | *Benchmark in progress* | *Minimal change expected* |

*Results based on the [MAW pruning method](#neuron-selection-methods). Full benchmark results will be published shortly.*

### Visualizing Bias with the Python API
Generate visualizations to analyze how a model's activations differ when processing prompts with varying demographic attributes.

```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}")
```
The code above generates the following types of visualizations, allowing for a deep dive into model fairness:
| ![Mean Image Differences](/images/mean_image_differences.png) | ![activation_differences_layer.png](/images/activation_differences_layer.png) | ![pca_analysis.png](/images/pca_analysis.png)
| ---| --- | --- |

---
## Dive Deeper: Features & Concepts

For those who want to understand the nuts and bolts of OptiPFair.

### Supported Models
OptiPFair is designed to work with transformer-based language models that use GLU architecture in their MLP layers. This includes most modern LLMs:

* **LLaMA family** (LLaMA, LLaMA-2, LLaMA-3)
* **Mistral** models
* **Gemma** models
* **Qwen** models
* **Phi** models
* ... and other models with a similar GLU architecture.

### Neuron Selection Methods
OptiPFair provides multiple methods for calculating neuron importance, allowing you to tailor the pruning strategy to your specific needs:

1.  **MAW (Maximum Absolute Weight)**: The default and typically most effective method for GLU architectures. It identifies influential neurons based on the magnitude of their connections.
2.  **VOW (Variance of Weights)**: Identifies neurons based on the variance of their weights.
3.  **PON (Product of Norms)**: Uses the product of L1 norms to identify important neurons.

### Expansion Rate vs. Pruning Percentage
You can define the pruning target in two intuitive ways:

1.  **Pruning Percentage**: Directly specify the percentage of neurons to remove (e.g., `20%`).
2.  **Expansion Rate**: Specify the target MLP expansion rate as a percentage (e.g., target a `140%` expansion rate, down from a model's original `400%`). This is often more comparable across different model sizes.

---

## 🗺️ Roadmap & Community

The OptiPFair project is actively developed. Here's what's planned for the future.

### Future Roadmap
Our goal is to make OptiPFair the go-to toolkit for efficient and fair model optimization. Key upcoming features include:

* **Attention Pruning**: Implementing pruning methods for attention heads and layers. 
* **Depth Pruning**: Adding support for removing entire layers or blocks. 
* **Advanced Benchmarks**: Integrating more comprehensive performance and evaluation benchmarks.
* **GPU Optimizations**: Creating a v2.0 with significant GPU-specific optimizations for faster execution. 
* **Large-Scale Model Support**: Adding compatibility for DeepSpeed and FSDP to handle 70B+ models efficiently. 

### 🤝 Contributing
Contributions are welcome! Whether it's bug reports, feature requests, or code contributions, please check out our [contributing guidelines](CONTRIBUTING.md) to get started.

### Citation
If you use OptiPFair in your research or projects, please cite the library:

```bibtex
@software{optipfair,
  author = {Pere Martra},
  title = {OptiPFair: A Library for Structured Pruning and Bias Visualization of Large Language Models},
  year = {2024},
  url = {https://github.com/peremartra/optipfair}
}
```
### License
This project is licensed under the Apache 2.0 License. See the LICENSE file for details.
