@@ -29,7 +29,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()),
 )
 if args.large_datastore:
     dataset = load_dataset('stingning/ultrachat', split='train')
Original file line number	Diff line number	Diff line change
`@@ -29,7 +29,7 @@`
`29`	`29`	`writer = draftretriever.Writer(`
`30`	`30`	`index_file_path=datastore_path,`
`31`	`31`	`max_chunk_len=51210241024,`
`32`		`- vocab_size=tokenizer.vocab_size,`
	`32`	`+ vocab_size=tokenizer.vocab_size + len(tokenizer.get_added_vocab()),`
`33`	`33`	`)`
`34`	`34`	`if args.large_datastore:`
`35`	`35`	`dataset = load_dataset('stingning/ultrachat', split='train')`