add resize token embeddings option

turbo-llm · Nov 26, 2024 · ead02ad · ead02ad
1 parent 809eb0e
commit ead02ad
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/turbo_alignment/common/tf/loaders/model/model.py b/turbo_alignment/common/tf/loaders/model/model.py
@@ -85,7 +85,8 @@ def load_model(
     if model_settings.transformers_settings.load_in_8bit:
         model = prepare_model_for_int8_training(model)
 
-    model.resize_token_embeddings(len(tokenizer))
+    if model_settings.resize_token_embeddings:
+        model.resize_token_embeddings(len(tokenizer))
 
     if model_settings.embeddings_initialization_strategy is not None:
         with torch.no_grad():

diff --git a/turbo_alignment/settings/model.py b/turbo_alignment/settings/model.py
@@ -41,6 +41,8 @@ class PreTrainedModelSettings(ExtraFieldsNotAllowedBaseModel):
 
     transformers_settings: ModelTransformersSettings
 
+    resize_token_embeddings: bool = False
+
     embeddings_initialization_strategy: dict[str, str] | None = None
 
     liger_kernels_settings: LigerKernelSettings | None = None