Qwen - Fix Preset Loader + Add Causal LM Test (#2193)

* load tie embedding param from config

* add causal lm test for qwen + bug fix

* address comment

Author: kanpuriyanawab

keras_hub/src/models/qwen/qwen_backbone.py

@@ -168,7 +168,7 @@ def __init__(
         self.layer_norm_epsilon = layer_norm_epsilon
         self.dropout = dropout
         self.tie_word_embeddings = tie_word_embeddings
-        self.use_sliding_window_attention = (use_sliding_window_attention,)
+        self.use_sliding_window_attention = use_sliding_window_attention
         self.sliding_window_size = sliding_window_size
 
     def get_config(self):

keras_hub/src/models/qwen/qwen_causal_lm_test.py

@@ -0,0 +1,123 @@
+from unittest.mock import patch
+
+import pytest
+from keras import ops
+
+from keras_hub.src.models.qwen.qwen_backbone import QwenBackbone
+from keras_hub.src.models.qwen.qwen_causal_lm import QwenCausalLM
+from keras_hub.src.models.qwen.qwen_causal_lm_preprocessor import (
+    QwenCausalLMPreprocessor,
+)
+from keras_hub.src.models.qwen.qwen_tokenizer import QwenTokenizer
+from keras_hub.src.tests.test_case import TestCase
+
+
+class QwenCausalLMTest(TestCase):
+    def setUp(self):
+        self.vocab = ["!", "air", "Ġair", "plane", "Ġat", "port"]
+        self.vocab += ["<|endoftext|>"]
+        self.vocab += ["<|eot_id|>"]
+        self.vocab = dict([(token, i) for i, token in enumerate(self.vocab)])
+        self.merges = ["Ġ a", "Ġ t", "Ġ i", "Ġ b", "a i", "p l", "n e"]
+        self.merges += ["Ġa t", "p o", "r t", "Ġt h", "ai r", "pl a", "po rt"]
+        self.merges += ["Ġai r", "Ġa i", "pla ne"]
+        self.preprocessor = QwenCausalLMPreprocessor(
+            QwenTokenizer(vocabulary=self.vocab, merges=self.merges),
+            sequence_length=7,
+        )
+        self.backbone = QwenBackbone(
+            vocabulary_size=self.preprocessor.tokenizer.vocabulary_size(),
+            num_layers=2,
+            num_query_heads=4,
+            num_key_value_heads=2,
+            hidden_dim=8,
+            intermediate_dim=16,
+        )
+        self.init_kwargs = {
+            "preprocessor": self.preprocessor,
+            "backbone": self.backbone,
+        }
+        self.train_data = ([" airplane at airport", " airplane at airport"],)
+        self.input_data = self.preprocessor(*self.train_data)[0]
+
+    def test_causal_lm_basics(self):
+        self.run_task_test(
+            cls=QwenCausalLM,
+            init_kwargs=self.init_kwargs,
+            train_data=self.train_data,
+            expected_output_shape=(2, 7, 8),
+        )
+
+    def test_generate(self):
+        causal_lm = QwenCausalLM(**self.init_kwargs)
+        # String input.
+        prompt = " airplane at airport"
+        output = causal_lm.generate(" airplane at airport")
+        self.assertTrue(prompt in output)
+        # Int tensor input.
+        prompt_ids = self.preprocessor.generate_preprocess([prompt])
+        causal_lm.preprocessor = None
+        outputs = causal_lm.generate(prompt_ids, stop_token_ids=None)
+        # Assert prompt is in output in token id space.
+        self.assertAllEqual(
+            outputs["token_ids"][:, :5],
+            prompt_ids["token_ids"][:, :5],
+        )
+        self.assertAllEqual(
+            outputs["padding_mask"][:, :5],
+            prompt_ids["padding_mask"][:, :5],
+        )
+
+    def test_generate_strip_prompt(self):
+        causal_lm = QwenCausalLM(**self.init_kwargs)
+        prompt = " airplane at airport"
+        output = causal_lm.generate(prompt, strip_prompt=True)
+        self.assertFalse(output.startswith(prompt))
+
+    def test_early_stopping(self):
+        causal_lm = QwenCausalLM(**self.init_kwargs)
+        call_with_cache = causal_lm.call_with_cache
+
+        def wrapper(*args, **kwargs):
+            """Modify output logits to always favor end_token_id"""
+            logits, hidden_states, cache = call_with_cache(*args, **kwargs)
+            index = self.preprocessor.tokenizer.end_token_id
+            update = ops.ones_like(logits)[:, :, index] * 1.0e9
+            update = ops.expand_dims(update, axis=-1)
+            logits = ops.slice_update(logits, (0, 0, index), update)
+            return logits, hidden_states, cache
+
+        with patch.object(causal_lm, "call_with_cache", wraps=wrapper):
+            prompt = [" airplane at airport", " airplane"]
+            output = causal_lm.generate(prompt)
+            # We should immediately abort and output the prompt.
+            self.assertEqual(prompt, output)
+
+    def test_generate_compilation(self):
+        causal_lm = QwenCausalLM(**self.init_kwargs)
+        # Assert we do not recompile with successive calls.
+        causal_lm.generate(" airplane at airport")
+        first_fn = causal_lm.generate_function
+        causal_lm.generate(" airplane at airport")
+        second_fn = causal_lm.generate_function
+        self.assertEqual(first_fn, second_fn)
+        # Assert we do recompile after compile is called.
+        causal_lm.compile(sampler="greedy")
+        self.assertIsNone(causal_lm.generate_function)
+
+    @pytest.mark.large
+    def test_saved_model(self):
+        self.run_model_saving_test(
+            cls=QwenCausalLM,
+            init_kwargs=self.init_kwargs,
+            input_data=self.input_data,
+        )
+
+    @pytest.mark.extra_large
+    def test_all_presets(self):
+        for preset in QwenCausalLM.presets:
+            self.run_preset_test(
+                cls=QwenCausalLM,
+                preset=preset,
+                input_data=self.input_data,
+            )

keras_hub/src/utils/transformers/convert_qwen.py

@@ -18,6 +18,7 @@ def convert_backbone_config(transformers_config):
         "rope_max_wavelength": transformers_config["rope_theta"],
         "use_sliding_window": transformers_config["use_sliding_window"],
         "sliding_window_size": transformers_config["sliding_window"],
+        "tie_word_embeddings": transformers_config["tie_word_embeddings"],
     }