update to transformers 5 format

configuration_llada.py

@@ -2,6 +2,7 @@
 LLaDA configuration
 """
 from transformers import AutoConfig, PretrainedConfig
+from huggingface_hub.dataclasses import strict
 
 from enum import Enum
 from os import PathLike
@@ -430,9 +431,11 @@ class ActivationCheckpointingStrategy(StrEnum):
     """
 
 
+@strict
 class LLaDAConfig(PretrainedConfig):
     model_type = "llada"
     keys_to_ignore_at_inference = ["past_key_values"]  # TODO: confirm
+    tie_word_embeddings: bool = False
 
     def __init__(self, use_cache: bool = False, **kwargs):
         model_config = ModelConfig()

modeling_llada.py

@@ -27,6 +27,7 @@ import torch.nn as nn
 import torch.nn.functional as F
 from torch import einsum
 from transformers import PreTrainedModel
+from transformers.utils import can_return_tuple
 from transformers.modeling_outputs import CausalLMOutputWithPast
 from transformers.models.auto import AutoModel
 from transformers.cache_utils import Cache
@@ -1393,6 +1394,7 @@ class LLaDAModelLM(PreTrainedModel):
     config_class = LLaDAConfig
     base_model_prefix = "model"
     _no_split_modules = ["LLaDABlock", "LLaDASequentialBlock", "LLaDALlamaBlock"]
+    _tied_weights_keys = {"model.transformer.ff_out.weight": "model.transformer.wte"}
 
     def __init__(self, config: LLaDAConfig, model: Optional[LLaDAModel] = None, init_params: bool = False):
         super().__init__(config)
@@ -1404,10 +1406,12 @@ class LLaDAModelLM(PreTrainedModel):
             self.model = LLaDAModel(model_config, init_params=init_params)
         else:
             self.model = model
+        self.post_init()
 
+    @can_return_tuple
     def forward(
         self,
-        input_ids: torch.LongTensor = None,
+        input_ids: torch.LongTensor | None = None,
         inputs_embeds: Optional[torch.FloatTensor] = None,
         attention_mask: Optional[torch.Tensor] = None,
         attention_bias: Optional[torch.Tensor] = None,
@@ -1416,8 +1420,7 @@ class LLaDAModelLM(PreTrainedModel):
         use_cache: Optional[bool] = None,
         output_attentions: Optional[bool] = None,
         output_hidden_states: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
-        cache_position: Optional[Cache] = None,  # This is a hack mitigation of an issue in transformers `4.39.x`
+        **kwargs
     ) -> Union[Tuple, CausalLMOutputWithPast]:
         if use_cache is None:
             use_cache = self.config.use_cache
@@ -1425,7 +1428,6 @@ class LLaDAModelLM(PreTrainedModel):
         if output_attentions:
             raise ValueError("output_attentions is not yet supported in LLaDA")
 
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
         # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
         outputs = self.model.forward(
@@ -1441,14 +1443,10 @@ class LLaDAModelLM(PreTrainedModel):
         logits = outputs.logits
         hidden_states = outputs.hidden_states
 
-        loss = None
         if labels is not None:
             import warnings
             warnings.warn("Note that for LLaDA, you cannot calculate the loss here.", UserWarning)
-        if not return_dict:
-            output = (logits,) + outputs[1:]
-            return (loss,) + output if loss is not None else output
-
+        
         return CausalLMOutputWithPast(
             logits=logits,
             past_key_values=outputs.attn_key_values,
@@ -1498,9 +1496,5 @@ class LLaDAModelLM(PreTrainedModel):
         else:
             self.model.transformer.ff_out = value
 
-    def tie_weights(self):
-        if self.config.weight_tying:
-            self.model.transformer.ff_out = self.model.transformer.wte
-
 # Register the model so that it is available for transformer pipelines, auto-loading, etc.
 AutoModel.register(LLaDAConfig, LLaDAModelLM)