optipfair 0.2.2__tar.gz → 0.2.4__tar.gz

{optipfair-0.2.2/optipfair.egg-info → optipfair-0.2.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: optipfair
-Version: 0.2.2
+Version: 0.2.4
 Summary: A library for structured pruning & Bias visualization of large language models
 Home-page: https://github.com/peremartra/optipfair
 Author: Pere Martra
@@ -80,6 +80,8 @@ Dynamic: summary
 
 ---
 > **New to optipfair?** Use our [LLM Reference Manual](optipfair_llm_reference_manual.txt) - paste it into ChatGPT, Claude or your Favourite LLM for guided assistance with any optipfair task.
+
+> **Note on Terminology:** The default neuron selection method is **PPM (Peak-to-Peak Magnitude)**, which calculates neuron importance based on the full dynamic range of weights (max + |min|). This method is formally described in: *Martra, P. (2025). Fragile Knowledge, Robust Instruction-Following: The Width Pruning Dichotomy in Llama-3.2. ArXiv. https://arxiv.org/abs/2512.22671*. For backward compatibility, the parameter value `"MAW"` is still accepted and maps to PPM.
 ### 🚀 Interactive Demos: Try optipfair NOW
 
 Experience optipfair's capabilities directly in your browser.
@@ -149,17 +151,17 @@ 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.
+Prune 20% of the MLP neurons from a model using the Peak-to-Peak Magnitude (PPM) method.
 
 ```python
 from transformers import AutoModelForCausalLM
-from optipfair import prune_model
+import optipfair as opf
 
 # 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(
+pruned_model, stats = opf.prune_model(
     model=model,
     pruning_type="MLP_GLU",
     neuron_selection_method="MAW",
@@ -185,7 +187,7 @@ The pruning process yields tangible results in model size and performance. Here'
 | **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.*
+*Results based on the [PPM pruning method](#neuron-selection-methods) (parameter `"MAW"`). Full benchmark results will be published shortly.*
 
 ### Data-Driven Width Pruning (NEW in v0.2.0)
 
@@ -194,7 +196,7 @@ Enhance pruning decisions with activation statistics from calibration data. This
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from torch.utils.data import DataLoader, TensorDataset
 import torch
-from optipfair import prune_model
+import optipfair as opf
 
 # Load model and tokenizer
 model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B")
@@ -213,9 +215,9 @@ dataset = TensorDataset(inputs['input_ids'], inputs['attention_mask'])
 dataloader = DataLoader(dataset, batch_size=8)
 
 # Prune with data-driven importance calculation
-pruned_model, stats = prune_model(
+pruned_model, stats = opf.prune_model(
     model=model,
-    neuron_selection_method="MAW",  # Only MAW supports data-driven pruning
+    neuron_selection_method="MAW",  # Only PPM (parameter "MAW") supports data-driven pruning
     pruning_percentage=20,
     dataloader=dataloader,  # ← Enables hybrid pruning
     show_progress=True,
@@ -232,7 +234,7 @@ pruned_model.save_pretrained("./pruned-datadriven-model")
 - 🔬 **Research-Backed**: Based on CFSP methodology (arXiv:2409.13199v2)
 - ⚡ **Easy Integration**: Just add a dataloader - no other changes needed
 
-**Note:** Data-driven pruning is currently only available with `neuron_selection_method="MAW"`. Using a dataloader with "VOW" or "PON" will raise a `ValueError`.
+**Note:** Data-driven pruning is currently only available with `neuron_selection_method="MAW"` (PPM method). Using a dataloader with "VOW" or "PON" will raise a `ValueError`.
 
 ### Selective Layer Width Pruning (NEW in v0.2.0)
 
@@ -240,13 +242,13 @@ Prune neurons only in specific layers while leaving others unchanged. Perfect fo
 
 ```python
 from transformers import AutoModelForCausalLM
-from optipfair import prune_model
+import optipfair as opf
 
 # Load a pre-trained model
 model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B")
 
 # Prune neurons only in specific layers (e.g., middle layers)
-pruned_model, stats = prune_model(
+pruned_model, stats = opf.prune_model(
     model=model,
     pruning_type="MLP_GLU",
     neuron_selection_method="MAW",
@@ -282,13 +284,13 @@ The `expansion_divisor` parameter ensures that intermediate layer sizes are divi
 
 ```python
 from transformers import AutoModelForCausalLM
-from optipfair import prune_model
+import optipfair as opf
 
 # Load model
 model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B")
 
 # Prune with hardware optimization
-pruned_model, stats = prune_model(
+pruned_model, stats = opf.prune_model(
     model=model,
     neuron_selection_method="MAW",
     pruning_percentage=20,
@@ -317,13 +319,13 @@ Prune neurons only in specific layers while leaving others unchanged. Perfect fo
 
 ```python
 from transformers import AutoModelForCausalLM
-from optipfair import prune_model
+import optipfair as opf
 
 # Load a pre-trained model
 model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B")
 
 # Prune neurons only in specific layers (e.g., middle layers)
-pruned_model, stats = prune_model(
+pruned_model, stats = opf.prune_model(
     model=model,
     pruning_type="MLP_GLU",
     neuron_selection_method="MAW",
@@ -335,7 +337,7 @@ pruned_model, stats = prune_model(
 
 # Print pruning statistics
 print(f"Pruned {stats['pruned_layers']} of {stats['total_layers']} layers")
-print(f"Total reduction: {stats['reduction']:,} parameters ({stats['percentage_reduction']:.2f}%)
+print(f"Total reduction: {stats['reduction']:,} parameters ({stats['percentage_reduction']:.2f}%)")
 
 # Save the pruned model
 pruned_model.save_pretrained("./selective-pruned-llama")
@@ -359,13 +361,13 @@ Remove entire layers from a model for significant efficiency gains. Here, we rem
 
 ```python
 from transformers import AutoModelForCausalLM
-from optipfair import prune_model
+import optipfair as opf
 
 # Load a pre-trained model
 model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B")
 
 # Remove the last 4 transformer layers
-pruned_model, stats = prune_model(
+pruned_model, stats = opf.prune_model(
     model=model,
     pruning_type="DEPTH",
     num_layers_to_remove=4,
@@ -475,7 +477,7 @@ optipfair is designed to work with transformer-based language models that use GL
 optipfair offers two powerful structured pruning strategies:
 
 1.  **MLP Pruning (Width Pruning)**: Reduces the number of neurons within the MLP layers of GLU-based models. This is a fine-grained approach to improve efficiency. You can control it via `pruning_percentage` or a target `expansion_rate`. It uses several neuron importance metrics:
-    * **MAW (Maximum Absolute Weight)**: Default and most effective method.
+    * **PPM (Peak-to-Peak Magnitude)**: Default and most effective method (parameter `"MAW"` for backward compatibility).
     * **VOW (Variance of Weights)**
     * **PON (Product of Norms)**
 

{optipfair-0.2.2 → optipfair-0.2.4}/README.md

@@ -30,6 +30,8 @@
 
 ---
 > **New to optipfair?** Use our [LLM Reference Manual](optipfair_llm_reference_manual.txt) - paste it into ChatGPT, Claude or your Favourite LLM for guided assistance with any optipfair task.
+
+> **Note on Terminology:** The default neuron selection method is **PPM (Peak-to-Peak Magnitude)**, which calculates neuron importance based on the full dynamic range of weights (max + |min|). This method is formally described in: *Martra, P. (2025). Fragile Knowledge, Robust Instruction-Following: The Width Pruning Dichotomy in Llama-3.2. ArXiv. https://arxiv.org/abs/2512.22671*. For backward compatibility, the parameter value `"MAW"` is still accepted and maps to PPM.
 ### 🚀 Interactive Demos: Try optipfair NOW
 
 Experience optipfair's capabilities directly in your browser.
@@ -99,17 +101,17 @@ 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.
+Prune 20% of the MLP neurons from a model using the Peak-to-Peak Magnitude (PPM) method.
 
 ```python
 from transformers import AutoModelForCausalLM
-from optipfair import prune_model
+import optipfair as opf
 
 # 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(
+pruned_model, stats = opf.prune_model(
     model=model,
     pruning_type="MLP_GLU",
     neuron_selection_method="MAW",
@@ -135,7 +137,7 @@ The pruning process yields tangible results in model size and performance. Here'
 | **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.*
+*Results based on the [PPM pruning method](#neuron-selection-methods) (parameter `"MAW"`). Full benchmark results will be published shortly.*
 
 ### Data-Driven Width Pruning (NEW in v0.2.0)
 
@@ -144,7 +146,7 @@ Enhance pruning decisions with activation statistics from calibration data. This
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from torch.utils.data import DataLoader, TensorDataset
 import torch
-from optipfair import prune_model
+import optipfair as opf
 
 # Load model and tokenizer
 model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B")
@@ -163,9 +165,9 @@ dataset = TensorDataset(inputs['input_ids'], inputs['attention_mask'])
 dataloader = DataLoader(dataset, batch_size=8)
 
 # Prune with data-driven importance calculation
-pruned_model, stats = prune_model(
+pruned_model, stats = opf.prune_model(
     model=model,
-    neuron_selection_method="MAW",  # Only MAW supports data-driven pruning
+    neuron_selection_method="MAW",  # Only PPM (parameter "MAW") supports data-driven pruning
     pruning_percentage=20,
     dataloader=dataloader,  # ← Enables hybrid pruning
     show_progress=True,
@@ -182,7 +184,7 @@ pruned_model.save_pretrained("./pruned-datadriven-model")
 - 🔬 **Research-Backed**: Based on CFSP methodology (arXiv:2409.13199v2)
 - ⚡ **Easy Integration**: Just add a dataloader - no other changes needed
 
-**Note:** Data-driven pruning is currently only available with `neuron_selection_method="MAW"`. Using a dataloader with "VOW" or "PON" will raise a `ValueError`.
+**Note:** Data-driven pruning is currently only available with `neuron_selection_method="MAW"` (PPM method). Using a dataloader with "VOW" or "PON" will raise a `ValueError`.
 
 ### Selective Layer Width Pruning (NEW in v0.2.0)
 
@@ -190,13 +192,13 @@ Prune neurons only in specific layers while leaving others unchanged. Perfect fo
 
 ```python
 from transformers import AutoModelForCausalLM
-from optipfair import prune_model
+import optipfair as opf
 
 # Load a pre-trained model
 model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B")
 
 # Prune neurons only in specific layers (e.g., middle layers)
-pruned_model, stats = prune_model(
+pruned_model, stats = opf.prune_model(
     model=model,
     pruning_type="MLP_GLU",
     neuron_selection_method="MAW",
@@ -232,13 +234,13 @@ The `expansion_divisor` parameter ensures that intermediate layer sizes are divi
 
 ```python
 from transformers import AutoModelForCausalLM
-from optipfair import prune_model
+import optipfair as opf
 
 # Load model
 model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B")
 
 # Prune with hardware optimization
-pruned_model, stats = prune_model(
+pruned_model, stats = opf.prune_model(
     model=model,
     neuron_selection_method="MAW",
     pruning_percentage=20,
@@ -267,13 +269,13 @@ Prune neurons only in specific layers while leaving others unchanged. Perfect fo
 
 ```python
 from transformers import AutoModelForCausalLM
-from optipfair import prune_model
+import optipfair as opf
 
 # Load a pre-trained model
 model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B")
 
 # Prune neurons only in specific layers (e.g., middle layers)
-pruned_model, stats = prune_model(
+pruned_model, stats = opf.prune_model(
     model=model,
     pruning_type="MLP_GLU",
     neuron_selection_method="MAW",
@@ -285,7 +287,7 @@ pruned_model, stats = prune_model(
 
 # Print pruning statistics
 print(f"Pruned {stats['pruned_layers']} of {stats['total_layers']} layers")
-print(f"Total reduction: {stats['reduction']:,} parameters ({stats['percentage_reduction']:.2f}%)
+print(f"Total reduction: {stats['reduction']:,} parameters ({stats['percentage_reduction']:.2f}%)")
 
 # Save the pruned model
 pruned_model.save_pretrained("./selective-pruned-llama")
@@ -309,13 +311,13 @@ Remove entire layers from a model for significant efficiency gains. Here, we rem
 
 ```python
 from transformers import AutoModelForCausalLM
-from optipfair import prune_model
+import optipfair as opf
 
 # Load a pre-trained model
 model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B")
 
 # Remove the last 4 transformer layers
-pruned_model, stats = prune_model(
+pruned_model, stats = opf.prune_model(
     model=model,
     pruning_type="DEPTH",
     num_layers_to_remove=4,
@@ -425,7 +427,7 @@ optipfair is designed to work with transformer-based language models that use GL
 optipfair offers two powerful structured pruning strategies:
 
 1.  **MLP Pruning (Width Pruning)**: Reduces the number of neurons within the MLP layers of GLU-based models. This is a fine-grained approach to improve efficiency. You can control it via `pruning_percentage` or a target `expansion_rate`. It uses several neuron importance metrics:
-    * **MAW (Maximum Absolute Weight)**: Default and most effective method.
+    * **PPM (Peak-to-Peak Magnitude)**: Default and most effective method (parameter `"MAW"` for backward compatibility).
     * **VOW (Variance of Weights)**
     * **PON (Product of Norms)**
 

{optipfair-0.2.2 → optipfair-0.2.4}/docs/api.md

@@ -1,5 +1,13 @@
 # API Reference
 
+---
+
+**Note on Terminology:** The default neuron selection method is **PPM (Peak-to-Peak Magnitude)**, which calculates neuron importance based on the full dynamic range of weights (max + |min|). This method is formally described in: *Martra, P. (2025). Fragile Knowledge, Robust Instruction-Following: The Width Pruning Dichotomy in Llama-3.2. ArXiv. https://arxiv.org/abs/2512.22671*
+
+For backward compatibility, the parameter value `"MAW"` is still accepted and maps to PPM.
+
+---
+
 ## Core Functions
 
 ### `prune_model`
@@ -24,13 +32,13 @@ def prune_model(
     Prune a pre-trained language model using the specified pruning method.
     
     Supports both width pruning (neuron-level) and depth pruning (layer-level).
-    For width pruning with MAW method, can use static (weight-only) or hybrid
+    For width pruning with PPM method (parameter "MAW"), can use static (weight-only) or hybrid
     (weight + activation) importance calculation.
     
     Args:
         model: Pre-trained model to prune
         pruning_type: Type of pruning to apply ("MLP_GLU" or "DEPTH")
-        neuron_selection_method: Method to calculate neuron importance ("MAW", "VOW", or "PON") - for MLP_GLU only
+        neuron_selection_method: Method to calculate neuron importance ("MAW"/PPM, "VOW", "PON", or "L2") - for MLP_GLU only
         pruning_percentage: Percentage of neurons to prune (0-100) - for MLP_GLU only
         expansion_rate: Target expansion rate in percentage (mutually exclusive with pruning_percentage) - for MLP_GLU only
         dataloader: Optional PyTorch DataLoader for data-driven pruning (MLP_GLU with MAW only).
@@ -40,8 +48,8 @@ def prune_model(
             or as tuples of (input_ids, attention_mask). Typically 100-1000 samples
             from your target domain yield best results.
             
-            **Compatibility:** Only works with neuron_selection_method='MAW'.
-            Will raise ValueError if used with 'VOW' or 'PON'.
+            **Compatibility:** Only works with neuron_selection_method='MAW' (PPM method).
+            Will raise ValueError if used with 'VOW', 'PON', or 'L2'.
             
             **Example:**
                 >>> from torch.utils.data import DataLoader, TensorDataset
@@ -62,7 +70,7 @@ def prune_model(
         
     Raises:
         ValueError: If parameters are invalid or incompatible
-        ValueError: If dataloader is provided with non-MAW method
+        ValueError: If dataloader is provided with non-PPM method (use "MAW" parameter)
         
     Examples:
         >>> # Static width pruning (traditional)
@@ -328,6 +336,18 @@ def compute_neuron_pair_importance_pon(gate_weight: torch.Tensor, up_weight: tor
     """
     Compute neuron pair importance scores using Product of Norms method.
     
+    Args:
+        gate_weight: Weight matrix from the gate_proj layer
+        up_weight: Weight matrix from the up_proj layer
+        
+    Returns:
+        importance_scores: Importance scores for each neuron pair
+    """
+
+def compute_neuron_pair_importance_l2(gate_weight: torch.Tensor, up_weight: torch.Tensor) -> torch.Tensor:
+    """
+    Compute neuron pair importance scores using L2 norm method.
+    
     Args:
         gate_weight: Weight matrix from the gate_proj layer
         up_weight: Weight matrix from the up_proj layer

{optipfair-0.2.2 → optipfair-0.2.4}/docs/examples.md

@@ -2,6 +2,14 @@
 
 This page provides detailed examples of using OptiPFair for pruning language models and visualizing bias.
 
+---
+
+**Note on Terminology:** The default neuron selection method is **PPM (Peak-to-Peak Magnitude)**, which calculates neuron importance based on the full dynamic range of weights (max + |min|). This method is formally described in: *Martra, P. (2025). Fragile Knowledge, Robust Instruction-Following: The Width Pruning Dichotomy in Llama-3.2. ArXiv. https://arxiv.org/abs/2512.22671*
+
+For backward compatibility, the parameter value `"MAW"` is still accepted and maps to PPM.
+
+---
+
 ## Basic Pruning Example
 
 This example demonstrates how to prune a LLaMA model using the default settings:
@@ -29,7 +37,7 @@ outputs = model.generate(inputs.input_ids, max_length=50)
 original_output = tokenizer.decode(outputs[0], skip_special_tokens=True)
 print(f"Original model output: {original_output}")
 
-# Apply pruning with default settings (10% pruning, MAW method)
+# Apply pruning with default settings (10% pruning, PPM method)
 pruned_model = prune_model(model)
 
 # Get pruned parameter count

{optipfair-0.2.2 → optipfair-0.2.4}/docs/index.md

@@ -26,12 +26,14 @@ Understanding bias in language models is crucial for:
 - **Layer Importance Analysis**: Analyze transformer layers using cosine similarity to inform pruning decisions
 - **Bias Visualization Tools**: Comprehensive analysis of how models process demographic attributes
 - **Quantitative Bias Metrics**: Numeric measurements for consistent evaluation
-- **Multiple Neuron Selection Methods**: MAW, VOW, and PON for different pruning strategies
+- **Multiple Neuron Selection Methods**: PPM (Peak-to-Peak Magnitude), VOW, and PON for different pruning strategies
 - **Flexible Pruning Targets**: Support for both pruning percentage and target expansion rate
 - **Layer Selection Methods**: Choose specific layers or use automatic selection strategies
 - **Simple API and CLI**: Easy to use interfaces for Python and command line
 - **Progress Tracking**: Visual progress bars and detailed statistics
 
+**Note:** The default neuron selection method is PPM (Peak-to-Peak Magnitude), formally described in: *Martra, P. (2025). Fragile Knowledge, Robust Instruction-Following: The Width Pruning Dichotomy in Llama-3.2. ArXiv. https://arxiv.org/abs/2512.22671*. For backward compatibility, the parameter value `"MAW"` is still accepted.
+
 ## Getting Started
 
 - [Installation](installation.md): How to install OptiPFair

{optipfair-0.2.2 → optipfair-0.2.4}/docs/usage.md

@@ -1,6 +1,14 @@
 # Usage Guide
 This guide provides detailed instructions on how to use the core functionalities of OptiPFair, from pruning models to analyzing bias.
 
+---
+
+**Note on Terminology:** The default neuron selection method is **PPM (Peak-to-Peak Magnitude)**, which calculates neuron importance based on the full dynamic range of weights (max + |min|). This method is formally described in: *Martra, P. (2025). Fragile Knowledge, Robust Instruction-Following: The Width Pruning Dichotomy in Llama-3.2. ArXiv. https://arxiv.org/abs/2512.22671*
+
+For backward compatibility, the parameter value `"MAW"` is still accepted and maps to PPM.
+
+---
+
 ## Python API
 
 OptiPFair provides a simple Python API for pruning models.
@@ -14,7 +22,7 @@ from optipfair import prune_model
 # Load a pre-trained model
 model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B")
 
-# Prune the model with default settings (10% pruning, MAW method)
+# Prune the model with default settings (10% pruning, PPM method)
 pruned_model = prune_model(model=model)
 
 # Save the pruned model
@@ -48,7 +56,7 @@ OptiPFair provides a command-line interface for pruning models:
 ### Basic Usage
 
 ```bash
-# Prune a model with default settings (10% pruning, MAW method)
+# Prune a model with default settings (10% pruning, PPM method)
 optipfair prune --model-path meta-llama/Llama-3.2-1B --output-path ./pruned-model
 ```
 
@@ -75,19 +83,58 @@ optipfair analyze --model-path meta-llama/Llama-3.2-1B
 
 ## Neuron Selection Methods
 
-OptiPFair supports three methods for calculating neuron importance:
+OptiPFair supports four methods for calculating neuron importance:
 
-### MAW (Maximum Absolute Weight)
+### PPM (Peak-to-Peak Magnitude)
 
-The MAW method identifies neurons based on the maximum absolute weight values in their connections. This is typically the most effective method for GLU architectures.
+The PPM method identifies neurons based on the peak-to-peak magnitude of weights (max + |min|), capturing the full dynamic range of each neuron's weight values. This is typically the most effective method for GLU architectures. Use parameter value `"MAW"` for backward compatibility.
 
 ```python
 pruned_model = prune_model(
     model=model,
-    neuron_selection_method="MAW",
+    neuron_selection_method="MAW",  # PPM method ("MAW" for compatibility)
+    pruning_percentage=20
+)
+```
+
+### VOW (Variance of Weights)
+
+The VOW method identifies neurons based on the variance of their weight values.
+
+```python
+pruned_model = prune_model(
+    model=model,
+    neuron_selection_method="VOW",
+    pruning_percentage=20
+)
+```
+
+### PON (Product of Norms)
+
+The PON method uses the product of L1 norms to identify important neurons.
+
+```python
+pruned_model = prune_model(
+    model=model,
+    neuron_selection_method="PON",
+    pruning_percentage=20
+)
+```
+
+### L2 (L2 Norm)
+
+The L2 method calculates neuron importance using L2 norms of weight values.
+
+```python
+pruned_model = prune_model(
+    model=model,
+    neuron_selection_method="L2",
     pruning_percentage=20
 )
 ```
+
+**Note:** Data-driven pruning (hybrid mode) is only available with the PPM method (`"MAW"`). VOW, PON, and L2 support static (weight-only) pruning only.
+
 ## Data-Driven Pruning (v0.2.0+)
 
 ### Overview
@@ -260,9 +307,9 @@ pruned_datadriven = prune_model(
 ### Troubleshooting
 
 #### Error: "Data-driven pruning with dataloader is only supported for 'MAW' method"
-**Solution:** Change `neuron_selection_method` to `"MAW"`:
+**Solution:** Change `neuron_selection_method` to `"MAW"` (PPM method):
 ```python
-pruned = prune_model(model, neuron_selection_method="MAW", dataloader=dl)
+pruned = prune_model(model, neuron_selection_method="MAW", dataloader=dl)  # PPM method
 ```
 
 #### Out of Memory during calibration
@@ -354,6 +401,35 @@ This approach is often more intuitive when comparing across different model scal
 
 ## Depth Pruning
 
+Depth pruning removes entire transformer layers. When calling `prune_model(..., pruning_type="DEPTH", return_stats=True)`, the returned `stats` dictionary includes depth-specific fields:
+
+```python
+pruned_model, stats = prune_model(
+    model=model,
+    pruning_type="DEPTH",
+    num_layers_to_remove=3,          # or depth_pruning_percentage / layer_indices
+    layer_selection_method="last",  # "last" (default), "first", or "custom"
+    return_stats=True,
+)
+
+print(stats)
+# {
+#   'original_parameters': int,           # Parameter count before pruning
+#   'pruned_parameters': int,             # Parameter count after pruning
+#   'reduction': int,                     # Absolute reduction in parameters
+#   'percentage_reduction': float,        # Percentage reduction of parameters
+#   'original_layer_count': int,          # Layers before pruning
+#   'final_layer_count': int,             # Layers after pruning
+#   'layers_removed': int,                # Number of removed layers
+#   'layer_reduction_percentage': float   # Percentage of layers removed
+# }
+```
+
+- Depth pruning stats do not include `expansion_rate` (only relevant for MLP/GLU width pruning).
+- Internally, stats are captured before modifying the model to ensure correctness and avoid deepcopy issues.
+
+## Depth Pruning
+
 OptiPFair also supports depth pruning, which removes entire transformer layers from models. This is more aggressive than neuron-level pruning but can lead to significant efficiency gains.
 
 ### Python API
@@ -577,4 +653,88 @@ pruned_model = prune_model(
     pruning_type="DEPTH",
     layer_indices=layers_to_remove
 )
-```
+```
+
+### DataLoader Format Support (v0.2.4+)
+
+Starting from OptiPFair v0.2.4, `analyze_layer_importance` automatically handles multiple DataLoader batch formats, making it compatible with both HuggingFace datasets and native PyTorch structures.
+
+#### Supported Batch Formats
+
+**1. Dictionary Format (HuggingFace)**
+
+```python
+from datasets import load_dataset
+from torch.utils.data import DataLoader
+
+# HuggingFace datasets return dict batches
+dataset = load_dataset('wikitext', 'wikitext-2-raw-v1', split='train[:100]')
+tokenized = dataset.map(tokenize_function, batched=True)
+tokenized.set_format(type='torch', columns=['input_ids', 'attention_mask'])
+dataloader = DataLoader(tokenized, batch_size=8)
+
+# Batch format: {'input_ids': tensor, 'attention_mask': tensor}
+importance_scores = analyze_layer_importance(model, dataloader)
+```
+
+**2. Tuple Format (TensorDataset)**
+
+```python
+from torch.utils.data import DataLoader, TensorDataset
+
+# Tokenize texts manually
+inputs = tokenizer(
+    texts,
+    truncation=True,
+    padding='max_length',
+    max_length=512,
+    return_tensors='pt'
+)
+
+# TensorDataset returns tuples
+dataset = TensorDataset(inputs['input_ids'], inputs['attention_mask'])
+dataloader = DataLoader(dataset, batch_size=8)
+
+# Batch format: (input_ids, attention_mask)
+# Automatically mapped: [0]=input_ids, [1]=attention_mask
+importance_scores = analyze_layer_importance(model, dataloader)
+```
+
+**3. List Format (Custom Datasets)**
+
+```python
+class CustomDataset(Dataset):
+    def __getitem__(self, idx):
+        return [self.input_ids[idx], self.attention_mask[idx]]
+
+# Batch format: [input_ids, attention_mask]
+# Same positional mapping as tuples
+importance_scores = analyze_layer_importance(model, dataloader)
+```
+
+**4. Single Tensor Format**
+
+```python
+# Dataset with only input_ids (no attention_mask)
+dataset = TensorDataset(input_ids_tensor)
+dataloader = DataLoader(dataset, batch_size=8)
+
+# Batch format: single tensor
+# Automatically treated as input_ids
+importance_scores = analyze_layer_importance(model, dataloader)
+```
+
+#### Positional Mapping for Tuple/List Formats
+
+When using tuple or list batches, elements are automatically mapped to standard transformer arguments:
+
+- `[0]` → `input_ids` (required)
+- `[1]` → `attention_mask` (optional)
+- `[2]` → `token_type_ids` (optional)
+- `[3]` → `position_ids` (optional)
+- `[4]` → `head_mask` (optional)
+- `[5]` → `inputs_embeds` (optional)
+
+**Note**: All formats are fully backward compatible. Existing code continues to work without modifications.
+
+---