fusion-bench 0.2.13__py3-none-any.whl → 0.2.15__py3-none-any.whl

fusion_bench/models/modeling_smile_qwen2/__init__.py

@@ -0,0 +1,8 @@
+from . import register
+from .configuration_smile_qwen2 import SmileQwen2Config
+from .modeling_smile_qwen2 import (
+    SmileQwen2ForCausalLM,
+    SmileQwen2ForQuestionAnswering,
+    SmileQwen2ForSequenceClassification,
+    SmileQwen2Model,
+)

fusion_bench/models/modeling_smile_qwen2/configuration_smile_qwen2.py

@@ -0,0 +1,21 @@
+from transformers import PretrainedConfig
+from transformers.models.qwen2.configuration_qwen2 import Qwen2Config
+
+
+class SmileQwen2Config(Qwen2Config):
+    model_type = "smile_qwen2"
+
+    def __init__(
+        self,
+        num_experts_per_tok: int = 1,
+        rank_of_router: int = None,
+        rank_of_expert: int = None,
+        num_local_experts: int = None,
+        **kwargs,
+    ):
+        self.num_experts_per_tok = num_experts_per_tok
+        self.rank_of_router = rank_of_router
+        self.rank_of_expert = rank_of_expert
+        self.num_local_experts = num_local_experts
+
+        super().__init__(**kwargs)