Update app.py

app.py

@@ -94,19 +94,30 @@ def load_model(model_path, progress=gr.Progress()):
         progress(0.3, desc="Loading tokenizer...")
         config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
         if 'MoE' in model_path:
-            config_moe = config
-            config_moe.auto_map["AutoConfig"] = "./configuration_upcycling_qwen2_moe.UpcyclingQwen2MoeConfig"
-            config_moe.auto_map["AutoModelForCausalLM"] = "./modeling_upcycling_qwen2_moe.UpcyclingQwen2MoeForCausalLM"
+            from configuration_upcycling_qwen2_moe import UpcyclingQwen2MoeConfig
+            config = UpcyclingQwen2MoeConfig.from_pretrained(model_path, trust_remote_code=True)
+            # config_moe.auto_map["AutoConfig"] = "./configuration_upcycling_qwen2_moe.UpcyclingQwen2MoeConfig"
+            # config_moe.auto_map["AutoModelForCausalLM"] = "./modeling_upcycling_qwen2_moe.UpcyclingQwen2MoeForCausalLM"
         current_tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False,trust_remote_code=True)
         
         progress(0.5, desc="Loading model...")
-        current_model = AutoModelForCausalLM.from_pretrained(
-            model_path, 
-            device_map="auto", 
-            torch_dtype=torch.float16,
-            config=config_moe if 'MoE' in model_path else config,
-            trust_remote_code=True
-        )
+        if 'MoE' in model_path:
+            from modeling_upcycling_qwen2_moe import UpcyclingQwen2MoeForCausalLM
+            current_model = UpcyclingQwen2MoeForCausalLM.from_pretrained(
+                model_path, 
+                device_map="auto", 
+                torch_dtype=torch.float16,
+                config=config,
+                trust_remote_code=True
+            )
+        else:
+            current_model = AutoModelForCausalLM.from_pretrained(
+                model_path, 
+                device_map="auto", 
+                torch_dtype=torch.float16,
+                config=config,
+                trust_remote_code=True
+            )
         
         current_model_path = model_path
         progress(1.0, desc="Model loading complete!")