Upload slimshard_gui.py

slimshard_gui.py

@@ -46,10 +46,14 @@ if __package__ in (None, ""):
 from .cuda_backend import cuda_device_name, cuda_is_available
 from .slimshard import (
     HPCConfig,
+    _resolve_shard_paths,
     build_slimshard,
     export_hf_bundle,
     export_safetensors_to_gguf,
+    fetch_minimum_shards_from_hf,
     inspect_recovery_plan,
+    parse_hf_repo_url,
+    plan_minimum_shards_from_hf,
     recover_from_shard,
     regenerate,
     run_deepseek_v41_generate,
@@ -98,6 +102,12 @@ class SlimShardGUI:
 
     def _init_vars(self) -> None:
         self.op_mode_var = tk.StringVar(value="recover")
+        self.hf_url_var = tk.StringVar(
+            value="https://huggingface.co/deepseek-ai/DeepSeek-V4.1-Flash-1017/tree/main"
+        )
+        self.hf_fetch_mode_var = tk.StringVar(value="surgical")
+        self.num_anchor_layers_var = tk.IntVar(value=3)
+        self.anchor_layers_var = tk.StringVar(value="")
         self.shard_var = tk.StringVar()
         self.index_var = tk.StringVar()
         self.config_var = tk.StringVar()
@@ -138,6 +148,10 @@ class SlimShardGUI:
 
         for v in (
             self.op_mode_var,
+            self.hf_url_var,
+            self.hf_fetch_mode_var,
+            self.num_anchor_layers_var,
+            self.anchor_layers_var,
             self.shard_var,
             self.index_var,
             self.config_var,
@@ -253,10 +267,16 @@ class SlimShardGUI:
         )
         ttk.Radiobutton(
             mode_bar,
-            text="CheckFree Single-Shard Recovery (1 Shard -> Full Model)",
+            text="CheckFree Multi/Single-Shard Recovery (Shards -> Full Model)",
             variable=self.op_mode_var,
             value="recover",
         ).pack(side=tk.LEFT, padx=(0, 14))
+        ttk.Radiobutton(
+            mode_bar,
+            text="Retrieve Minimum Shards from HF Link",
+            variable=self.op_mode_var,
+            value="fetch-hf",
+        ).pack(side=tk.LEFT, padx=(0, 14))
         ttk.Radiobutton(
             mode_bar,
             text="Build Self-Contained SlimShard (Multi-Shard -> 1 SlimShard)",
@@ -276,44 +296,71 @@ class SlimShardGUI:
             value="export-gguf",
         ).pack(side=tk.LEFT)
 
+        # Hugging Face Link Minimum-Shard Retriever Bar
+        hf_bar = ttk.Frame(io_frame)
+        hf_bar.grid(row=1, column=0, columnspan=3, sticky=tk.EW, pady=(0, 6))
+        hf_bar.columnconfigure(1, weight=1)
+        ttk.Label(
+            hf_bar, text="Hugging Face Link / Repo:", font=("Sans", 9, "bold")
+        ).grid(row=0, column=0, sticky=tk.W, padx=(0, 6))
+        ttk.Entry(hf_bar, textvariable=self.hf_url_var).grid(
+            row=0, column=1, sticky=tk.EW, padx=(0, 8)
+        )
+        ttk.Combobox(
+            hf_bar,
+            textvariable=self.hf_fetch_mode_var,
+            values=("surgical", "full_shards"),
+            width=12,
+            state="readonly",
+        ).grid(row=0, column=2, padx=(0, 6))
+        ttk.Button(
+            hf_bar, text="Analyze Min Shards", command=self._analyze_hf_min_shards
+        ).grid(row=0, column=3, padx=(0, 6))
+        ttk.Button(
+            hf_bar,
+            text="Retrieve Min Shards",
+            style="Accent.TButton",
+            command=self._retrieve_hf_min_shards,
+        ).grid(row=0, column=4)
+
         self._add_path_row(
             io_frame,
-            1,
-            "Surviving Shard / Input (.safetensors):",
+            2,
+            "Surviving Shard(s) / Input (.safetensors, comma-sep):",
             self.shard_var,
             self._pick_shard,
         )
         self._add_path_row(
             io_frame,
-            2,
+            3,
             "Index JSON (optional, auto-inferred if omitted):",
             self.index_var,
             self._pick_index,
         )
         self._add_path_row(
             io_frame,
-            3,
+            4,
             "Config JSON (auto-detected from folder):",
             self.config_var,
             self._pick_config,
         )
         self._add_path_row(
             io_frame,
-            4,
-            "Source Model Dir (for SlimShard Build):",
+            5,
+            "Source / Download Model Dir:",
             self.model_dir_var,
             self._pick_model_dir,
         )
         self._add_path_row(
             io_frame,
-            5,
-            "Output .safetensors File:",
+            6,
+            "Output .safetensors / .gguf File:",
             self.output_var,
             self._pick_output,
         )
         self._add_path_row(
             io_frame,
-            6,
+            7,
             "Export Ready-to-Run HF Bundle Dir:",
             self.export_dir_var,
             self._pick_export_dir,
@@ -375,6 +422,27 @@ class SlimShardGUI:
             width=5,
         ).pack(side=tk.LEFT, padx=(0, 14))
 
+        ttk.Label(row_opt, text="# Anchor Layers:").pack(
+            side=tk.LEFT, padx=(0, 4)
+        )
+        ttk.Spinbox(
+            row_opt,
+            from_=1,
+            to=128,
+            increment=1,
+            textvariable=self.num_anchor_layers_var,
+            width=4,
+        ).pack(side=tk.LEFT, padx=(0, 10))
+
+        ttk.Label(row_opt, text="Anchor Layers (e.g. 0,14,39):").pack(
+            side=tk.LEFT, padx=(0, 4)
+        )
+        ttk.Entry(
+            row_opt,
+            textvariable=self.anchor_layers_var,
+            width=12,
+        ).pack(side=tk.LEFT, padx=(0, 12))
+
         ttk.Label(row_opt, text="SlimMoE Keep Ratio:").pack(
             side=tk.LEFT, padx=(0, 6)
         )
@@ -638,14 +706,97 @@ class SlimShardGUI:
             self.test_model_dir_var.set(self.export_dir_var.get())
 
     def _pick_shard(self) -> None:
-        p = filedialog.askopenfilename(
-            title="Select Surviving .safetensors Shard",
+        ps = filedialog.askopenfilenames(
+            title="Select Surviving .safetensors Shard(s) (Multi-Select Supported)",
             filetypes=[("SafeTensors", "*.safetensors"), ("All files", "*.*")],
         )
-        if p:
-            self._set_shard_and_companions(Path(p))
+        if ps:
+            paths = [Path(p) for p in ps]
+            self._set_shard_and_companions(paths[0])
+            self.shard_var.set(", ".join(str(p) for p in paths))
             self._inspect_plan()
 
+    def _analyze_hf_min_shards(self) -> None:
+        url = self.hf_url_var.get().strip()
+        if not url:
+            messagebox.showwarning("SlimShard", "Please enter a Hugging Face link or repo ID 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:
+            n_anc = int(self.num_anchor_layers_var.get())
+        except Exception:
+            n_anc = 3
+        anc_str = self.anchor_layers_var.get().strip() or None
+
+        self._log(f"Analyzing minimum shards on Hugging Face for {url}...")
+        try:
+            plan = plan_minimum_shards_from_hf(
+                url,
+                target_params_b=tp if tp > 0.0 else None,
+                max_experts=me if me > 0 else None,
+                num_anchor_layers=n_anc,
+                anchor_layers=anc_str,
+            )
+        except Exception as exc:
+            self._log(f"HF Analyze Failed: {exc}")
+            messagebox.showerror("HF Minimum Shards Analysis Failed", str(exc))
+            return
+
+        repo_name = plan["repo_id"].split("/")[-1]
+        if not self.model_dir_var.get().strip():
+            self.model_dir_var.set(str(Path.cwd() / repo_name))
+        if not self.output_var.get().strip():
+            self.output_var.set(str(Path.cwd() / repo_name / "model.safetensors"))
+        if not self.export_dir_var.get().strip():
+            self.export_dir_var.set(str(Path.cwd() / f"recovered_{repo_name.lower().replace('-', '_')}"))
+
+        min_sh = plan["minimum_shards"]
+        tot_sh = plan["total_repo_shards"]
+        anchors = plan["anchor_layers"]
+        archs = ",".join(plan["architectures"]) or plan["model_type"]
+        self.arch_summary_var.set(
+            f"{plan['repo_id']} ({archs} | ~{plan['estimated_params_b']}B target | "
+            f"layers={plan['num_hidden_layers']}, experts={plan['kept_experts']}/{plan['orig_num_experts']})"
+        )
+        self.layers_summary_var.set(
+            f"Minimum Shards Needed: {len(min_sh)} / {tot_sh} ({', '.join(min_sh)})  |  "
+            f"Anchor Layers ({len(anchors)}): {anchors}"
+        )
+        self.tele_breakdown_var.set(
+            f"min_shards={len(min_sh)}/{tot_sh} | surgical_tensors={plan['surgical_tensor_count']} | "
+            f"full_repo={plan['total_repo_size_gb']} GB"
+        )
+        roles_lines = [
+            f"  - {sh}: {', '.join(plan['shard_roles'].get(sh, []))}" for sh in min_sh
+        ]
+        msg = (
+            f"HF Repo: {plan['repo_id']}@{plan['revision']}\n"
+            f"Minimum Shards Required: {len(min_sh)} of {tot_sh} shards\n"
+            f"Selected Anchor Layers: {anchors}\n"
+            f"Surgical Tensor Count: {plan['surgical_tensor_count']} tensors\n"
+            + "\n".join(roles_lines)
+        )
+        self._log(msg)
+
+    def _retrieve_hf_min_shards(self) -> None:
+        self.op_mode_var.set("fetch-hf")
+        if not self.model_dir_var.get().strip():
+            try:
+                repo_id, _ = parse_hf_repo_url(self.hf_url_var.get().strip())
+                self.model_dir_var.set(str(Path.cwd() / repo_id.split("/")[-1]))
+            except Exception:
+                self.model_dir_var.set(str(Path.cwd() / "hf_min_shards"))
+        if not self.output_var.get().strip():
+            self.output_var.set(str(Path(self.model_dir_var.get().strip()) / "model.safetensors"))
+        self._start_run()
+
     def _pick_index(self) -> None:
         p = filedialog.askopenfilename(
             title="Select model.safetensors.index.json",