from transformers import RobertaConfig
from transformers import RobertaForMaskedLM
from pytorch_block_sparse import BlockSparseModelPatcher
import re

config = RobertaConfig(
    vocab_size=52_000,
    max_position_embeddings=514,
    num_attention_heads=12,
    num_hidden_layers=6,
    type_vocab_size=1,
)

model = RobertaForMaskedLM(config=config).cuda()

# =>84 million parameters
print(f"Initial model parameters count={model.num_parameters()}")
 # Create a model patcher
mp = BlockSparseModelPatcher()

# Show names that can be used: this returns a list of all names in the network that are patchable.
# These names are escaped to be used as regexps in mp.add_pattern()
patchables = mp.get_patchable_layers(model)

dedup_layers = []

# Pretty print the regexps: replace layer number with regexp matching numbers, and dedup them
# This is a bit specific to Roberta, but should work for most transformers, it's just for ease of reading.
for patchable in patchables:
    r = patchable["regexp"]
    r = re.sub(r'[0-9]+', '[0-9]+', r)
    if r not in dedup_layers:
        dedup_layers.append(r)
        layer = patchable['layer']
        print(f"{r}\n   => {layer.in_features}x{layer.out_features}, bias={layer.bias is not None}")

# Selecting some layers to sparsify.
# This is the "artful" part, as some parts are more prone to be sparsified, other may impact model precision too much.

# Match layers using regexp (we escape the ., just because, it's more correct, but it does not change anything here)
# the [0-9]+ match any layer number.
# We setup a density of 0.5 on these layers, you can test other layers / densities .
mp.add_pattern("roberta\.encoder\.layer\.[0-9]+\.intermediate\.dense", {"density":0.5})
mp.add_pattern("roberta\.encoder\.layer\.[0-9]+\.output\.dense", {"density":0.5})
mp.add_pattern("roberta\.encoder\.layer\.[0-9]+\.attention\.output\.dense", {"density":0.5})
mp.patch_model(model)

print(f"Final model parameters count={model.num_parameters()}")

# => 68 million parameters instead of 84 million parameters (embeddings are taking a lof space in Roberta)