setup: --draft-vocab fr - the English/code subset plus the tokens of French text (#597)

French answers drafted poorly with the shipped subsets: 23.6% of the token occurrences of French text fall outside
draft_vocab_en.bin, and the draft head cannot propose those (no subset adds the accented words and word pieces).

- tools/draft_vocab.py: a coverage mode. --corpus FILE... adds to the base the most frequent tokens of the corpus that
  together cover --coverage (default 0.99) of its occurrences, and prints the share outside the base and the new list.
- data/draft_vocab_fr.bin: draft_vocab_en.bin plus 5,686 tokens (46,211 ids), from 10 French Wikipedia articles
  (France, Paris, Napoleon Ier, Langue francaise, Union europeenne, Litterature francaise, Cuisine francaise,
  Economie de la France, Revolution francaise, Seconde Guerre mondiale; 1.24 M characters, 326,926 tokens) at 0.99:
    python tools/draft_vocab.py --gguf <IQ3_XXS shard 1> --base data/draft_vocab_en.bin --corpus <the articles> \
        --coverage 0.99 --out data/draft_vocab_fr.bin
  Outside the list on that corpus: 0.74% (23.65% with en).
- setup.py: "fr" in DRAFT_VOCABS / --draft-vocab (~151 MiB head on IQ3_S, between en and cyrillic). The sha256 check
  knows it as a shipped subset, so a later --draft-vocab cjk/en replaces it and a subset made by hand stays.
- tests: tools/test_draft_vocab.py (the coverage mode), tools/test_setup_draft_vocab.py (fr replaces a shipped subset
  and back, keeps the en ids first in their order, its MiB estimate).

The default subset (cjk) is unchanged.

setup.py

@@ -2737,7 +2737,8 @@ def start(cfg_path: Path, port: int | None, gpu: int | list | None = None, open_
 OLD_DRAFT_VOCABS = {"369151522226a5edaa5f12cfd1e2ae7db8f4fbdbd222f3dcf327dced9597fb25"}   # to 0.1.26: 27 Han tokens
 
 
-DRAFT_VOCABS = {"cjk": "draft_vocab.bin", "en": "draft_vocab_en.bin", "cyrillic": "draft_vocab_cyrillic.bin"}
+DRAFT_VOCABS = {"cjk": "draft_vocab.bin", "en": "draft_vocab_en.bin", "cyrillic": "draft_vocab_cyrillic.bin",
+                "fr": "draft_vocab_fr.bin"}
 
 
 def saved_draft_vocab(cfg_path: Path) -> str | None:
@@ -2750,7 +2751,7 @@ def saved_draft_vocab(cfg_path: Path) -> str | None:
     return v if v in DRAFT_VOCABS else None
 
 
-DRAFT_VOCAB_MIB = {"cjk": 348, "cyrillic": 193, "en": 133}   # the draft head's VRAM per subset (IQ3_S: the largest)
+DRAFT_VOCAB_MIB = {"cjk": 348, "cyrillic": 193, "fr": 151, "en": 133}   # the draft head's VRAM per subset (IQ3_S: the largest)
 SMALL_DRAFT_VRAM_GB = 14   # #474: below this the default subset's head can be what does not fit
 
 
@@ -2825,7 +2826,7 @@ 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
     `cyrillic` (data/draft_vocab_cyrillic.bin: English/code and the whole Cyrillic script, for Ukrainian, Russian,
-    Bulgarian, Serbian... answers).
+    Bulgarian, Serbian... answers) or `fr` (data/draft_vocab_fr.bin: English/code and the tokens of French text, #597).
     Copied when missing or when a shipped subset other than the chosen one is there; a subset made by hand is kept."""
     new, dst = ROOT / "data" / DRAFT_VOCABS.get(choice, "draft_vocab.bin"), rt / "draft_vocab.bin"
     if not new.exists() or not rt.is_dir():
@@ -2837,7 +2838,8 @@ def refresh_draft_vocab(rt: Path, choice: str = "cjk") -> None:
         if old not in shipped or old == hashlib.sha256(new.read_bytes()).hexdigest():
             return
         ok("draft layer: the token subset " + {"cjk": "with Chinese, Japanese and Korean",
-                                               "cyrillic": "with the Cyrillic script"}.get(choice,
+                                               "cyrillic": "with the Cyrillic script",
+                                               "fr": "for French"}.get(choice,
                                                                                           "for English and code (less VRAM)"))
     shutil.copyfile(new, dst)
 
@@ -2968,7 +2970,8 @@ def main() -> int:
     ap.add_argument("--draft-vocab", choices=list(DRAFT_VOCABS),
                     help="the draft layer's tokens: cjk = with Chinese, Japanese and Korean (default), en = English "
                          "and code only (~110 MiB less VRAM, English answers 1-2%% faster), cyrillic = English, code "
-                         "and the Cyrillic script (Ukrainian, Russian... answers decode ~30%% faster)")
+                         "and the Cyrillic script (Ukrainian, Russian... answers decode ~30%% faster), fr = English, "
+                         "code and French (French answers draft more)")
     ap.add_argument("--low-ram", choices=["auto", "on", "off", "resident", "mmap"], default="auto",
                     help="read the model's experts from one file in its folder instead of copying them all into RAM "
                          "(for a PC with a big GPU and little RAM); auto: when the experts would not fit the RAM. In "
@@ -3532,7 +3535,7 @@ def main() -> int:
              "--out", str(mtp / "mtp-q2_0.gguf")], env=env)
         run([sys.executable, str(ROOT / "tools" / "mtp_rt.py"), "--gguf", str(mtp / "mtp-q2_0.gguf"), "--out", str(rt)],
             env=env)
-    # a setup run again without --draft-vocab keeps the subset this model's config chose before (cyrillic, en)
+    # a setup run again without --draft-vocab keeps the subset this model's config chose before (cyrillic, fr, en)
     draft_vocab = a.draft_vocab or saved_draft_vocab(ROOT / f"strata-{tag.lower()}.json")
     refresh_draft_vocab(rt, draft_vocab or "cjk")
     ok(f"MTP draft layer: {rt}")

tools/draft_vocab.py

@@ -10,12 +10,23 @@ Chinese drafted almost nothing (#137).  This adds whole scripts to a subset:
 
 --add takes han, kana, hangul, cjk_punct, cjk (those four) or cyrillic.  The base ids keep their order; the added ones follow in
 id order.  --stats prints what a subset holds.
+
+A language written in a script the base already holds (French, Spanish, German... in Latin letters) needs specific
+words and word pieces, not a whole script.  --corpus adds the tokens of a text corpus in that language: the most
+frequent ones that together cover --coverage (default 0.99) of the corpus's token occurrences (#597):
+
+    python tools/draft_vocab.py --gguf <model>-00001-of-0000N.gguf --base data/draft_vocab_en.bin \
+        --corpus fr1.txt fr2.txt ... --coverage 0.99 --out data/draft_vocab_fr.bin
+
+It prints the share of the corpus's occurrences outside the base and outside the new subset.  --add and --corpus
+can be combined.
 """
 from __future__ import annotations
 
 import argparse
 import sys
 from array import array
+from collections import Counter
 from pathlib import Path
 
 sys.path.insert(0, str(Path(__file__).resolve().parent))
@@ -50,11 +61,41 @@ def token_text(tok, i: int) -> str | None:
         return None
 
 
+def corpus_counts(tok, paths) -> Counter:
+    """Token occurrences over the corpus files (UTF-8 text, tokenized as plain text: no special tokens)."""
+    counts: Counter = Counter()
+    for path in paths:
+        counts.update(tok.encode(Path(path).read_text(encoding="utf-8")))
+    return counts
+
+
+def covering(counts: Counter, coverage: float) -> list[int]:
+    """The most frequent ids (ties: the lower id first) that together cover `coverage` of the occurrences."""
+    total = sum(counts.values())
+    picked, covered = [], 0
+    for i, c in sorted(counts.items(), key=lambda kv: (-kv[1], kv[0])):
+        if covered >= coverage * total:
+            break
+        picked.append(i)
+        covered += c
+    return picked
+
+
+def outside(counts: Counter, ids) -> float:
+    """The share of the occurrences whose token is not in `ids`."""
+    total = sum(counts.values())
+    have = set(ids)
+    return sum(c for i, c in counts.items() if i not in have) / total if total else 0.0
+
+
 def main() -> int:
     ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
     ap.add_argument("--gguf", required=True, help="the model's (first) GGUF file: its vocabulary")
     ap.add_argument("--base", help="the subset to extend (int32 ids)")
     ap.add_argument("--add", default="", help="comma list: han, kana, hangul, cjk_punct, cjk, cyrillic")
+    ap.add_argument("--corpus", nargs="+", default=[], help="UTF-8 text files in the language to add (#597)")
+    ap.add_argument("--coverage", type=float, default=0.99,
+                    help="with --corpus: the share of its token occurrences the added tokens cover (default 0.99)")
     ap.add_argument("--out", help="where to write the new subset")
     ap.add_argument("--stats", action="store_true", help="print what the base (and the new) subset holds")
     a = ap.parse_args()
@@ -72,8 +113,12 @@ def main() -> int:
         if any(nm not in SCRIPTS for nm in names):
             sys.exit(f"--add {x}: expected one of {', '.join([*SCRIPTS, *GROUPS])}")
         want.update(names)
+    if not 0 < a.coverage <= 1:
+        sys.exit("--coverage: expected a share in (0, 1]")
+    counts = corpus_counts(tok, a.corpus) if a.corpus else Counter()
+    from_corpus = set(covering(counts, a.coverage)) if counts else set()
     have = set(base)
-    added = [i for i in range(n) if i not in have and kinds[i] & want]
+    added = [i for i in range(n) if i not in have and (kinds[i] & want or i in from_corpus)]
     ids = base + added
 
     def stats(label, sel):
@@ -86,6 +131,11 @@ def main() -> int:
             stats("base", base)
         if added:
             stats("new", ids)
+    if counts:
+        print(f"corpus: {sum(counts.values())} token occurrences, {len(counts)} distinct; "
+              f"{len(from_corpus)} tokens cover {a.coverage:g}, {len(from_corpus - have)} of them not in the base")
+        print(f"corpus occurrences outside the base: {outside(counts, base):.2%}, outside the new subset: "
+              f"{outside(counts, ids):.2%}")
     if a.out:
         Path(a.out).write_bytes(array("i", ids).tobytes())
         print(f"wrote {a.out}: {len(ids)} ids ({len(added)} added)")

tools/test_draft_vocab.py

@@ -0,0 +1,47 @@
+"""Tests for tools/draft_vocab.py's corpus coverage mode (#597): the most frequent tokens of a corpus that together cover
+a given share of its occurrences.  No model, no GPU.
+
+    python -m unittest tools.test_draft_vocab
+"""
+from __future__ import annotations
+
+import sys
+import tempfile
+import unittest
+from collections import Counter
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+import draft_vocab as DV  # noqa: E402
+
+
+class Coverage(unittest.TestCase):
+    def test_most_frequent_first_until_the_share_is_covered(self):
+        c = Counter({7: 50, 3: 30, 9: 15, 4: 5})              # 100 occurrences
+        self.assertEqual(DV.covering(c, 0.5), [7])
+        self.assertEqual(DV.covering(c, 0.8), [7, 3])
+        self.assertEqual(DV.covering(c, 0.81), [7, 3, 9])
+        self.assertEqual(DV.covering(c, 1.0), [7, 3, 9, 4])
+
+    def test_ties_take_the_lower_id_first(self):
+        self.assertEqual(DV.covering(Counter({5: 10, 2: 10, 8: 10}), 0.5), [2, 5])
+
+    def test_share_outside_a_subset(self):
+        c = Counter({1: 60, 2: 30, 3: 10})
+        self.assertAlmostEqual(DV.outside(c, [1]), 0.4)
+        self.assertAlmostEqual(DV.outside(c, [1, 2, 3]), 0.0)
+        self.assertEqual(DV.outside(Counter(), [1]), 0.0)
+
+    def test_corpus_counts_reads_every_file(self):
+        class Tok:                                            # one id per character
+            def encode(self, text):
+                return [ord(ch) for ch in text]
+        with tempfile.TemporaryDirectory() as d:
+            a, b = Path(d) / "a.txt", Path(d) / "b.txt"
+            a.write_text("été", encoding="utf-8")
+            b.write_text("eté", encoding="utf-8")
+            self.assertEqual(DV.corpus_counts(Tok(), [a, b]), Counter({ord("é"): 3, ord("t"): 2, ord("e"): 1}))
+
+
+if __name__ == "__main__":
+    unittest.main()

tools/test_setup_draft_vocab.py

@@ -1,5 +1,5 @@
 """Tests for setup.py's draft-vocabulary choice (#287): a setup run again without --draft-vocab keeps the subset the
-model's config chose before (cyrillic, en), and refresh_draft_vocab copies the chosen shipped subset.  Pure file work
+model's config chose before (cyrillic, fr, en), and refresh_draft_vocab copies the chosen shipped subset.  Pure file work
 in a temporary folder - no GPU, no downloads, no prompts.
 
     python -m unittest tools.test_setup_draft_vocab
@@ -49,6 +49,37 @@ class Refresh(unittest.TestCase):
             self.assertEqual(hashlib.sha256((rt / "draft_vocab.bin").read_bytes()).hexdigest(), want)
 
 
+    def test_fr_replaces_a_shipped_subset_and_back(self):
+        # #597: the French subset is a shipped one too - a later --draft-vocab cjk (or en) replaces it again, while a
+        # subset made by hand stays
+        data = ROOT / "data"
+        if not (data / setup.DRAFT_VOCABS["fr"]).exists():
+            self.skipTest("data/draft_vocab_fr.bin is not in this checkout")
+        sha = lambda p: hashlib.sha256(p.read_bytes()).hexdigest()   # noqa: E731
+        with tempfile.TemporaryDirectory() as d:
+            rt = Path(d)
+            (rt / "draft_vocab.bin").write_bytes((data / setup.DRAFT_VOCABS["en"]).read_bytes())
+            setup.refresh_draft_vocab(rt, "fr")
+            self.assertEqual(sha(rt / "draft_vocab.bin"), sha(data / setup.DRAFT_VOCABS["fr"]))
+            setup.refresh_draft_vocab(rt, "cjk")
+            self.assertEqual(sha(rt / "draft_vocab.bin"), sha(data / setup.DRAFT_VOCABS["cjk"]))
+            (rt / "draft_vocab.bin").write_bytes(bytes([1, 0, 0, 0]))   # made by hand: kept
+            setup.refresh_draft_vocab(rt, "fr")
+            self.assertEqual((rt / "draft_vocab.bin").read_bytes(), bytes([1, 0, 0, 0]))
+
+    def test_fr_holds_the_en_subset_in_its_order(self):
+        # built with tools/draft_vocab.py --base data/draft_vocab_en.bin --corpus ...: the base first, unchanged
+        from array import array
+        data = ROOT / "data"
+        if not (data / setup.DRAFT_VOCABS["fr"]).exists():
+            self.skipTest("data/draft_vocab_fr.bin is not in this checkout")
+        en = list(array("i", (data / setup.DRAFT_VOCABS["en"]).read_bytes()))
+        fr = list(array("i", (data / setup.DRAFT_VOCABS["fr"]).read_bytes()))
+        self.assertEqual(fr[:len(en)], en)
+        self.assertEqual(len(set(fr)), len(fr))
+        self.assertEqual(fr[len(en):], sorted(fr[len(en):]))
+
+
 class SmallCardNote(unittest.TestCase):
     """#474: a card under 14 GB is told about a smaller draft subset - a note only, and not when one was chosen."""
 
@@ -57,7 +88,8 @@ class SmallCardNote(unittest.TestCase):
         self.assertTrue(note)
         self.assertIn("--draft-vocab en", " ".join(note))
         self.assertIn("the draft head does not fit", " ".join(note))
-        for vram, chosen in ((16.0, None), (24.0, None), (12.0, "en"), (12.0, "cyrillic"), (12.0, "cjk"), (0.0, None)):
+        for vram, chosen in ((16.0, None), (24.0, None), (12.0, "en"), (12.0, "cyrillic"), (12.0, "fr"), (12.0, "cjk"),
+                             (0.0, None)):
             self.assertEqual(setup.draft_vocab_note(vram, chosen), [], (vram, chosen))
 
     def test_sizes_follow_the_shipped_subsets(self):
@@ -68,7 +100,7 @@ class SmallCardNote(unittest.TestCase):
             if not p.exists():
                 self.skipTest(f"data/{name} is not in this checkout")
             sizes[choice] = p.stat().st_size // 4
-        for choice in ("en", "cyrillic"):
+        for choice in ("en", "cyrillic", "fr"):
             want = setup.DRAFT_VOCAB_MIB["cjk"] * sizes[choice] / sizes["cjk"]
             self.assertAlmostEqual(setup.DRAFT_VOCAB_MIB[choice], want, delta=3)