fix(data): Handle integer labels in DataCollatorWithFlattening
Author: yashisthebatman
Pull request: yes

# What does this PR do?

This PR fixes a `TypeError` in `DataCollatorWithFlattening` that occurs when processing datasets with integer labels, which is a common format for sequence classification tasks.

The collator was written with the assumption that the `labels` field would always be a sequence, and it attempted to slice it (`sample["labels"][1:]`). When `sample["labels"]` is an integer, this causes the application to crash.

The fix introduces a check to determine if the `labels` are sequence-like.
- If the labels are a sequence (the original expected format), the collator preserves the exact original logic.
- If the labels are integers, the collator now correctly broadcasts the integer label across all tokens of the corresponding `input_ids`.

This change makes the collator more robust to common data formats without altering its behavior for existing use cases, thereby preventing unexpected crashes. A new test case has been added to verify this fix.

Fixes #41652


## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [X] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [X] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [X] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

This change relates to data collators, which are often used with the Trainer. Tagging from the original issue and the trainer section.
@ArthurZucker @SunMarc

Comments:
@LPirch Thanks for the review i have removed the unecessary comments and also have also followed the more consistent naming scheme but regarding the Why not batch["labels"].append(sample["labels"])? :
The entire purpose of DataCollatorWithFlattening is to create one single long sequence, so the labels tensor must have the same length as the final input_ids tensor. Your suggestion (.append()) would create mismatched lengths (e.g., input_ids of length 7 but labels of length 2), which would immediately crash any model's loss function.

Let me give a bit more context to prevent misunderstandings: In sequence classification, you don't have a label per token but a label per sample. Now, when sequences get very long, you might want to use flash attention (as in my case). The only way to process a batch of long sequences is to pack them using the flattening collator and run them through the language model. The needs to be unpacked using the `position_ids` tensor (see example below) to recover the original samples.

In my case, you then need a pooling that creates an embedding per sequence and a classification head (linear layers) producing an output score. The loss function is then just a regular classification loss.

unpacking example:
```python
def pool_last_token_flat(last_hidden_states: torch.Tensor, position_ids: torch.Tensor) -> torch.Tensor:
    end_idx = torch.argwhere(position_ids == 0)[:, 1][1:]
    end_idx = torch.cat([end_idx, torch.tensor([last_hidden_states.shape[1] - 1], device=last_hidden_states.device)])
    return last_hidden_states[0, end_idx]
```

This pooling takes the last layers hidden vector of the last token of each (sub-)sequence in the flattened sequence. 

TLDR: We don't create flat sequences for the sake of flattening but because flash attention 2 can't handle padding. Recovering the batch structure is required anyways in the case of sequence classification.



Hi @LPirch ,
Thank you for that context! That makes total sense. I didn't consider the flash attention -> pack -> unpack -> pool workflow.
You're right, my original fix is useful for preventing a crash but actively unhelpful for your use case, since you'd just have to undo it. What do you think of this?
We add a boolean flag to the collator: pack_sequence_labels=False.
Default (False): It does my original fix (broadcasts the label). This is the "safe" option for anyone not doing what you're doing, and it guarantees input_ids and labels have the same length.
Your case (True): You'd set pack_sequence_labels=True. The collator would then just collect the integer labels into a simple list [label1, label2, ...], giving you the (batch_size,) tensor you need for your pooling logic.
This feels like the cleanest way to support both scenarios without breaking things for anyone.
If that sounds good to you, I'll go ahead and implement it.

Sounds good! Thanks for your support :raised_hands: 

@LPirch I have done all the changes that we discussed please do check it 

@LPirch thank you for accepting my pr i am new to open source could you tell me where i could improve myself or any suggestions that could help me become a better contributor 

@yashisthebatman I think some steps in the CI pipeline failed (code quality). Can you fix those?

About your question: You're doing well! Keep the positivity in communication and stay curious / be open to learn from others and keep a focus on understanding the bigger picture before finalizing a PR - that are the key ingredients if you ask me. The rest is experience.

@LPirch fixed those formatting errors this time it must work thank you for the patience

@LPirch all the tests have passed successfully 

@LPirch could you please merge this problem I have had successfully solved all the formatting issues 

Just to clarify: I'm not a maintainer of this project. I can't assign, resolve or merge anything. I'm just a user who reported an issue.

From now on, @Rocketknight1 would need to take the lead. I suggest we wait patiently until he finds the time and see if he has additional feedback. Also, let's not bloat the conversation too much.

Hey all! I'm sorry for the delays, but I finally got around to doing a proper review! I have one main question: I realize that packing is used as a way of improving performance with some FlashAttention backends, but do our `SequenceClassification` models actually support packed inputs?

Packing works for CLM training because the objective for those models is token-level, so you can just pack sequences together and concatenate `input_ids` with a separator, then do the same with their labels. However, most `SequenceClassification` models emit one logit distribution per sequence, not per token. If you pack inputs together, it's not clear to me how sequence classification training is supposed to work, and I think the user would have to unpack the outputs before passing them to the model anyway. What exactly is the workflow supposed to look like when a `Flattening` collator is combined with a `SequenceClassification` model?

 @Rocketknight1 A `SequenceClassification` model won't know what to do with a packed input straight out of the box. This feature is really for an advanced, manual workflow that some users need for performance, especially with things like Flash Attention 2 that don't like padding.

So the idea is, a user would grab this collator with `pack_sequence_labels=True`. This gives them:
1.  The packed `input_ids` and `position_ids` they need to feed the base model.
2.  And crucially, a clean `labels` tensor of shape `(batch_size,)`. **This is the missing piece my PR provides.**

From there, the workflow looks like this on the user's end:

First, they'd get the raw hidden states by calling just the base model, bypassing the classification head:
```python
# User's code
outputs = model.base_model(packed_input_ids)
last_hidden_state = outputs.last_hidden_state # A single giant tensor
```

Next, they have to do the manual part: use the `position_ids` to find the last token of each original sequence in that giant tensor and pool them. This turns the flattened output back into a standard `(batch_size, hidden_size)` tensor.

And now they have everything they need. They can feed their pooled outputs into `model.classifier`, and the `labels` tensor from the collator will match up perfectly for the loss calculation.

So basically this PR doesn't make the models  understand packed inputs rather provides the correctly shaped `labels` tensor that makes this whole manual packing/unpacking dance possible for sequence classification.


Hmmn. I see what you're going for, but there's a distribution shift there: If a model is fine-tuned for sequence classification, it will generally not see concatenated documents with separators during training. Packing works for CLM training because documents are often concatenated across document boundaries anyway, so the model is familiar with that, but I expect the performance of a sequence classification model to sharply drop if you use the hidden states from a later sequence that follows multiple earlier sequences and separator tokens.

In other words, packed sequences aren't actually independent and the later sequences can still attend to the earlier ones, they just learn not to do that if they see enough concatenated docs in training!

As such, I'm not sure we want to add this feature - I think it's very niche and likely to be confusing to users who expect it to work, when really it requires a lot of manual effort to separate the outputs for the classification head, and probably also degrades performance in most cases as well.

Oh, I was not aware that there are cross-sequence information leaks in a packed sequence. I didn't think of this as an option because (as you said) this would also contaminate gradients on regular token predictions, i.e. attending to tokens of previous sequences in a batch. From a user perspective, I find it unintuitive that flash attention implicitly sacrifices soundness for efficiency but that maybe belongs to a different issue.

Generally, I can understand your argument though and I would be fine if you close it. I have already monkey-patched the collator to my needs in the meantime and only thought it would be a useful feature for others as well.

Thanks for clarifying your position, @LPirch. Given the valid concerns about attention leakage and the potential for user confusion, I agree that the best path forward is to close this PR. This was a really insightful discussion, and I appreciate the detailed feedback from both of you. I'll go ahead and close this now.

Hang on, my mistake! I dug into the FlashAttention docs and I was wrong; as long as `position_ids` and `cu_seq_lens` are set correctly then there shouldn't be information leakage and performance shouldn't degrade. I didn't realize that `cu_seq_lens` controlled attention masking like that.

I'm debating whether it's worth adding this after all, in that case

@Rocketknight1 Hey, wow, thank you so much for digging into the FlashAttention docs and for that update! That's a super important clarification and really great to know.
No problem at all, I'll keep the PR open. Please take your time to think it over. If you do decide it's a worthwhile addition, just let me know if there's anything else you'd like me to do. No rush at all!

@Rocketknight1 any update ?

Will try to review it tomorrow!