fusion-bench 0.2.19__py3-none-any.whl → 0.2.21__py3-none-any.whl

fusion_bench/models/modeling_smile_gemma2/register.py

@@ -0,0 +1,26 @@
+from transformers import (
+    AutoConfig,
+    AutoModel,
+    AutoModelForCausalLM,
+    AutoModelForSequenceClassification,
+    AutoModelForTokenClassification,
+)
+
+from .configuration_smile_gemma2 import SmileGemma2Config
+from .modeling_smile_gemma2 import (
+    SmileGemma2ForCausalLM,
+    SmileGemma2ForSequenceClassification,
+    SmileGemma2ForTokenClassification,
+    SmileGemma2Model,
+    SmileGemma2PreTrainedModel,
+)
+
+AutoConfig.register("smile_gemma2", SmileGemma2Config)
+AutoModel.register(SmileGemma2Config, SmileGemma2Model)
+AutoModelForCausalLM.register(SmileGemma2Config, SmileGemma2ForCausalLM)
+AutoModelForSequenceClassification.register(
+    SmileGemma2Config, SmileGemma2ForSequenceClassification
+)
+AutoModelForTokenClassification.register(
+    SmileGemma2Config, SmileGemma2ForTokenClassification
+)

fusion_bench/models/modeling_smile_llama/configuration_smile_llama.py

@@ -0,0 +1,20 @@
+from transformers import LlamaConfig
+
+
+class SmileLlamaConfig(LlamaConfig):
+    model_type = "smile_llama"
+
+    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)