Update modeling_neollm.py

modeling_neollm.py

@@ -33,7 +33,7 @@ from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
 from transformers.processing_utils import Unpack
 from transformers.utils import TransformersKwargs, logging
 from transformers.utils.generic import check_model_inputs
-from .configuration_neollm import NeoLLMConfig
+from configuration_neollm import NeoLLMConfig
 
 from transformers import AutoConfig, AutoModel, AutoModelForCausalLM
 
@@ -953,42 +953,72 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
         attention_mask: Optional[torch.Tensor] = None,
         position_ids: Optional[torch.LongTensor] = None,
         inputs_embeds: Optional[torch.FloatTensor] = None,
-        labels: Optional[torch.LongTensor] = None,
-        logits_to_keep: Union[int, torch.Tensor] = 0,
+        output_hidden_states: Optional[bool] = None, 
+        return_dict: Optional[bool] = None,           
         **kwargs: Unpack[TransformersKwargs],
-    ) -> CausalLMOutputWithPast:
-        outputs: BaseModelOutputWithPast = self.model(
-            input_ids=input_ids,
+    ) -> BaseModelOutputWithPast:
+        
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None 
+            else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        if position_ids is None:
+            position_ids = torch.arange(0, inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0)
+
+        causal_mask = create_causal_mask(
+            config=self.config,
+            input_embeds=inputs_embeds,
             attention_mask=attention_mask,
+            cache_position=position_ids.squeeze(0),
+            past_key_values=None,
             position_ids=position_ids,
-            inputs_embeds=inputs_embeds,
-            **kwargs,
         )
+
+        hidden_states = inputs_embeds
         
-        hidden_states = outputs.last_hidden_state
-        
-        # CCE Loss computation for training
-        if labels is not None:
-            loss = compute_cce_loss(
-                hidden_states, 
-                labels, 
-                self.lm_head.weight,
-                getattr(self.lm_head, 'bias', None),
-                self.config.pad_token_id
+        all_hidden_states = () if output_hidden_states else None
+
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+        self.first_layer_fan = None
+
+        for decoder_layer in self.layers[: self.config.num_hidden_layers]:
+
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (hidden_states,)
+
+            hidden_states = decoder_layer(
+                hidden_states,
+                position_embeddings=position_embeddings,
+                attention_mask=causal_mask,
+                first_layer_fan=self.first_layer_fan,
+                **kwargs,
             )
-            logits = None
-        else:
-            # Inference mode - compute logits normally
-            slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
-            logits = self.lm_head(hidden_states[:, slice_indices, :])
-            loss = None
-        
-        return CausalLMOutputWithPast(
-            loss=loss,
-            logits=logits,
+            
+            if self.first_layer_fan is None and hasattr(decoder_layer, 'current_layer_fan'):
+                self.first_layer_fan = decoder_layer.current_layer_fan
+
+        hidden_states = self.norm(hidden_states)
+
+        if output_hidden_states:
+            all_hidden_states = all_hidden_states + (hidden_states,)
+
+        if not return_dict:
+            return tuple(v for v in [hidden_states, None, all_hidden_states] if v is not None)
+
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
             past_key_values=None,
-            hidden_states=outputs.hidden_states,
-            attentions=outputs.attentions,
+            hidden_states=all_hidden_states,
+            attentions=None,
         )
 
 
@@ -1009,4 +1039,4 @@ __all__ = [
 # Register the configuration and model for AutoClass support
 AutoConfig.register("neollm", NeoLLMConfig)
 AutoModel.register(NeoLLMConfig, NeoLLMModel)
-AutoModelForCausalLM.register(NeoLLMConfig, NeoLLMForCausalLM)
+AutoModelForCausalLM.register(NeoLLMConfig, NeoLLMForCausalLM)