Deci/DeciLM-7B-instruct · how to save the model as a pytorch or tensorflow model

Dec 27, 2023

Is there a wat to save the model as a pytorch model instance instead of loading it every time with the transformers module? I have tried with wrapping the code inside a class that inherits torch.nn.Module but when I try to save the model (all the model not only the state dict) it throws an error.
Thanks

NajeebDeci

Deci AI org Dec 31, 2023

Hi Rick,

Please provide more details on what you need exactly and code snippets in order for us to help.
Are you having issues with .save_pretrained ?

Thanks

Rick-29

Dec 31, 2023

Hi Najeeb,

No, I ment saving completly the model architecture and weigths as a .pth file, I tried using a model wrapper using torch.nn.Module, something like this:

import torch
from torch import nn

class Wrapper(nn.Module):
    def __init__(self, model, tokenizer):
        # The model and tokenizer are loaded exaclty like in the `DeciLM-7B-Instruct.ipynb` colab notebook 
        super().__init__()
        self.model = model
        self.tokenizer = tokenizer
    
    def forward(self, x):
        inputs = self.tokenizer(SYSTEM_PROMPT_TEMPLATE.format(instruction=x), return_tensors="pt")
        if torch.cuda.is_available():  # Ensure input tensors are on the GPU if model is on GPU
            inputs = inputs.to('cuda')
        output = self.model.generate(**inputs,
                                max_new_tokens=3000,
                                num_beams=5,
                                no_repeat_ngram_size=4,
                                early_stopping=True
                                )
        return self.tokenizer.decode(output[0], skip_special_tokens=True)

wrapper = Wrapper(model, tokenizer)
model = torch.jit.script(wrapper)
torch.jit.save(model, "model.pth")

But it doesn't work.

What should I do?

Thanks