Temporary fix to the TunixMaxTextAdapter padding logic.

Author: igorts-git

src/maxtext/integration/tunix/tunix_adapter.py

@@ -25,6 +25,7 @@
 
 from flax import nnx
 from jax import Array
+import jax.numpy as jnp
 from maxtext.checkpoint_conversion.utils.hf_model_configs import HF_MODEL_CONFIGS  # pylint: disable=ungrouped-imports
 from maxtext.integration.tunix.utils import VllmWeightMapping
 from maxtext.models.models import Transformer
@@ -63,6 +64,11 @@ def __call__(
     """Forward compatible with Tunix Trainers default loss.
     Returns logits, None.
     """
+    # Bug B Fix: Dynamically build segment IDs based on padding tokens to block causal attention leaks
+    # Extract pad_id dynamically from the left-padding prefix
+    pad_id = input_tokens[0, 0]
+    decoder_segment_ids = jnp.where(input_tokens == pad_id, 0, 1).astype(jnp.int32)
+
     logits = self.base(
         decoder_input_tokens=input_tokens,
         decoder_positions=positions,