Add Qwen Moe (#2163)

* qwen moe init commit

* wip

* wip

* weight conversion wip

* weight matching complete

* update the docstrings + configs

* remove incorrect import

* wip

* updates

* updates

* updates

* bug fix

* add aux loss

* address comments

* causal lm test

* add tests

* add docstrings

* small bug fix

* bug fix in aux loss

* update init.py

* address comments

* bug fixes

* address comments

* code format

* update

* update

* address comments

* update

* format

---------

Co-authored-by: Anshuman Mishra <[email protected]>
Co-authored-by: Anshuman Mishra <[email protected]>

Author: kanpuriyanawab

keras_hub/api/models/__init__.py

@@ -420,6 +420,15 @@
 from keras_hub.src.models.qwen.qwen_tokenizer import (
     QwenTokenizer as QwenTokenizer,
 )
+from keras_hub.src.models.qwen_moe.qwen_moe_backbone import (
+    QwenMoeBackbone as QwenMoeBackbone,
+)
+from keras_hub.src.models.qwen_moe.qwen_moe_causal_lm import (
+    QwenMoeCausalLM as QwenMoeCausalLM,
+)
+from keras_hub.src.models.qwen_moe.qwen_moe_causal_lm_preprocessor import (
+    QwenMoeCausalLMPreprocessor as QwenMoeCausalLMPreprocessor,
+)
 from keras_hub.src.models.resnet.resnet_backbone import (
     ResNetBackbone as ResNetBackbone,
 )

keras_hub/api/tokenizers/__init__.py

@@ -68,6 +68,9 @@
 from keras_hub.src.models.qwen.qwen_tokenizer import (
     QwenTokenizer as QwenTokenizer,
 )
+from keras_hub.src.models.qwen_moe.qwen_moe_tokenizer import (
+    QwenMoeTokenizer as QwenMoeTokenizer,
+)
 from keras_hub.src.models.roberta.roberta_tokenizer import (
     RobertaTokenizer as RobertaTokenizer,
 )

keras_hub/src/models/qwen/qwen_attention.py

@@ -287,7 +287,9 @@ def _compute_attention(
         if self.use_sliding_window_attention:
             attention_mask = self._mask_sliding_window(
                 attention_mask,
-                cache_update_index=cache_update_index,
+                cache_update_index=cache_update_index
+                if cache_update_index
+                else 0,
             )
         attention_scores = self._masked_softmax(
             attention_scores, attention_mask