Text Generation
Transformers
PyTorch
Safetensors
English
hf_olmo
custom_code

Can I use AutoModel instead of AutoModelForCausalLM ?

#9
by Mengyao00 - opened

Hello, I don't need the final emgedding layer:(ff_out): Embedding(50304, 2048)

But when I tried AutoModel I meet this error:
local variable 'AutoModel' referenced before assignment

Do you know how I can delete the final embedding layer?

Ai2 org

We don't have a simple way to remove the final embedding layer. The options I can think of:

  1. You might be able to call model.set_output_embeddings with a 50304-dimensional identity linear layer. Then the embedding layer would have essentially no effect as long as you are not running gradient descent to update model weights. You will need to set weight_tying to false to make this work.
  2. You can grab the source code directly from https://github.com/allenai/OLMo/tree/main/olmo and make whatever changes you would like to. A minimal change would be modifying https://github.com/allenai/OLMo/blob/d7d4de4cf881f1a4a4a988cce26e8be114c9ff52/olmo/model.py#L1303 to not use a linear layer.

Sign up or log in to comment