Update modeling_neollm.py

modeling_neollm.py

@@ -788,7 +788,6 @@ class NeoLLMPreTrainedModel(PreTrainedModel):
             # scale adaptations from data without initial bias
             if hasattr(module, 'multiplier'):
                 module.multiplier.data.fill_(1.0)
-
 class NeoLLMModel(NeoLLMPreTrainedModel):
     """
     NeoLLM base model with transformer decoder architecture.
@@ -842,8 +841,16 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
         attention_mask: Optional[torch.Tensor] = None,
         position_ids: Optional[torch.LongTensor] = None,
         inputs_embeds: Optional[torch.FloatTensor] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
         **kwargs: Unpack[TransformersKwargs],
     ) -> 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")
 
@@ -867,6 +874,7 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
         )
 
         hidden_states = inputs_embeds
+        all_hidden_states = () if output_hidden_states else None
 
         # create position embeddings to be shared across the decoder layers
         position_embeddings = self.rotary_emb(hidden_states, position_ids)
@@ -875,6 +883,9 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
         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,
@@ -890,9 +901,16 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
         # Apply SeeDNorm for final normalization
         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=all_hidden_states,
         )
 
 
@@ -953,75 +971,48 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
         attention_mask: Optional[torch.Tensor] = None,
         position_ids: Optional[torch.LongTensor] = None,
         inputs_embeds: Optional[torch.FloatTensor] = None,
-        output_hidden_states: Optional[bool] = None, 
-        return_dict: Optional[bool] = 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],
-    ) -> 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,
+    ) -> CausalLMOutputWithPast:
+        outputs: BaseModelOutputWithPast = self.model(
+            input_ids=input_ids,
             attention_mask=attention_mask,
-            cache_position=position_ids.squeeze(0),
-            past_key_values=None,
             position_ids=position_ids,
+            inputs_embeds=inputs_embeds,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            **kwargs,
         )
-
-        hidden_states = inputs_embeds
         
-        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,
+        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
             )
-            
-            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,
+            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,
             past_key_values=None,
-            hidden_states=all_hidden_states,
-            attentions=None,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
         )
 
-
 # ==================== AUTOMODEL REGISTRATION ====================
 
 __all__ = [