Add backwards compatibility for transformers v4.57

Author: Maxusmusti

pyproject.toml

@@ -34,7 +34,7 @@ dependencies = [
     "deprecated",
     "numpy",
     "datasets",
-    "transformers>=5.0.0",
+    "transformers>=4.57.0",
     "ninja",
     "numba>=0.62.0",
     "rich",

research_scratch/fsdp1_dummy_script.py

@@ -53,7 +53,7 @@ def wrap_model_with_fsdp1(model: torch.nn.Module) -> FSDP:
     # Get transformer layer class for auto-wrap
     auto_wrap_policy = None
     if hasattr(model, "_no_split_modules") and model._no_split_modules:
-        layer_name = model._no_split_modules[0]
+        layer_name = next(iter(model._no_split_modules))
         layer_cls = get_module_class_from_name(model, layer_name)
         if layer_cls:
             auto_wrap_policy = functools.partial(transformer_auto_wrap_policy, transformer_layer_cls={layer_cls})

research_scratch/fsdp1_wrapper.py

@@ -38,7 +38,7 @@ def wrap_fsdp1(model: torch.nn.Module) -> torch.nn.Module:
     model.gradient_checkpointing_enable()
 
     # Determine the block class to auto-wrap (first no-split module)
-    block_name = model._no_split_modules[0]
+    block_name = next(iter(model._no_split_modules))
     block_cls = get_module_class_from_name(model, block_name)
     if block_cls is None:
         raise ValueError(f"Could not find module class named {block_name}")

scripts/process_data.py

@@ -1,10 +1,15 @@
 from pathlib import Path
 
 import numpy as np
+import transformers
 import typer
 from datasets import load_dataset
 from transformers import AutoTokenizer
 
+# Transformers v5 renamed 'additional_special_tokens' to 'extra_special_tokens'
+_TRANSFORMERS_V5 = int(transformers.__version__.split(".")[0]) >= 5
+_SPECIAL_TOKENS_KEY = "extra_special_tokens" if _TRANSFORMERS_V5 else "additional_special_tokens"
+
 app = typer.Typer()
 
 # Import from main codebase instead of duplicating
@@ -235,7 +240,7 @@ def process_data(
 ):
     tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
     assistant_tk_ids, user_tk_ids = infer_special_token_sequences(tokenizer)
-    tokenizer.add_special_tokens({"extra_special_tokens": [string_for_printing_masks]})
+    tokenizer.add_special_tokens({_SPECIAL_TOKENS_KEY: [string_for_printing_masks]})
     string_for_printing_masks_tk = tokenizer.encode(string_for_printing_masks, add_special_tokens=False)[0]
 
     dataset = load_dataset("json", data_files=input_jsonl, split="train")