DFlash2-only installs use the draft-vocabulary subset the person chose in setup: the choice (en / fr / cyrillic) is copied beside the DFlash2 pack, refreshed when it changes and on update and start, kept by the drafter's fetch; tests for DFlash2-only and MTP-folder installs

setup.py

@@ -7307,6 +7307,8 @@ def _start(cfg_path: Path, port: int | None, gpu: int | list | None = None, open
     cfg_path.touch()                                     # the most recently used model
     if "--mtp" in cfg["args"][:-1] and not own_path_config(cfg):
         refresh_draft_vocab(Path(cfg["args"][cfg["args"].index("--mtp") + 1]), cfg.get("draft_vocab", "cjk"))
+        from strata_control import drafting
+        refresh_dflash2_vocab(drafting.data_dir_of(cfg["args"]), cfg.get("draft_vocab", "cjk"))
     cmd = [sys.executable, str(ROOT / "serve" / "server.py"), "--engine", "strata", "--config", str(cfg_path),
            "--port", str(port or cfg.get("port", 8080))]
     if cfg.get("backend") == "hip":                    # AMD, numbered as HIP numbers them (setup's KFD order)
@@ -7690,13 +7692,24 @@ def mtp_corrupt(mtp: Path, env=None) -> bool:
     return r.returncode == 3
 
 
-def provide_dflash2_vocab(data: Path) -> None:
-    """A DFlash2-only install has no MTP folder to hold the draft vocabulary: it is kept beside the DFlash2 pack (strata_control/drafting.py provide_vocab)."""
+def provide_dflash2_vocab(data: Path, choice: str = "cjk") -> None:
+    """A DFlash2-only install has no MTP folder to hold the draft vocabulary: it is kept beside the DFlash2 pack (strata_control/drafting.py provide_vocab), as the subset the person chose
+    in setup (--draft-vocab, saved in the model's config; the default file when none)."""
     from strata_control import drafting
-    if drafting.provide_vocab(data) is None:
+    if drafting.provide_vocab(data, choice or "cjk") is None:
         warn("the draft vocabulary file (data/draft_vocab.bin) is missing from this Strata folder, so DFlash2 cannot start; update Strata")
 
 
+def refresh_dflash2_vocab(data: Path | None, choice: str = "cjk") -> None:
+    """The DFlash2 copy of the draft subset (beside its pack) follows the person's choice too, when that copy exists: an update, a start and setup run again keep or apply the
+    choice for a DFlash2-only install, which has no MTP folder for refresh_draft_vocab to work on."""
+    if data is None:
+        return
+    from strata_control import drafting
+    if drafting.vocab_target(data).is_file():
+        drafting.refresh_vocab(data, choice or "cjk")
+
+
 def refresh_draft_vocab(rt: Path, choice: str = "cjk") -> None:
     """The draft layer's token subset in the MTP folder: `cjk` (data/draft_vocab.bin, since 0.1.27, #137), `en`
     (data/draft_vocab_en.bin, the English/code subset before it: ~110 MiB less VRAM, English answers 1-2% faster) or
@@ -7730,6 +7743,13 @@ def refresh_shared_draft_vocab(cfg: dict, done: set) -> None:
     args = cfg.get("args") or []
     if "--mtp" not in args[:-1]:
         return
+    from strata_control import drafting
+    data = drafting.data_dir_of(args)
+    if data is not None:                                  # the DFlash2 copy, once per data folder, the config used last first like the MTP folder below
+        dkey = ("dflash2", os.path.normcase(os.path.abspath(data)))
+        if dkey not in done:
+            done.add(dkey)
+            refresh_dflash2_vocab(data, cfg.get("draft_vocab", "cjk"))
     rt = Path(args[args.index("--mtp") + 1])
     key = os.path.normcase(os.path.abspath(rt))
     if key in done:
@@ -9365,7 +9385,7 @@ def mac_control_setup(a) -> int:
         use_mtp = False
     elif use_mtp and getattr(a, "drafter", None) == "dflash2" and not (mtp / "mtp-q2_0.gguf").exists():
         ok("DFlash2 is the drafter: the MTP draft layer's ~5 GB are not downloaded (choosing the MTP later fetches it once)")
-        provide_dflash2_vocab(data)
+        provide_dflash2_vocab(data, a.draft_vocab or "cjk")
         use_mtp = False
     if use_mtp and not (mtp / "mtp-q2_0.gguf").exists():
         say("  The MTP draft layer (speculative decoding) comes from the original Qwen checkpoint: only its ~5 GB of MTP tensors are downloaded.")
@@ -10450,7 +10470,7 @@ def main() -> int:
         ok("the draft layer (MTP) is off: its ~5 GB are not downloaded (`--mtp on` fetches it later)")   # round 8, N3
     elif getattr(a, "drafter", None) == "dflash2" and not (rt / "experts.bin").exists():
         ok("DFlash2 is the drafter: the MTP draft layer's ~5 GB are not downloaded (choosing the MTP later fetches it once)")
-        provide_dflash2_vocab(data)
+        provide_dflash2_vocab(data, a.draft_vocab or saved_draft_vocab(home() / f"strata-{tag.lower()}.json") or "cjk")
     elif corrupt or not (rt / "experts.bin").exists():
         say("  The MTP draft layer (speculative decoding, ~2x faster output) comes from the original Qwen checkpoint:")
         say("  only its ~5 GB of MTP tensors are downloaded.")
@@ -10464,6 +10484,7 @@ def main() -> int:
     if (rt / "experts.bin").exists():
         refresh_draft_vocab(rt, draft_vocab or "cjk")
         ok(f"MTP draft layer: {rt}")
+    refresh_dflash2_vocab(data, draft_vocab or "cjk")        # an existing DFlash2 copy follows the choice too
     for line in draft_vocab_note(gpu.get("vram_gb", 0.0), draft_vocab):   # #474: a recommendation, nothing changes
         say("  " + line)
 

strata_control/drafting.py

@@ -552,9 +552,19 @@ def mtp_ready(args: list[Any]) -> bool | None:
 VOCAB_NAME = "draft_vocab.bin"
 
 
-def shipped_vocab() -> Path:
-    """The default draft vocabulary Strata ships with (``data/draft_vocab.bin``, also in the packaged Control)."""
-    return Path(__file__).resolve().parents[1] / "data" / VOCAB_NAME
+#: The subsets a person can choose in setup (``--draft-vocab``), by name: the files Strata ships in data/ (setup.DRAFT_VOCABS names the same ones).
+VOCAB_FILES = {"cjk": VOCAB_NAME, "en": "draft_vocab_en.bin", "cyrillic": "draft_vocab_cyrillic.bin", "fr": "draft_vocab_fr.bin"}
+#: The sha256 of the 27-Han-token subset up to 0.1.26 (setup.OLD_DRAFT_VOCABS): a copy that is any shipped subset may be replaced by the person's later choice, a made-by-hand one never.
+OLD_VOCAB_SHA256 = {"369151522226a5edaa5f12cfd1e2ae7db8f4fbdbd222f3dcf327dced9597fb25"}
+
+
+def shipped_vocab(choice: str | None = None) -> Path:
+    """The draft vocabulary Strata ships with (``data/draft_vocab.bin``, also in the packaged Control), or the shipped subset named ``choice`` (cjk, en, cyrillic, fr)."""
+    return Path(__file__).resolve().parents[1] / "data" / VOCAB_FILES.get(choice or "cjk", VOCAB_NAME)
+
+
+def _sha256(path: Path) -> str:
+    return hashlib.sha256(path.read_bytes()).hexdigest()
 
 
 def vocab_target(data_dir: str | Path) -> Path:
@@ -562,12 +572,40 @@ def vocab_target(data_dir: str | Path) -> Path:
     return Path(data_dir) / FOLDER / VOCAB_NAME
 
 
-def provide_vocab(data_dir: str | Path | None) -> Path | None:
-    """Make sure the DFlash2 folder holds the draft vocabulary (a local copy of the file Strata ships; no network). Called by the drafter's fetch and by
-    setup's DFlash2 step, so a fresh install that selected only DFlash2 has everything the engine needs. None when it cannot be had."""
+def refresh_vocab(data_dir: str | Path | None, choice: str = "cjk") -> Path | None:
+    """Put the subset the person chose in setup (``--draft-vocab``) beside the DFlash2 pack, where a DFlash2 start reads it when the MTP folder has none. The copy is made when
+    it is missing and replaced when it is a shipped subset other than the chosen one (a later choice, or an install from before the choice reached DFlash2); a subset made by hand
+    stays. Same rule as setup.refresh_draft_vocab for the MTP folder. None when it cannot be had."""
     if not data_dir:
         return None
     target = vocab_target(data_dir)
+    try:
+        new = shipped_vocab(choice)
+        if not new.is_file():
+            return provide_vocab(data_dir)
+        if target.is_file() and target.stat().st_size > 0:
+            old = _sha256(target)
+            shipped_hashes = OLD_VOCAB_SHA256 | {_sha256(shipped_vocab(c)) for c in VOCAB_FILES if shipped_vocab(c).is_file()}
+            if old == _sha256(new) or old not in shipped_hashes:
+                return target
+        target.parent.mkdir(parents=True, exist_ok=True)
+        tmp = target.with_suffix(".tmp")
+        tmp.write_bytes(new.read_bytes())
+        tmp.replace(target)
+        return target
+    except OSError:
+        return None
+
+
+def provide_vocab(data_dir: str | Path | None, choice: str | None = None) -> Path | None:
+    """Make sure the DFlash2 folder holds the draft vocabulary (a local copy of the file Strata ships; no network). Called by the drafter's fetch and by
+    setup's DFlash2 step, so a fresh install that selected only DFlash2 has everything the engine needs. None when it cannot be had.
+    Without ``choice`` a copy that is there is never touched (the fetch must not undo the person's setup choice); with it (setup) see refresh_vocab."""
+    if not data_dir:
+        return None
+    if choice is not None:
+        return refresh_vocab(data_dir, choice)
+    target = vocab_target(data_dir)
     try:
         if target.is_file() and target.stat().st_size > 0:
             return target

tools/test_dflash2_vocab.py

@@ -0,0 +1,158 @@
+"""The draft-vocabulary subset a person chose in setup (--draft-vocab en / fr / cyrillic) reaches DFlash2, which is the default drafter: on an install that has NO MTP folder
+(a fresh DFlash2-only install) as well as on one that has it.  Pure file work in temporary folders, no GPU, no downloads, no prompts.
+
+    python -m unittest tools.test_dflash2_vocab -v
+
+What is checked, per choice (cjk = the default file, en, cyrillic, fr):
+  - setup's vocab step (provide_dflash2_vocab / refresh_draft_vocab / refresh_shared_draft_vocab) and then strata_control.drafting.apply: the file `--mtp-draft-vocab` points at is the chosen subset
+    byte for byte, for an explicit `dflash2` and for `auto`;
+  - a later choice replaces a shipped subset, a subset made by hand stays;
+  - an update (refresh_shared_draft_vocab) and the drafter's own fetch (provide_vocab without a choice) keep the person's choice.
+"""
+from __future__ import annotations
+
+import hashlib
+import inspect
+import json
+import sys
+import tempfile
+import unittest
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+import setup  # noqa: E402
+from strata_control import drafting  # noqa: E402
+
+CHOICES = ("cjk", "en", "cyrillic", "fr")
+
+
+def sha(p) -> str:
+    return hashlib.sha256(Path(p).read_bytes()).hexdigest()
+
+
+def shipped(choice: str) -> Path:
+    return ROOT / "data" / setup.DRAFT_VOCABS[choice]
+
+
+def vocab_arg(cfg: dict) -> Path:
+    a = cfg["args"]
+    return Path(a[a.index("--mtp-draft-vocab") + 1])
+
+
+@unittest.skipUnless(all(shipped(c).exists() for c in CHOICES), "the shipped draft subsets are not in this checkout")
+class Base(unittest.TestCase):
+    with_mtp = False
+
+    def setUp(self):
+        self._tmp = tempfile.TemporaryDirectory()
+        self.addCleanup(self._tmp.cleanup)
+        self.data = Path(self._tmp.name) / "data"
+        (self.data / "packs" / "m").mkdir(parents=True)
+        (self.data / "dflash2").mkdir(parents=True)
+        (self.data / "dflash2" / drafting.PACK_NAME).write_bytes(b"\0" * ((1 << 20) + 16))   # "downloaded" (size is all drafting.pack_ready looks at)
+        self.rt = self.data / "mtp" / "rt"
+        if self.with_mtp:
+            self.rt.mkdir(parents=True)
+            (self.rt / "experts.bin").write_bytes(b"e")
+            (self.rt / "draft_vocab.bin").write_bytes(shipped("cjk").read_bytes())      # what an install of the MTP layer holds first
+
+    # ---- what setup does for a model with this draft_vocab choice
+    def setup_step(self, choice: str) -> None:
+        """setup.py's vocab step on an install with a DFlash2 drafter: the DFlash2 copy (beside the pack) and, when there is an MTP folder, the folder's file."""
+        takes_choice = "choice" in inspect.signature(setup.provide_dflash2_vocab).parameters
+        if takes_choice:
+            setup.provide_dflash2_vocab(self.data, choice)
+        else:
+            setup.provide_dflash2_vocab(self.data)               # the code before this branch
+        if (self.rt / "experts.bin").exists():
+            setup.refresh_draft_vocab(self.rt, choice)
+
+    def config(self, mode: str = "dflash2", choice: str | None = None) -> dict:
+        args = ["--pack", str(self.data / "packs" / "m"), "--native", "x.gguf", "--mtp", str(self.rt), "--spec", "4", "--spec-min-p", "0.5"]
+        cfg = {"args": args, "model_name": "qwen3.8-flash-next-iq3_s", "drafting": mode}
+        if choice:
+            cfg["draft_vocab"] = choice
+        return cfg
+
+    def applied(self, mode: str = "dflash2") -> dict:
+        cfg = self.config(mode)
+        d = drafting.apply(cfg, data_dir=self.data, family="qwen", backend="cuda", archs=["sm_121"], slots=0)
+        self.assertEqual(d.drafter, "dflash2", d)
+        return cfg
+
+
+class DFlash2Only(Base):
+    with_mtp = False
+
+    def test_each_choice_reaches_the_engine(self):
+        for choice in CHOICES:
+            with self.subTest(choice=choice):
+                self.setup_step(choice)
+                self.assertEqual(sha(vocab_arg(self.applied())), sha(shipped(choice)), f"--mtp-draft-vocab is not the {choice} subset")
+
+    def test_auto_uses_the_choice(self):
+        for choice in CHOICES:
+            with self.subTest(choice=choice):
+                self.setup_step(choice)
+                self.assertEqual(sha(vocab_arg(self.applied("auto"))), sha(shipped(choice)))
+
+    def test_a_later_choice_replaces_a_shipped_subset_and_a_hand_made_one_stays(self):
+        self.setup_step("fr")
+        self.assertEqual(sha(vocab_arg(self.applied())), sha(shipped("fr")))
+        self.setup_step("en")
+        self.assertEqual(sha(vocab_arg(self.applied())), sha(shipped("en")))
+        self.setup_step("cjk")
+        self.assertEqual(sha(vocab_arg(self.applied())), sha(shipped("cjk")))
+        mine = drafting.vocab_target(self.data)
+        mine.write_bytes(bytes([1, 0, 0, 0]))
+        self.setup_step("fr")
+        self.assertEqual(mine.read_bytes(), bytes([1, 0, 0, 0]), "a subset made by hand is kept")
+
+    def test_an_update_and_the_drafters_fetch_keep_the_choice(self):
+        self.setup_step("fr")
+        cfg = self.config(choice="fr")
+        setup.refresh_shared_draft_vocab(cfg, set())                       # setup.py --update, the saved config names fr
+        self.assertEqual(sha(vocab_arg(self.applied())), sha(shipped("fr")))
+        drafting.provide_vocab(self.data)                                  # the fetch (Load, `strata drafting fetch`) never overwrites a copy that is there
+        self.assertEqual(sha(vocab_arg(self.applied())), sha(shipped("fr")))
+
+    def test_an_update_applies_the_saved_choice_to_an_older_copy(self):
+        drafting.provide_vocab(self.data)                                  # an install made before the choice reached DFlash2: the default copy
+        self.assertEqual(sha(drafting.vocab_target(self.data)), sha(shipped("cjk")))
+        setup.refresh_shared_draft_vocab(self.config(choice="cyrillic"), set())
+        self.assertEqual(sha(vocab_arg(self.applied())), sha(shipped("cyrillic")))
+
+    def test_no_choice_means_the_default_file(self):
+        setup.refresh_shared_draft_vocab(self.config(), set())             # no "draft_vocab" key: nothing to apply, nothing broken
+        drafting.provide_vocab(self.data)
+        self.assertEqual(sha(vocab_arg(self.applied())), sha(shipped("cjk")))
+
+
+class WithMtpFolder(Base):
+    with_mtp = True
+
+    def test_each_choice_reaches_the_engine(self):
+        for choice in CHOICES:
+            with self.subTest(choice=choice):
+                self.setup_step(choice)
+                self.assertEqual(sha(vocab_arg(self.applied())), sha(shipped(choice)))
+                self.assertEqual(sha(self.rt / "draft_vocab.bin"), sha(shipped(choice)), "the MTP layer's file follows too")
+                if drafting.vocab_target(self.data).exists():
+                    self.assertEqual(sha(drafting.vocab_target(self.data)), sha(shipped(choice)), "the DFlash2 copy agrees with the folder")
+
+    def test_auto_uses_the_choice(self):
+        self.setup_step("fr")
+        self.assertEqual(sha(vocab_arg(self.applied("auto"))), sha(shipped("fr")))
+
+    def test_the_two_files_never_disagree_after_an_update(self):
+        self.setup_step("en")
+        setup.refresh_shared_draft_vocab(self.config(choice="fr"), set())
+        self.assertEqual(sha(self.rt / "draft_vocab.bin"), sha(shipped("fr")))
+        if drafting.vocab_target(self.data).exists():
+            self.assertEqual(sha(drafting.vocab_target(self.data)), sha(shipped("fr")))
+        self.assertEqual(sha(vocab_arg(self.applied())), sha(shipped("fr")))
+
+
+if __name__ == "__main__":
+    unittest.main()