🎞️ Support sequence classification models in clone_chat_template (#4097)

Author: qgallouedec

tests/test_dataset_formatting.py

@@ -15,7 +15,7 @@
 from typing import Callable
 
 from datasets import Dataset, load_dataset
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoModelForSequenceClassification, AutoTokenizer
 
 from trl.extras.dataset_formatting import get_formatting_func_from_dataset
 from trl.models.utils import ChatMlSpecialTokens, clone_chat_template, setup_chat_format
@@ -159,47 +159,59 @@ def test_example_with_setup_model(self):
 
 
 class CloneChatTemplateTestCase(TrlTestCase):
-    def setUp(self):
-        super().setUp()
+    def test_clone(self):
         # This tokenizer doesn't have a chat_template by default
-        self.tokenizer = AutoTokenizer.from_pretrained("trl-internal-testing/tiny-BloomForCausalLM")
-        self.model = AutoModelForCausalLM.from_pretrained("trl-internal-testing/tiny-BloomForCausalLM")
+        tokenizer = AutoTokenizer.from_pretrained("trl-internal-testing/tiny-BloomForCausalLM")
+        model = AutoModelForCausalLM.from_pretrained("trl-internal-testing/tiny-BloomForCausalLM")
         # This one has a chat_template by default
-        self.source = "trl-internal-testing/tiny-Qwen3ForCausalLM"
-
-    def test_clone(self):
-        _, modified_tokenizer, _ = clone_chat_template(self.model, self.tokenizer, self.source)
+        source = "trl-internal-testing/tiny-Qwen3ForCausalLM"
+        _, modified_tokenizer, _ = clone_chat_template(model, tokenizer, source)
 
         # Check if special tokens are correctly set
         self.assertEqual(modified_tokenizer.eos_token, "<|im_end|>")
 
     def test_clone_with_resize(self):
+        # This tokenizer doesn't have a chat_template by default
+        tokenizer = AutoTokenizer.from_pretrained("trl-internal-testing/tiny-BloomForCausalLM")
+        model = AutoModelForCausalLM.from_pretrained("trl-internal-testing/tiny-BloomForCausalLM")
+        # This one has a chat_template by default
+        source = "trl-internal-testing/tiny-Qwen3ForCausalLM"
         modified_model, modified_tokenizer, _ = clone_chat_template(
-            self.model, self.tokenizer, self.source, resize_to_multiple_of=123
+            model, tokenizer, source, resize_to_multiple_of=123
         )
 
         # Check that the input embeddings have been resized to a multiple of 123
         self.assertEqual((modified_model.vocab_size % 123), 0)
         # Check that the input embeddings size matches the tokenizer vocabulary size
-        self.assertEqual(self.model.vocab_size, len(modified_tokenizer.vocab))
+        self.assertEqual(model.vocab_size, len(modified_tokenizer.vocab))
 
     def test_clone_with_resize_and_extra_tokens_already_in_vocab(self):
+        # This tokenizer doesn't have a chat_template by default
+        tokenizer = AutoTokenizer.from_pretrained("trl-internal-testing/tiny-BloomForCausalLM")
+        model = AutoModelForCausalLM.from_pretrained("trl-internal-testing/tiny-BloomForCausalLM")
+        # This one has a chat_template by default
+        source = "trl-internal-testing/tiny-Qwen3ForCausalLM"
         # This will add <extra_id_0>, <extra_id_1>, ... to the tokenizer
         modified_model, modified_tokenizer, _ = clone_chat_template(
-            self.model, self.tokenizer, self.source, resize_to_multiple_of=123
+            model, tokenizer, source, resize_to_multiple_of=123
         )
         # Try if we can resize a tokenizer that already has extra these extra tokens
         modified_model, modified_tokenizer, _ = clone_chat_template(
-            modified_model, modified_tokenizer, self.source, resize_to_multiple_of=124
+            modified_model, modified_tokenizer, source, resize_to_multiple_of=124
         )
 
         # Check that the input embeddings have been resized to a multiple of 123
         self.assertEqual((modified_model.vocab_size % 124), 0)
         # Check that the input embeddings size matches the tokenizer vocabulary size
-        self.assertEqual(self.model.vocab_size, len(modified_tokenizer.vocab))
+        self.assertEqual(model.vocab_size, len(modified_tokenizer.vocab))
 
     def test_apply_new_chat_template(self):
-        _, modified_tokenizer, _ = clone_chat_template(self.model, self.tokenizer, self.source)
+        # This tokenizer doesn't have a chat_template by default
+        tokenizer = AutoTokenizer.from_pretrained("trl-internal-testing/tiny-BloomForCausalLM")
+        model = AutoModelForCausalLM.from_pretrained("trl-internal-testing/tiny-BloomForCausalLM")
+        # This one has a chat_template by default
+        source = "trl-internal-testing/tiny-Qwen3ForCausalLM"
+        _, modified_tokenizer, _ = clone_chat_template(model, tokenizer, source)
         messages = [
             {"role": "system", "content": "You are helpful"},
             {"role": "user", "content": "Hello"},
@@ -211,3 +223,16 @@ def test_apply_new_chat_template(self):
             prompt,
             "<|im_start|>system\nYou are helpful<|im_end|>\n<|im_start|>user\nHello<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\nHi, how can I help you?<|im_end|>\n",
         )
+
+    def test_clone_with_sequence_classification_model(self):
+        # This tokenizer doesn't have a chat_template by default
+        tokenizer = AutoTokenizer.from_pretrained("trl-internal-testing/tiny-GptNeoXForSequenceClassification")
+        model = AutoModelForSequenceClassification.from_pretrained(
+            "trl-internal-testing/tiny-GptNeoXForSequenceClassification"
+        )
+        # This one has a chat_template by default
+        source = "trl-internal-testing/tiny-Qwen3ForCausalLM"
+        _, modified_tokenizer, _ = clone_chat_template(model, tokenizer, source)
+
+        # Check if special tokens are correctly set
+        self.assertEqual(modified_tokenizer.eos_token, "<|im_end|>")

trl/models/utils.py

@@ -223,7 +223,8 @@ def clone_chat_template(
     # Set the EOS token from the source tokenizer (important for generation)
     tokenizer.eos_token = tokenizer_source.eos_token
     model.config.eos_token_id = tokenizer.eos_token_id
-    model.generation_config.eos_token_id = tokenizer.eos_token_id
+    if model.generation_config is not None:  # for SequenceClassification models, generation_config is None
+        model.generation_config.eos_token_id = tokenizer.eos_token_id
 
     # Resize model embeddings to include any new tokens, optionally rounding up to a multiple
     model.resize_token_embeddings(