### Install Dependencies

Source: https://github.com/princeton-nlp/cofipruning/blob/main/README.md

Run this script to install all necessary dependencies for CoFiPruning. Ensure you have a compatible version of transformers installed.

```bash
pip install -r requirements.txt
```

--------------------------------

### Initialize CoFiTrainer for Pruning

Source: https://context7.com/princeton-nlp/cofipruning/llms.txt

Initializes the CoFiTrainer for a minimal pruning run on the SST-2 dataset. This setup includes model configuration, tokenizer, L0Module, training arguments, and additional pruning-specific arguments.

```python
from transformers import (
    AutoConfig, AutoTokenizer, TrainingArguments,
    HfArgumentParser, default_data_collator
)
from datasets import load_dataset
from models.l0_module import L0Module
from models.modeling_bert import CoFiBertForSequenceClassification
from trainer.trainer import CoFiTrainer
from args import AdditionalArguments, DataTrainingArguments

# --- Minimal pruning run on SST-2 ---

model_name = "bert-base-uncased"
config = AutoConfig.from_pretrained(model_name, num_labels=2,
                                     finetuning_task="sst2")
config.do_distill = True
config.do_layer_distill = True

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = CoFiBertForSequenceClassification.from_pretrained(model_name, config=config)
teacher_model = CoFiBertForSequenceClassification.from_pretrained(model_name, config=config)

l0_module = L0Module(
    config,
    lagrangian_warmup=500,
    target_sparsity=0.90,
    pruning_type="structured_heads+structured_mlp+hidden+layer",
)

training_args = TrainingArguments(
    output_dir="./out/sst2/CoFi/sparsity0.9",
    num_train_epochs=20,
    per_device_train_batch_size=32,
    per_device_eval_batch_size=32,
    learning_rate=2e-5,
    evaluation_strategy="steps",
    eval_steps=500,
    logging_steps=100,
    save_steps=0,
    seed=57,
)

additional_args = AdditionalArguments(
    pruning_type="structured_heads+structured_mlp+hidden+layer",
    reg_learning_rate=0.01,
    target_sparsity=0.90,
    lagrangian_warmup_epochs=2,
    prepruning_finetune_epochs=1,
    distillation_path=model_name,
    do_distill=True,
    do_layer_distill=True,
    layer_distill_version=4,
    distill_loss_alpha=0.9,
    distill_ce_loss_alpha=0.1,
)

raw_dataset = load_dataset("glue", "sst2")
def preprocess(examples):
    return tokenizer(examples["sentence"], padding="max_length",
                     max_length=128, truncation=True)
tokenized = raw_dataset.map(preprocess, batched=True)

trainer = CoFiTrainer(
    model=model,
    args=training_args,
    additional_args=additional_args,
    train_dataset=tokenized["train"],
    eval_dataset=tokenized["validation"],
    tokenizer=tokenizer,
    data_collator=default_data_collator,
    l0_module=l0_module,
    teacher_model=teacher_model,
)

trainer.train()
# Best model saved to: ./out/sst2/CoFi/sparsity0.9/best/
#   - pytorch_model.bin  (full weights with learned masks)
#   - zs.pt              (binary mask tensors)
#   - l0_module.pt       (full L0Module state)
```

--------------------------------

### Train with CoFiPruning

Source: https://github.com/princeton-nlp/cofipruning/blob/main/README.md

Example script for training a model with CoFiPruning. This script supports single-GPU training and various pruning configurations. Arguments specify the task, experiment details, pruning types, sparsity targets, and distillation parameters.

```bash
TASK=MNLI
SUFFIX=sparsity0.95
EX_CATE=CoFi
PRUNING_TYPE=structured_heads+structured_mlp+hidden+layer
SPARSITY=0.95
DISTILL_LAYER_LOSS_ALPHA=0.9
DISTILL_CE_LOSS_ALPHA=0.1
LAYER_DISTILL_VERSION=4
SPARSITY_EPSILON=0.01

bash scripts/run_CoFi.sh $TASK $SUFFIX $EX_CATE $PRUNING_TYPE $SPARSITY [DISTILLATION_PATH] $DISTILL_LAYER_LOSS_ALPHA $DISTILL_CE_LOSS_ALPHA $LAYER_DISTILL_VERSION $SPARSITY_EPSILON
```

--------------------------------

### Fine-tune Pruned Model

Source: https://github.com/princeton-nlp/cofipruning/blob/main/README.md

Example script for fine-tuning a model after it has been pruned. Set `pruning_type` to `None` for standard fine-tuning. This script requires the path to the pruned model and a learning rate.

```bash
PRUNED_MODEL_PATH=$proj_dir/$TASK/$EX_CATE/${TASK}_${SUFFIX}/best
PRUNING_TYPE=None # Setting the pruning type to be None for standard fine-tuning.
LEARNING_RATE=3e-5

bash scripts/run_CoFi.sh $TASK $SUFFIX $EX_CATE $PRUNING_TYPE $SPARSITY [DISTILLATION_PATH] $DISTILL_LAYER_LOSS_ALPHA $DISTILL_CE_LOSS_ALPHA $LAYER_DISTILL_VERSION $SPARSITY_EPSILON [PRUNED_MODEL_PATH] $LEARNING_RATE
```

--------------------------------

### Initialize and Use L0Module for Pruning

Source: https://context7.com/princeton-nlp/cofipruning/llms.txt

Initializes an `L0Module` for a specified model configuration, targeting a desired sparsity. Demonstrates how to obtain stochastic masks during training and deterministic masks during inference. Includes usage of Lagrangian regularization loss and model size calculation.

```python
from transformers import AutoConfig
from models.l0_module import L0Module

# Initialize an L0Module for bert-base-uncased targeting 95% sparsity
config = AutoConfig.from_pretrained("bert-base-uncased")

l0_module = L0Module(
    config,
    droprate_init=0.5,          # initial dropout rate for loga initialization
    temperature=2.0 / 3.0,     # temperature of Hard Concrete distribution
    lagrangian_warmup=2000,     # number of steps to linearly ramp target sparsity
    start_sparsity=0.0,
    target_sparsity=0.95,
    pruning_type="structured_heads+structured_mlp+hidden+layer",
)

# --- Training forward: stochastic masks ---
l0_module.train()
zs_train = l0_module(training=True)
# zs_train keys: 'head_z', 'head_layer_z', 'intermediate_z', 'mlp_z', 'hidden_z'
# Each value is a float tensor in [0, 1] used as a soft gate inside the model.

# --- Inference forward: deterministic masks ---
l0_module.eval()
zs_eval = l0_module(training=False)

# --- Lagrangian regularization loss (added to task loss during pruning) ---
pruned_steps = 500
lag_loss, expected_sparsity, target_sparsity = l0_module.lagrangian_regularization(pruned_steps)
# lag_loss: scalar tensor — add to the main training loss
# expected_sparsity: float ~0.70 (current expected sparsity)
# target_sparsity:  float ~0.95 (current ramped target)
print(f"Expected sparsity: {expected_sparsity:.4f}, Target: {target_sparsity:.4f}")

# --- Inspect how many parameters remain ---
size_info = l0_module.calculate_model_size(zs_eval)
# Returns dict with keys:
#   head_layers, mlp_layers (per-layer binary flags)
#   hidden_dims (int), intermediate_dims (list), head_nums (list)
#   pruned_params, remaining_params, pruned_model_sparsity
print(f"Remaining params: {size_info['remaining_params']}")
print(f"Model sparsity:   {size_info['pruned_model_sparsity']:.4f}")

# Constrain loga parameters to prevent numerical overflow (call every step)
l0_module.constrain_parameters()
```

--------------------------------

### Benchmark Sparsity, Latency, and Accuracy with `evaluate.py`

Source: https://context7.com/princeton-nlp/cofipruning/llms.txt

Command-line script to evaluate CoFi pruned models. It loads a model, runs validation multiple times for stable latency, and reports size, sparsity, task score, and inference speed.

```bash
# Evaluate a released model on MNLI
python evaluation.py MNLI princeton-nlp/CoFi-MNLI-s95
# Expected output:
#   Task: mnli
#   Model path: princeton-nlp/CoFi-MNLI-s95
#   Model size: 4920106
#   Sparsity: 0.943
#   mnli/acc: 0.8055
#   seconds/example: 0.010151

# Evaluate a locally saved model on SST-2
python evaluation.py SST2 ./out/sst2/CoFi/sparsity0.9/best
```

--------------------------------

### Load and Apply Pruning Masks with `load_model_with_zs`

Source: https://context7.com/princeton-nlp/cofipruning/llms.txt

Loads a full-size checkpoint, applies soft mask weights, and permanently removes pruned units to return a structurally smaller model. Requires the model path and the classification model class.

```python
from models.modeling_bert import CoFiBertForSequenceClassification
from utils.cofi_utils import load_model, load_zs

model_path = "./out/sst2/CoFi/sparsity0.9/best"

# Load masks
zs = load_zs(model_path)

# Load the full checkpoint and produce a structurally pruned model
model = load_model(model_path, CoFiBertForSequenceClassification, zs)
# Output:
#   Load weights from ./out/sst2/CoFi/sparsity0.9/best
#   Model Size before pruning: 85054464
#   Model Size after pruning:  4920106
#   Model Size: 4920106

model.eval()
import torch
inputs = {
    "input_ids":      torch.randint(0, 30522, (1, 64)),
    "attention_mask": torch.ones(1, 64, dtype=torch.long),
}
with torch.no_grad():
    out = model(**inputs)
print(out.logits)  # tensor([[0.12, -0.34]])
```

--------------------------------

### Run CoFi Pruning Script (Bash)

Source: https://context7.com/princeton-nlp/cofipruning/llms.txt

Shell script to configure and launch single-GPU GLUE or SQuAD pruning jobs. Automatically selects epoch counts and evaluation intervals based on task size. Used for end-to-end pruning training.

```bash
# Prune BERT-base on MNLI to 95% sparsity with layer-wise distillation (version 4)
TASK=MNLI
SUFFIX=sparsity0.95
EX_CATE=CoFi
PRUNING_TYPE=structured_heads+structured_mlp+hidden+layer
SPARSITY=0.95
DISTILL_PATH=bert-base-uncased   # teacher = full fine-tuned model directory
DISTILL_LAYER_ALPHA=0.9
DISTILL_CE_ALPHA=0.1
LAYER_DISTILL_VERSION=4
SPARSITY_EPSILON=0.01            # start saving when sparsity > 0.94

bash scripts/run_CoFi.sh \
  $TASK $SUFFIX $EX_CATE $PRUNING_TYPE $SPARSITY \
  $DISTILL_PATH $DISTILL_LAYER_ALPHA $DISTILL_CE_ALPHA \
  $LAYER_DISTILL_VERSION $SPARSITY_EPSILON

# Fine-tune the best pruned model (no pruning, just standard FT)
PRUNED_MODEL=./out/MNLI/CoFi/MNLI_sparsity0.95/best
bash scripts/run_CoFi.sh \
  $TASK $SUFFIX $EX_CATE None $SPARSITY \
  $DISTILL_PATH $DISTILL_LAYER_ALPHA $DISTILL_CE_ALPHA \
  $LAYER_DISTILL_VERSION $SPARSITY_EPSILON \
  $PRUNED_MODEL 3e-5
```

--------------------------------

### Load and Evaluate a CoFi Model

Source: https://context7.com/princeton-nlp/cofipruning/llms.txt

Loads a CoFi model from a specified path, handling zero-shot (zs) loading if available. Evaluates model size and calculates sparsity against a full model configuration.

```python
import sys, torch
from utils.cofi_utils import load_zs, load_model
from utils.utils import calculate_parameters
from models.modeling_bert import CoFiBertForSequenceClassification

model_path = "princeton-nlp/CoFi-SST2-s95"
zs = load_zs(model_path)
if zs is None:
    model = CoFiBertForSequenceClassification.from_pretrained(model_path)
else:
    model = load_model(model_path, CoFiBertForSequenceClassification, zs)

model = model.cuda().eval()
model_size = calculate_parameters(model)
full_size  = calculate_parameters(CoFiBertForSequenceClassification(model.config))
sparsity   = 1 - round(model_size / full_size, 3)
print(f"Parameters: {model_size}  |  Sparsity: {sparsity:.3f}")
```

--------------------------------

### Evaluate Pruned Model

Source: https://github.com/princeton-nlp/cofipruning/blob/main/README.md

Use this script to obtain sparsity, inference time, and development set results for a pruned model. Specify the task and the model name or directory.

```bash
python evaluation.py [TASK] [MODEL_NAME_OR_DIR]
```

```bash
python evaluation.py MNLI princeton-nlp/CoFi-MNLI-s95 
```

--------------------------------

### Print Model Parameter Names

Source: https://github.com/princeton-nlp/cofipruning/blob/main/notebook/function_test.ipynb

Iterates through and prints the names of all parameters within the loaded model. This is useful for understanding the model's architecture and identifying specific layers.

```python
for n, p in model.named_parameters():
    print(n)
```

--------------------------------

### Import Transformers Library

Source: https://github.com/princeton-nlp/cofipruning/blob/main/notebook/function_test.ipynb

Imports the necessary components from the transformers library for model loading.

```python
import transformers
from transformers import AutoModelForSequenceClassification
```

--------------------------------

### Prepare Tokenized Inputs for Model

Source: https://context7.com/princeton-nlp/cofipruning/llms.txt

Prepares tokenized inputs for a model, including input IDs, attention masks, and token type IDs. This is typically done before passing inputs to the model for inference or training.

```python
inputs = {
    "input_ids":      torch.randint(0, 30522, (2, 32)),
    "attention_mask": torch.ones(2, 32, dtype=torch.long),
    "token_type_ids": torch.zeros(2, 32, dtype=torch.long),
}

with torch.no_grad():
    output = model(**inputs)
print(output.logits.shape)   # torch.Size([2, 3])  — 3 classes for MNLI
```

--------------------------------

### Load Pre-trained Sequence Classification Model

Source: https://github.com/princeton-nlp/cofipruning/blob/main/notebook/function_test.ipynb

Loads a pre-trained model for sequence classification. This model is initialized from the 'roberta-base' checkpoint. Note that some weights may not be used and new weights may be initialized, indicating the model is ready for downstream task training.

```python
model = AutoModelForSequenceClassification.from_pretrained("roberta-base")
```

--------------------------------

### BibTeX Citation for CoFiPruning

Source: https://github.com/princeton-nlp/cofipruning/blob/main/README.md

Cite this paper when using CoFiPruning in your work. This provides the necessary details for academic referencing.

```bibtex
@inproceedings{xia2022structured,
   title={Structured Pruning Learns Compact and Accurate Models},
   author={Xia, Mengzhou and Zhong, Zexuan and Chen, Danqi},
   booktitle={Association for Computational Linguistics (ACL)},
   year={2022}
}

```

--------------------------------

### Load Saved Sparsity Masks

Source: https://context7.com/princeton-nlp/cofipruning/llms.txt

Loads the `zs.pt` mask tensor dictionary from a model directory. If `zs.pt` is not found, it falls back to loading `l0_module.pt` and regenerating the masks.

```python
from utils.cofi_utils import load_zs

# Load from a pruned checkpoint directory
zs = load_zs("./out/sst2/CoFi/sparsity0.9/best")
# Returns dict with keys like: 'head_z', 'head_layer_z', 'intermediate_z', 'mlp_z', 'hidden_z'
# Each value is a CPU tensor.

if zs is None:
    print("No zs found — model is already structurally pruned.")
else:
    print("head_layer_z:", zs["head_layer_z"])   # shape: [12] — which attention layers survive
    print("mlp_z:       ", zs["mlp_z"])           # shape: [12] — which MLP layers survive
```

--------------------------------

### Load Pre-trained CoFiPruning Model

Source: https://github.com/princeton-nlp/cofipruning/blob/main/README.md

Load a pre-trained CoFiPruning model for sequence classification using the Hugging Face interface. Ensure the 'models.modeling_bert' module is available.

```python
from models.modeling_bert import CoFiBertForSequenceClassification
model = CoFiBertForSequenceClassification.from_pretrained("princeton-nlp/CoFi-MNLI-s95") 
output = model(**inputs)
```

--------------------------------

### Calculate Model Parameters (Python)

Source: https://context7.com/princeton-nlp/cofipruning/llms.txt

Counts non-embedding parameters in a model, excluding specific components to match CoFiPruning paper's sparsity calculation convention. Useful for verifying compression ratios.

```python
from utils.utils import calculate_parameters
from transformers import AutoConfig
from models.modeling_bert import CoFiBertForSequenceClassification

config = AutoConfig.from_pretrained("bert-base-uncased", num_labels=3)
full_model = CoFiBertForSequenceClassification(config)

full_size = calculate_parameters(full_model)
print(f"Full BERT-base (encoder only): {full_size:,}")
# Full BERT-base (encoder only): 85,054,464

# After structural pruning to 95% sparsity:
from utils.cofi_utils import load_model, load_zs
pruned = load_model("princeton-nlp/CoFi-MNLI-s95",
                    CoFiBertForSequenceClassification,
                    load_zs("princeton-nlp/CoFi-MNLI-s95"))
pruned_size = calculate_parameters(pruned)
sparsity    = 1 - pruned_size / full_size
print(f"Pruned model: {pruned_size:,}  |  Sparsity: {sparsity:.3f}")
# Pruned model: 4,920,106  |  Sparsity: 0.943
```

--------------------------------

### Load Released HuggingFace Compressed Model with `load_pruned_model`

Source: https://context7.com/princeton-nlp/cofipruning/llms.txt

Reconstructs sparse mask tensors from weight shapes of a structurally small HuggingFace checkpoint and applies structural pruning to a freshly initialized model. Use this when the checkpoint lacks a `zs.pt` file.

```python
import torch
from models.modeling_bert import CoFiBertForSequenceClassification
from utils.cofi_utils import load_pruned_model

# Load released model weights
weights = torch.load("pytorch_model.bin", map_location="cpu")

# Create a full-size shell
from transformers import AutoConfig
config = AutoConfig.from_pretrained("princeton-nlp/CoFi-MNLI-s95")
model = CoFiBertForSequenceClassification(config)

# Reconstruct zs from weight shapes and apply structural pruning
model = load_pruned_model(model, weights)
print(model.config.hidden_size)  # 768 (config unchanged)
# But internal weight shapes are smaller — fast at inference
```

--------------------------------

### Pass ZS Masks into Forward Pass

Source: https://context7.com/princeton-nlp/cofipruning/llms.txt

During pruning training, this demonstrates how to pass zero-sparsity (zs) masks into the forward pass of a model. The CoFiTrainer automatically handles this via `fill_inputs_with_zs`.

```python
from transformers import AutoConfig
from models.l0_module import L0Module

config = AutoConfig.from_pretrained("bert-base-uncased")
l0_module = L0Module(config, target_sparsity=0.9,
                     pruning_type="structured_heads+structured_mlp+hidden+layer")
l0_module.eval()
zs = l0_module(training=False)   # deterministic masks

inputs_with_zs = {**inputs, **zs}   # merge masks into inputs dict
with torch.no_grad():
    output = model(**inputs_with_zs)
```

--------------------------------

### Load a Released Compressed Model

Source: https://context7.com/princeton-nlp/cofipruning/llms.txt

Loads a pre-trained compressed BERT model directly from the HuggingFace Hub using `CoFiBertForSequenceClassification`. This is useful for directly using models that have already undergone the CoFiPruning process.

```python
from models.modeling_bert import (
    CoFiBertForSequenceClassification,
    CoFiBertForQuestionAnswering,
)
from models.modeling_roberta import CoFiRobertaForSequenceClassification
import torch

# --- Load a released compressed model directly from HuggingFace Hub ---
model = CoFiBertForSequenceClassification.from_pretrained("princeton-nlp/CoFi-MNLI-s95")
model.eval()
```

--------------------------------

### Permanently Prune Model Weights with `prune_model_with_z`

Source: https://context7.com/princeton-nlp/cofipruning/llms.txt

Applies binary masks to permanently remove pruned units from weight matrices, shrinking the model's weight tensors. This function requires the mask dictionary `zs` and the model itself. It's typically preceded by `update_params`.

```python
import torch
from models.modeling_bert import CoFiBertForSequenceClassification
from utils.cofi_utils import prune_model_with_z, update_params

model = CoFiBertForSequenceClassification.from_pretrained("bert-base-uncased")

# Construct example zs (normally loaded from zs.pt)
zs = {
    "head_z":         torch.ones(12, 12),   # keep all heads
    "head_layer_z":   torch.ones(12),
    "intermediate_z": torch.ones(12, 3072),
    "mlp_z":          torch.ones(12),
    "hidden_z":       torch.ones(768),
}
# Zero out some units to simulate pruning decisions
zs["head_z"][0, :6] = 0      # prune first 6 heads of layer 0
zs["mlp_z"][11] = 0          # prune entire MLP of last layer

# Step 1: weight masking (zero out outgoing weights of pruned units)
update_params(model, zs)
# Step 2: structural removal (shrinks linear layer weight shapes)
prune_model_with_z(zs, model)

# Verify layer 0 now has 6 attention heads instead of 12
print(model.bert.encoder.layer[0].attention.self.query.weight.shape)
# torch.Size([384, 768])  — 6 heads × 64 dims_per_head = 384
```

=== COMPLETE CONTENT === This response contains all available snippets from this library. No additional content exists. Do not make further requests.