fix bug for vocab size

Author: zhenyuhe00

datastore/get_datastore_code.py

@@ -41,7 +41,7 @@
 writer = draftretriever.Writer(
     index_file_path=datastore_path,
     max_chunk_len=512 * 1024 * 1024,
-    vocab_size=tokenizer.vocab_size,
+    vocab_size=tokenizer.vocab_size + len(tokenizer.get_added_vocab()),
 )
 
 total_length = len(dataset)
@@ -51,4 +51,4 @@
     token_list = tokenizer.encode(sample['content'])
     writer.add_entry(token_list)
 
-writer.finalize()
+writer.finalize()