Upload 2 files

slimshard_gui.py

@@ -107,6 +107,8 @@ class SlimShardGUI:
         self.device_var = tk.StringVar(value="auto")
         self.quant_format_var = tk.StringVar(value="mixed")
         self.expert_keep_var = tk.DoubleVar(value=1.0)
+        self.target_params_b_var = tk.DoubleVar(value=30.0)
+        self.max_experts_var = tk.IntVar(value=0)
 
         self.filter_action_var = tk.StringVar(value="ALL")
         self.search_var = tk.StringVar()
@@ -144,6 +146,8 @@ class SlimShardGUI:
             self.device_var,
             self.quant_format_var,
             self.expert_keep_var,
+            self.target_params_b_var,
+            self.max_experts_var,
         ):
             v.trace_add("write", lambda *_: self._update_cli_preview())
 
@@ -339,7 +343,31 @@ class SlimShardGUI:
             state="readonly",
         ).pack(side=tk.LEFT, padx=(0, 16))
 
-        ttk.Label(row_opt, text="SlimMoE Expert Keep Ratio:").pack(
+        ttk.Label(row_opt, text="Target Params (B, 0=full):").pack(
+            side=tk.LEFT, padx=(0, 6)
+        )
+        ttk.Spinbox(
+            row_opt,
+            from_=0.0,
+            to=600.0,
+            increment=1.0,
+            textvariable=self.target_params_b_var,
+            width=6,
+        ).pack(side=tk.LEFT, padx=(0, 14))
+
+        ttk.Label(row_opt, text="Max Experts/Layer (0=auto):").pack(
+            side=tk.LEFT, padx=(0, 6)
+        )
+        ttk.Spinbox(
+            row_opt,
+            from_=0,
+            to=512,
+            increment=1,
+            textvariable=self.max_experts_var,
+            width=5,
+        ).pack(side=tk.LEFT, padx=(0, 14))
+
+        ttk.Label(row_opt, text="SlimMoE Keep Ratio:").pack(
             side=tk.LEFT, padx=(0, 6)
         )
         ttk.Spinbox(
@@ -349,13 +377,13 @@ class SlimShardGUI:
             increment=0.05,
             textvariable=self.expert_keep_var,
             width=6,
-        ).pack(side=tk.LEFT, padx=(0, 16))
+        ).pack(side=tk.LEFT, padx=(0, 14))
 
         ttk.Checkbutton(
             row_opt,
-            text="Export Ready-to-Load HF Directory (symlinks config/tokenizer + model.safetensors)",
+            text="Export HF Bundle Dir",
             variable=self.auto_bundle_var,
-        ).pack(side=tk.LEFT, padx=(0, 16))
+        ).pack(side=tk.LEFT, padx=(0, 8))
 
         btn_bar = ttk.Frame(opt_frame)
         btn_bar.pack(fill=tk.X, pady=(8, 0))
@@ -665,6 +693,18 @@ class SlimShardGUI:
                 cmd += ["--output", self.output_var.get()]
             if self.auto_bundle_var.get() and self.export_dir_var.get():
                 cmd += ["--export-hf-dir", self.export_dir_var.get()]
+            try:
+                tp = float(self.target_params_b_var.get())
+                if tp > 0.0:
+                    cmd += ["--target-params-b", f"{tp:g}"]
+            except Exception:
+                pass
+            try:
+                me = int(self.max_experts_var.get())
+                if me > 0:
+                    cmd += ["--max-experts", str(me)]
+            except Exception:
+                pass
             cmd += ["--device", self.device_var.get()]
         elif mode == "build":
             if self.model_dir_var.get():
@@ -695,21 +735,37 @@ class SlimShardGUI:
         if not shard or not Path(shard).exists():
             messagebox.showwarning("SlimShard", "Please select a valid .safetensors shard first.")
             return
+        try:
+            tp = float(self.target_params_b_var.get())
+        except Exception:
+            tp = 0.0
+        try:
+            me = int(self.max_experts_var.get())
+        except Exception:
+            me = 0
         try:
             plan = inspect_recovery_plan(
                 shard,
                 index=self.index_var.get().strip() or None,
                 config=self.config_var.get().strip() or None,
+                target_params_b=tp if tp > 0.0 else None,
+                max_experts=me if me > 0 else None,
             )
         except Exception as exc:
             messagebox.showerror("Inspect Failed", str(exc))
             return
 
         archs = ",".join(plan["architectures"]) or plan["model_type"]
+        orig_e = plan.get("orig_num_experts", plan["num_experts"])
+        exp_str = (
+            f"{plan['num_experts']}/{orig_e}"
+            if plan["num_experts"] != orig_e
+            else str(plan["num_experts"])
+        )
         self.arch_summary_var.set(
-            f"{archs} (type={plan['model_type']}, hidden={plan['hidden_size']}, "
-            f"layers={plan['num_hidden_layers']}, experts={plan['num_experts']}, "
-            f"vocab={plan['vocab_size']}, total_tensors={plan['total_tensors']})"
+            f"{archs} (~{plan.get('estimated_params_b', 0)}B params | type={plan['model_type']}, "
+            f"hidden={plan['hidden_size']}, layers={plan['num_hidden_layers']}, "
+            f"experts={exp_str}, vocab={plan['vocab_size']}, tensors={plan['total_tensors']})"
         )
         surv = plan["surviving_layers"]
         miss = plan["missing_layers"]
@@ -787,6 +843,14 @@ class SlimShardGUI:
                         if self.auto_bundle_var.get() and self.export_dir_var.get().strip()
                         else None
                     )
+                    try:
+                        tp = float(self.target_params_b_var.get())
+                    except Exception:
+                        tp = 0.0
+                    try:
+                        me = int(self.max_experts_var.get())
+                    except Exception:
+                        me = 0
                     res = recover_from_shard(
                         shard,
                         idx,
@@ -797,6 +861,8 @@ class SlimShardGUI:
                         progress_cb=progress_cb,
                         stop_event=self._stop_event,
                         export_hf_dir=exp_dir,
+                        target_params_b=tp if tp > 0.0 else None,
+                        max_experts=me if me > 0 else None,
                     )
                     if exp_dir:
                         self._queue.put(("set_test_dir", exp_dir))