"""How many numbers does LoRA train, for each size of add-on?

    python lora_numbers.py

It loads the model and prints the sizes of the two matrices that get an add-on
in train.py. Then it really builds the add-ons at seven sizes, counts their
numbers, and prints the count beside what the formula on the page gives.
"""
import torch
from peft import LoraConfig, get_peft_model
from transformers import AutoModelForCausalLM

MODEL = "HuggingFaceTB/SmolLM2-135M-Instruct"
RANKS = [1, 2, 4, 8, 16, 32, 64]
EVERY_KIND = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]


def count_numbers(tensors):
    total = 0
    for tensor in tensors:
        total = total + tensor.numel()
    return total


def trained_numbers(rank, layers):
    """Build the add-ons for real and count the numbers that training may change."""
    model = AutoModelForCausalLM.from_pretrained(MODEL, dtype=torch.float32)
    config = LoraConfig(r=rank, lora_alpha=2 * rank, target_modules=layers,
                        task_type="CAUSAL_LM")
    tuned = get_peft_model(model, config)
    return count_numbers(t for t in tuned.parameters() if t.requires_grad)


def main():
    model = AutoModelForCausalLM.from_pretrained(MODEL, dtype=torch.float32)
    whole = count_numbers(model.parameters())
    blocks = model.config.num_hidden_layers
    attention = model.model.layers[0].self_attn
    print(f"{MODEL}")
    print(f"{whole:,} numbers, in {blocks} blocks one after another")
    print()
    print("In each block, train.py puts an add-on on two matrices:")
    per_block = 0
    for name in ["q_proj", "v_proj"]:
        rows, columns = getattr(attention, name).weight.shape
        addon = 8 * columns + rows * 8
        per_block = per_block + rows + columns
        print(f"  {name}: {rows} x {columns} = {rows * columns:,} numbers."
              f" Its add-on at r=8: 8 x {columns} and {rows} x 8 = {addon:,} numbers.")
    print(f"So each 1 in r is {per_block:,} numbers in a block,"
          f" and {blocks} x {per_block:,} = {blocks * per_block:,} numbers in the model.")
    print()
    print("   r   numbers trained   by the formula   share of the model")
    for rank in RANKS:
        trained = trained_numbers(rank, ["q_proj", "v_proj"])
        formula = rank * blocks * per_block
        print(f"{rank:4}   {trained:15,}   {formula:14,}   {100 * trained / whole:16.2f} %")
    print()
    wide = trained_numbers(8, EVERY_KIND)
    print(f"With an add-on on all seven kinds of matrix in a block, at r=8:"
          f" {wide:,} numbers, {100 * wide / whole:.2f} % of the model.")


if __name__ == "__main__":
    main()