Upload slimshard_gui.py

slimshard_gui.py

@@ -51,6 +51,7 @@ from .slimshard import (
     inspect_recovery_plan,
     recover_from_shard,
     regenerate,
+    run_deepseek_v41_generate,
     slimshard_info,
 )
 
@@ -922,19 +923,52 @@ class SlimShardGUI:
 
         def test_worker() -> None:
             try:
+                import json
                 import torch
                 import torch.nn.functional as F
                 from safetensors.torch import safe_open
                 from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
 
                 t0 = time.time()
+                mpath = Path(mdir)
+                if mpath.is_file() and mpath.suffix == ".safetensors":
+                    st_file = mpath
+                    base_dir = mpath.parent
+                else:
+                    base_dir = mpath
+                    st_file = base_dir / "model.safetensors"
+                    if not st_file.exists():
+                        cands = sorted(base_dir.glob("*.safetensors"))
+                        if cands:
+                            st_file = cands[0]
+
+                cfg_file = base_dir / "config.json"
+                raw_cfg = json.loads(cfg_file.read_text()) if cfg_file.exists() else {}
+                if raw_cfg.get("model_type") == "deepseek_v41":
+                    decoded, out_list = run_deepseek_v41_generate(
+                        st_file,
+                        cfg_file,
+                        base_dir,
+                        prompt=prompt,
+                        max_new_tokens=max_new,
+                        use_chat_template=use_chat,
+                        device=self.device_var.get(),
+                    )
+                    dt_s = time.time() - t0
+                    self._queue.put(
+                        (
+                            "test_done",
+                            f"Completed in {dt_s:.2f}s:\n{decoded}\n\nToken IDs: {out_list}",
+                        )
+                    )
+                    return
+
                 tok = AutoTokenizer.from_pretrained(
-                    mdir, trust_remote_code=True, local_files_only=True
+                    str(base_dir), trust_remote_code=True, local_files_only=True
                 )
                 cfg = AutoConfig.from_pretrained(
-                    mdir, trust_remote_code=True, local_files_only=True
+                    str(base_dir), trust_remote_code=True, local_files_only=True
                 )
-                st_file = Path(mdir) / "model.safetensors"
                 dev = "cuda:0" if torch.cuda.is_available() else "cpu"
 
                 # Fast on-demand expert streaming for MXFP8 MoE models, or standard load for generic models