NameError when trying to trim MBART

[tatiana-iazykova]

Hi!

I'm trying to use you library to trim MBART but it produces the following error:

# trim model
mt = MBartTrimmer(model, config, tt.trimmed_tokenizer)
mt.make_weights(tt.trimmed_vocab_ids)
mt.make_model()

---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
Input In [13], in <cell line: 4>()
      2 mt = MBartTrimmer(model, config, tt.trimmed_tokenizer)
      3 mt.make_weights(tt.trimmed_vocab_ids)
----> 4 mt.make_model()

File ~/translator/og/hf_trim/hftrim/model_trimmers/BaseTrimmer.py:21, in BaseTrimmer.make_model(self)
     19 def make_model(self):
     20     self.initialize_new_model()
---> 21     self.trim_model()
     22     return self.trimmed_model

File ~/translator/og/hf_trim/hftrim/model_trimmers/BartTrimmer.py:95, in BartTrimmer.trim_model(self)
     92 self.trimmed_model.set_input_embeddings(prunedEmbeddingMatrix)
     94 if 'lm_head' in self.trimmed_weights:
---> 95     prunedLMHeadMatrix = torch.Tensor(new_lmh)
     96     _ = self.trimmed_model.lm_head.weight.data.copy_(prunedLMHeadMatrix)
     98 # tie weights as described in model config

NameError: name 'new_lmh' is not defined

I've looked through the source code and indeed there's no such variable. Maybe you can help me with solving this issue?

[IamAdiSri]

Hi @tatiana-iazykova, this looks like a bug. Let me have a look and get back to you.

[IamAdiSri]

@tatiana-iazykova Hi, I've made a fix that works on my machine. I also noticed that the code in the readme file was wrong earlier and have fixed that too. I suspect that you couldn't import the modules from the pip package earlier but it should work now.

Can you clone the repository and reinstall it? The commands are given below;

$ pip uninstall hf-trim
$ git clone https://github.com/IamAdiSri/hf-trim
$ cd hf-trim
$ pip install -e .

I'll also paste the code I used to test it below, in case you want to try this;

from transformers import MBartForConditionalGeneration, MBartConfig, MBartTokenizer, MBart50Tokenizer
from hftrim.TokenizerTrimmer import TokenizerTrimmer
from hftrim.ModelTrimmers import MBartTrimmer

data = [
        " UN Chief Says There Is No Military Solution in Syria", 
        "Şeful ONU declară că nu există o soluţie militară în Siria"
]

# load pretrained config, tokenizer and model
config = MBartConfig.from_pretrained("facebook/mbart-large-50-many-to-one-mmt")
model = MBartForConditionalGeneration.from_pretrained("facebook/mbart-large-50-many-to-one-mmt")
tokenizer = MBart50Tokenizer.from_pretrained("facebook/mbart-large-50-many-to-one-mmt")

# trim tokenizer
tt = TokenizerTrimmer(tokenizer)
tt.make_vocab(data)
tt.make_tokenizer()

# trim model
mt = MBartTrimmer(model, config, tt.trimmed_tokenizer)
mt.make_weights(tt.trimmed_vocab_ids)
mt.make_model()

Let me know if this works now.

[tatiana-iazykova]

Yes! Thanks a lot!

[IamAdiSri]

@tatiana-iazykova Hi, the PyPI package is also functional now and I'll recommend uninstalling the repo and switching to the latest version with pip.

You'll have to change your import statement from from hftrim.model_trimmers.MBartTrimmer import MBartTrimmer to from hftrim.ModelTrimmers import MBartTrimmer

Everything else functions the same. I have updated the code in my previous comment to reflect the changes as well.

[tatiana-iazykova]

Hi again! Caught the same error while trimming MT5 :(

[IamAdiSri]

@tatiana-iazykova Should be fixed now. You'll have to upgrade your package to v3.0.1