This commit is contained in:
hiyouga 2023-12-11 18:09:40 +08:00
parent 64744dde89
commit 0239d29fa0
1 changed files with 4 additions and 4 deletions

View File

@ -189,7 +189,7 @@ def resize_embedding_layer(model: "PreTrainedModel", tokenizer: "PreTrainedToken
Resize token embeddings.
"""
old_vocab_size = model.get_input_embeddings().weight.size(0)
new_vocab_size = len(tokenizer)
if new_vocab_size != old_vocab_size:
model.resize_token_embeddings(new_vocab_size, pad_to_multiple_of=64)
logger.info("Resized embedding tokens from {} to {}.".format(old_vocab_size, new_vocab_size))
if len(tokenizer) != old_vocab_size:
model.resize_token_embeddings(len(tokenizer), pad_to_multiple_of=64)
new_vocab_size = model.get_input_embeddings().weight.size(0)
logger.info("Resized token embeddings from {} to {}.".format(old_vocab_size, new_vocab_size))