ppak10/AdditiveLLM-Notebooks

1from transformers import T5EncoderModel, T5Config2from huggingface_hub import hf_hub_download3import torch.nn as nn4import torch5 6NUM_LABELS = 47 8class T5ClassificationModel(nn.Module):9    def __init__(self, model_path="t5-small", freeze_weights=True):10        super(T5ClassificationModel, self).__init__()11        if model_path == "t5-small":12            self.base_model = T5EncoderModel.from_pretrained(model_path)13        else:14            pytorch_model_path = hf_hub_download(15                repo_id=model_path,16                repo_type="model",17                filename="pytorch_model.bin"18            )19            config = T5Config.from_pretrained(model_path)20            self.base_model = T5EncoderModel(config)21 22            # Load the state_dict and remove unwanted keys23            state_dict = torch.load(pytorch_model_path, map_location=torch.device("cpu"))24            filtered_state_dict = {25                k.replace("base_model.", ""): v26                for k, v in state_dict.items()27                if not k.startswith("classifier.")28            }29            self.base_model.load_state_dict(filtered_state_dict)30 31        # For push to hub.32        self.config = self.base_model.config33        34        # Freeze the base model's weights35        if freeze_weights:36            for param in self.base_model.parameters():37                param.requires_grad = False38        39        # Add a classification head40        self.classifier = nn.Linear(self.base_model.config.hidden_size, NUM_LABELS)41    42    def forward(self, input_ids, attention_mask, labels=None):43        with torch.no_grad():  # No gradients for the base model44            outputs = self.base_model(input_ids=input_ids, attention_mask=attention_mask)45        46        # Sum token representations47        summed_representation = outputs.last_hidden_state.sum(dim=1)  # Summing over the sequence length (dim=1)48        49        logits = self.classifier(summed_representation)  # Pass the summed representation to the classifier50        loss = None51        if labels is not None:52            loss_fn = nn.BCEWithLogitsLoss()53            loss = loss_fn(logits, labels.float())54        return {"loss": loss, "logits": logits}55