Fix bug in setting output vocab size in tokensurgeon

arcee-ai · Dec 1, 2024 · 2ab77ab · 2ab77ab
1 parent 6f7a79e
commit 2ab77ab
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/mergekit/scripts/tokensurgeon.py b/mergekit/scripts/tokensurgeon.py
@@ -190,7 +190,7 @@ def main(
     tokenizer.save_pretrained(out_path)
     cfg_out = arch_info.config
     try:
-        cfg_out.vocab_size = tokenizer.vocab_size
+        cfg_out.vocab_size = new_embed.shape[0]
     except AttributeError:
         LOG.error(
             "Could not set vocab size in config.json - you may need to update it manually."