Upload a custom model on Hub

Models
1

Hi, I have to modify a HF model from my repository changing the forward function, and I do it defining a new class ModifiedModel which takes the original model and defines new forward, and it works well.

The problem is that I don’t understand how to upload completely the new model on a new repo, because if I inherit the push_to_hub method from the child model it pushes only the child, so the pushed model is the original one, how can I push the model in order to have online all the new behaviour?

i hope to be clear, thanks…


import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch.nn.functional as F


def reshape_logits(logits):
    sorted_indices = torch.argsort(logits, descending=True)  # Ordina i token per probabilità
    original_probs = torch.softmax(logits, dim=-1)  # Converti logits in probabilità

    reshaped_probs = torch.zeros_like(original_probs)
    for i, idx in enumerate(sorted_indices):
        reshaped_probs[idx] = 1 / (2 ** (i + 1))  # i+1 per partire da 1

    reshaped_probs = reshaped_probs / reshaped_probs.sum()  # Riscala a sommare 1
    reshaped_logits = torch.log(reshaped_probs + 1e-12)  # Converti di nuovo in logits

    return reshaped_logits

class ModifiedModel(torch.nn.Module):
    def __init__(self, base_model):
        super().__init__()
        self.model = base_model

    def forward(self, *args, **kwargs):

        result = self.model.forward(*args, **kwargs)
        last_token_logits = result.logits[0, -1, :]
        probs = F.softmax(last_token_logits, dim=-1)

        reshaped_probs = torch.zeros_like(probs)
        sorted_indices = torch.argsort(probs, descending=True)

        for i, idx in enumerate(sorted_indices):
            reshaped_probs[idx] = 1 / (2 ** (i + 1))  # i+1 per partire da 1
        
        reshaped_probs = reshaped_probs / reshaped_probs.sum()  # Riscala a sommare 1
        reshaped_logits = torch.log(reshaped_probs + 1e-12)
        
        result.logits = reshaped_logits.unsqueeze(0).unsqueeze(0)  
          

        return result


# Modello originale e nome per il modello modificato
original_model_name = "dgambettavuw/M_gen5_run0_llama2-7b_wiki_doc1000_real64_synt64_vuw"

# Carica il modello e il tokenizer
model = AutoModelForCausalLM.from_pretrained(original_model_name, torch_dtype=torch.float16, device_map="auto")
tokenizer = AutoTokenizer.from_pretrained(original_model_name)


modified_model = ModifiedModel(model)
1 Like
2

I’ve never done it, but I think it’s probably something like this…