Repository: huggingface/transformers
Pull request: Add GptOssForTokenClassification for GPT-OSS models

# What does this PR do?

Add GptOssForTokenClassification for GPT-OSS models following generic practice 

Testing:

Tested with single-node option, loss curve and AUC are all expected with no issue.
- All testes passed
- make fixup passed

Fixes #40189 


## 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. Link #40189 
- [x] 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.

@ArthurZucker



Changed files:
filename: docs/source/en/model_doc/gpt_oss.md
patch: @@ -60,3 +60,8 @@ The original code can be found [here](<INSERT LINK TO GITHUB REPO HERE>).
 
 [[autodoc]] GptOssForSequenceClassification
     - forward
+
+## GptOssForTokenClassification
+
+[[autodoc]] GptOssForTokenClassification
+    - forward

filename: src/transformers/models/auto/modeling_auto.py
patch: @@ -1429,6 +1429,7 @@ class _BaseModelWithGenerate(PreTrainedModel, GenerationMixin):
         ("gpt_bigcode", "GPTBigCodeForTokenClassification"),
         ("gpt_neo", "GPTNeoForTokenClassification"),
         ("gpt_neox", "GPTNeoXForTokenClassification"),
+        ("gpt_oss", "GptOssForTokenClassification"),
         ("helium", "HeliumForTokenClassification"),
         ("ibert", "IBertForTokenClassification"),
         ("layoutlm", "LayoutLMForTokenClassification"),

filename: src/transformers/models/gpt_oss/modeling_gpt_oss.py
patch: @@ -28,7 +28,11 @@
 from ...generation import GenerationMixin
 from ...integrations.hub_kernels import use_kernel_forward_from_hub
 from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask
-from ...modeling_layers import GenericForSequenceClassification, GradientCheckpointingLayer
+from ...modeling_layers import (
+    GenericForSequenceClassification,
+    GenericForTokenClassification,
+    GradientCheckpointingLayer,
+)
 from ...modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast
 from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
 from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
@@ -707,4 +711,14 @@ class GptOssForSequenceClassification(GenericForSequenceClassification, GptOssPr
     pass
 
 
-__all__ = ["GptOssForCausalLM", "GptOssForSequenceClassification", "GptOssModel", "GptOssPreTrainedModel"]
+class GptOssForTokenClassification(GenericForTokenClassification, GptOssPreTrainedModel):
+    pass
+
+
+__all__ = [
+    "GptOssForCausalLM",
+    "GptOssForSequenceClassification",
+    "GptOssForTokenClassification",
+    "GptOssModel",
+    "GptOssPreTrainedModel",
+]

filename: src/transformers/models/gpt_oss/modular_gpt_oss.py
patch: @@ -41,7 +41,12 @@
     LlamaRotaryEmbedding,
     repeat_kv,
 )
-from ..mixtral.modeling_mixtral import MixtralForCausalLM, MixtralForSequenceClassification, MixtralModel
+from ..mixtral.modeling_mixtral import (
+    MixtralForCausalLM,
+    MixtralForSequenceClassification,
+    MixtralForTokenClassification,
+    MixtralModel,
+)
 from ..qwen2.modeling_qwen2 import Qwen2Attention
 from .configuration_gpt_oss import GptOssConfig
 
@@ -447,9 +452,14 @@ class GptOssForSequenceClassification(MixtralForSequenceClassification):
     pass
 
 
+class GptOssForTokenClassification(MixtralForTokenClassification):
+    pass
+
+
 __all__ = [
     "GptOssForCausalLM",
     "GptOssForSequenceClassification",
+    "GptOssForTokenClassification",
     "GptOssModel",
     "GptOssPreTrainedModel",
 ]

filename: tests/models/gpt_oss/test_modeling_gpt_oss.py
patch: @@ -50,6 +50,7 @@
     from transformers import (
         GptOssForCausalLM,
         GptOssForSequenceClassification,
+        GptOssForTokenClassification,
         GptOssModel,
     )
 
@@ -62,12 +63,14 @@ class GptOssModelTester(CausalLMModelTester):
         base_model_class = GptOssModel
         causal_lm_class = GptOssForCausalLM
         sequence_class = GptOssForSequenceClassification
+        token_class = GptOssForTokenClassification
 
     pipeline_model_mapping = (
         {
             "feature-extraction": GptOssModel,
             "text-classification": GptOssForSequenceClassification,
             "text-generation": GptOssForCausalLM,
+            "token-classification": GptOssForTokenClassification,
         }
         if is_torch_available()
         else {}
@@ -77,13 +80,16 @@ class GptOssModelTester(CausalLMModelTester):
 @require_torch
 class GptOssModelTest(CausalLMModelTest, unittest.TestCase):
     all_model_classes = (
-        (GptOssModel, GptOssForCausalLM, GptOssForSequenceClassification) if is_torch_available() else ()
+        (GptOssModel, GptOssForCausalLM, GptOssForSequenceClassification, GptOssForTokenClassification)
+        if is_torch_available()
+        else ()
     )
     pipeline_model_mapping = (
         {
             "feature-extraction": GptOssModel,
             "text-classification": GptOssForSequenceClassification,
             "text-generation": GptOssForCausalLM,
+            "token-classification": GptOssForTokenClassification,
         }
         if is_torch_available()
         else {}