Update qwen3_moe_model.py

qwen3_moe_model.py

@@ -1,4 +1,3 @@
-import os
 import torch
 from torch.nn import ModuleDict
 from transformers import (
@@ -9,92 +8,86 @@ from transformers import (
 )
 from .qwen3_moe_config import Qwen3MoEConfig
 from transformers.modeling_outputs import CausalLMOutput
-from typing import Optional, Tuple, Union, Dict
-
+from typing import Optional, Tuple, Union
 
 
 class Qwen3MoEForCausalLM(PreTrainedModel):
     config_class = Qwen3MoEConfig
 
-    def __init__(self, config: Qwen3MoEConfig):
+    def __init__(self, config: Qwen3MoEConfig, hub_repo_id: str):
+        """
+        hub_repo_id: str, ví dụ "beyoru/Qwen3-MaCoTo"
+        """
         super().__init__(config)
+        self.hub_repo_id = hub_repo_id
+
+        # Load router model + tokenizer từ subfolder "router"
         self.router = AutoModelForSequenceClassification.from_pretrained(
-            config.router_model_path,
+            hub_repo_id,
+            subfolder="router",
             torch_dtype=config.torch_dtype,
-            trust_remote_code=True,
-            local_files_only=True
+            trust_remote_code=True
         )
-
         self.router_tokenizer = AutoTokenizer.from_pretrained(
-            config.router_model_path,
-            trust_remote_code=True,
-            local_files_only=True
+            hub_repo_id,
+            subfolder="router",
+            trust_remote_code=True
         )
 
+        # Load expert models từ các subfolder tương ứng
         self.experts = ModuleDict({
             label: AutoModelForCausalLM.from_pretrained(
-                path,
+                hub_repo_id,
+                subfolder=folder,  # folder con code/math/if/tool
                 torch_dtype=config.torch_dtype,
-                trust_remote_code=True,
-                local_files_only=True
+                trust_remote_code=True
             )
-            for label, path in config.expert_model_paths.items()
+            for label, folder in config.expert_model_paths.items()
         })
 
+        # Load tokenizer chung của expert từ root repo
         self.expert_tokenizer = AutoTokenizer.from_pretrained(
-            config.tokenizer_path,
-            trust_remote_code=True,
-            local_files_only=True
+            hub_repo_id,
+            subfolder=".",
+            trust_remote_code=True
         )
 
     @classmethod
-    def from_pretrained(cls, pretrained_dir: str, config: Optional[Qwen3MoEConfig] = None, **kwargs):
+    def from_pretrained(cls, hub_repo_id: str, config: Optional[Qwen3MoEConfig] = None, **kwargs):
         if config is None:
-            config = Qwen3MoEConfig.from_pretrained(pretrained_dir)
-
-        base = pretrained_dir
-        config.router_model_path = os.path.join(base, config.router_model_path)
-        config.expert_model_paths = {
-            label: os.path.join(base, path)
-            for label, path in config.expert_model_paths.items()
-        }
-        config.tokenizer_path = os.path.join(base, config.tokenizer_path)
-
-        return cls(config)
+            # load config.json từ root repo
+            config = Qwen3MoEConfig.from_pretrained(hub_repo_id, **kwargs)
+        return cls(config, hub_repo_id=hub_repo_id)
 
     def get_tokenizer(self):
         return self.expert_tokenizer
 
     def route(self, plain_text: str) -> str:
+        """
+        Dùng router model để chọn expert phù hợp dựa trên input text.
+        """
         with torch.no_grad():
             inputs = self.router_tokenizer(plain_text, return_tensors="pt").to(self.router.device)
             logits = self.router(**inputs).logits
-
             if logits.dim() == 2:
                 class_id = torch.argmax(logits, dim=-1).item()
                 return self.config.labels[class_id]
-
         return self.config.labels[0]
 
-    def generate(
-        self,
-        text: str,
-        max_new_tokens: int = 50,
-        **kwargs
-    ) -> torch.LongTensor:
-        # 1. Route using router tokenizer
+    def generate(self, text: str, max_new_tokens: int = 50, **kwargs) -> torch.LongTensor:
+        """
+        Generate text từ expert được chọn bởi router.
+        """
         plain_text = text
         if "<|im_start|>" in plain_text:
-          temp = plain_text.split("<|im_start|>")[-2]
-          plain_text = temp[:temp.find("<|im_end|>")][4:]
+            temp = plain_text.split("<|im_start|>")[-2]
+            plain_text = temp[:temp.find("<|im_end|>")][4:]
 
         label = self.route(plain_text)
         expert = self.experts[label]
 
-        # 2. Tokenize once with the expert tokenizer
         inputs = self.expert_tokenizer(text, return_tensors="pt").to(expert.device)
 
-        # 3. Generate using selected expert
         return expert.generate(
             input_ids=inputs.input_ids,
             attention_mask=inputs.attention_mask,
@@ -109,4 +102,4 @@ class Qwen3MoEForCausalLM(PreTrainedModel):
         labels: Optional[torch.LongTensor] = None,
         **kwargs
     ) -> Union[Tuple, CausalLMOutput]:
-        raise NotImplementedError("Use `generate(text=...)` instead for inference.")
+        raise NotImplementedError("Use `generate(text=...)` instead for inference.")