Upload modeling_ivme_conversate.py with huggingface_hub

modeling_ivme_conversate.py

@@ -180,6 +180,17 @@ class IvmeConversateModel(PreTrainedModel):
             nn.init.normal_(module.weight, mean=0.0, std=0.02)
 
     def forward(self, input_ids, **kwargs):
+        # Defensive truncation: keep only the most recent context_len
+        # tokens. This is a second safety layer beyond exposing
+        # max_position_embeddings in the config (see
+        # IvmeConversateConfig) - if any caller still manages to pass a
+        # longer sequence, truncate cleanly to the trailing window
+        # rather than crash inside RoPE with a cryptic tensor-shape
+        # mismatch (confirmed reproducible: a 2048-token input against
+        # this model's 1024 context previously crashed deep inside
+        # apply_rope with no indication of the real cause).
+        if input_ids.shape[1] > self._context_len:
+            input_ids = input_ids[:, -self._context_len:]
         x = self.tok_embed(input_ids)
         cos, sin = precompute_rope(self._head_dim, self._context_len,
                                     theta=self._rope_theta, device=x.device)
@@ -231,6 +242,18 @@ class IvmeConversateForCausalLM(PreTrainedModel, GenerationMixin):
         # attention_mask accepted for interface compatibility but not
         # used - this model is trained/used causal-only with no padding
         # within a sequence (see ShardedTokenDataset / infer scripts).
+        #
+        # If input_ids exceeds context_len, self.model() truncates it
+        # internally (see IvmeConversateModel.forward). labels must be
+        # truncated identically here, BEFORE the loss is computed -
+        # otherwise a caller passing both an over-length input_ids and
+        # matching labels (e.g. a perplexity/loglikelihood task) would
+        # hit a shape mismatch between logits and labels, or worse, a
+        # silent misalignment if the shapes happened to coincide.
+        context_len = self.config.max_position_embeddings
+        if input_ids.shape[1] > context_len and labels is not None:
+            labels = labels[:, -context_len:]
+
         hidden = self.model(input_ids)
         logits = self.lm_head(hidden)