fix: use real hf tokenizer vocab size when adding new trainable token

aphp · Dec 4, 2023 · 3a632d4 · 3a632d4
1 parent 7a1b075
commit 3a632d4
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/edsnlp/pipes/trainable/embeddings/transformer/transformer.py b/edsnlp/pipes/trainable/embeddings/transformer/transformer.py
@@ -160,7 +160,9 @@ def __init__(
                 )
             )
             # and add a new entry to the model's embeddings
-            self.transformer.resize_token_embeddings(len(self.tokenizer))
+            self.transformer.resize_token_embeddings(
+                max(self.tokenizer.vocab.values()) + 1
+            )
 
     def to_disk(self, path, *, exclude: Optional[Set[str]]):
         repr_id = object.__repr__(self)