Merge pull request #11 from helpmefindaname/fix_output_vocab_size

set the vocab size correctly when recreating the full embedding

tests/test_contextual_reduce.py

@@ -69,6 +69,7 @@ def test_saving_while_reduction_can_be_loaded_afterwards():
         "Home sweet home",
         "ay ay ay",
     ]
+    initial_vocab_size = model.config.vocab_size
     with tempfile.TemporaryDirectory() as tdir:
         with reduce_train_vocab(model=model, tokenizer=tokenizer, texts=texts):
             model.save_pretrained(tdir)
@@ -77,3 +78,5 @@ def test_saving_while_reduction_can_be_loaded_afterwards():
         new_tokenizer = AutoTokenizer.from_pretrained(tdir)
         assert new_model.config.vocab_size == 13
         assert len(new_tokenizer) == 13
+    assert model.config.vocab_size == initial_vocab_size
+    assert len(tokenizer) == initial_vocab_size

transformer_smaller_training_vocab/modify_model.py

@@ -54,5 +54,6 @@ def recreate_embedding(
     for reduced_id, full_id in enumerate(keep_token_ids):
         saved_embeddings[full_id] = embedding_weights[reduced_id]
     new_input_embedding = nn.Embedding(saved_embeddings.size(0), saved_embeddings.size(1), _weight=saved_embeddings)
+    model.config.vocab_size = saved_embeddings.size(0)
     model.set_input_embeddings(new_input_embedding)
     model.get_input_embeddings().to(model_device)