import sys

import torch


def add_class_labels_to_checkpoint(input_path, classes_file_path):
    with open(classes_file_path, "r") as f:
        class_labels = [line.strip() for line in f.readlines()]

    assert len(class_labels) > 0, "The classes file is empty."

    # 1. Load the checkpoint
    checkpoint = torch.load(
        input_path, map_location=torch.device("cpu"), weights_only=False
    )

    if "classification_labels" in checkpoint:
        print(
            "Warning: 'classification_labels' key already exists in the checkpoint and will be overwritten."
        )


    # 2. Add your custom key/value pair
    checkpoint["classification_labels"] = class_labels

    # 3. Save the modified checkpoint
    output_path = input_path.replace(".ckpt", "_modified.ckpt")
    torch.save(checkpoint, output_path)
    print(f"Successfully added classification_labels and saved to {output_path}")


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python modify_checkpoints.py <input_checkpoint> <classes_file>")
        sys.exit(1)

    input_ckpt = sys.argv[1]
    classes_file = sys.argv[2]

    add_class_labels_to_checkpoint(
        input_path=input_ckpt, classes_file_path=classes_file
    )

import os
import tempfile

import torch

from chebai.models.electra import Electra

# Smallest viable config
model = Electra(
    model_type="classification",
    config={
        "vocab_size": 10,
        "max_position_embeddings": 1,
        "num_attention_heads": 1,
        "num_hidden_layers": 1,
        "type_vocab_size": 1,
        "hidden_size": 1,
        "intermediate_size": 1,
    },
    out_dim=10,
    input_dim=10,
)

# Count parameters
num_params = sum(p.numel() for p in model.parameters())
print(f"Total parameters: {num_params}")

# Save checkpoint and measure size
with tempfile.TemporaryDirectory() as tmpdir:
    ckpt_path = os.path.join(tmpdir, "electra_small.ckpt")
    torch.save({"state_dict": model.state_dict()}, ckpt_path)
    size_bytes = os.path.getsize(ckpt_path)
    print(f"Checkpoint size: {size_bytes} bytes")