esp32kb 0.1.0__tar.gz

esp32kb-0.1.0/PKG-INFO

@@ -0,0 +1,123 @@
+Metadata-Version: 2.3
+Name: esp32kb
+Version: 0.1.0
+Summary: Build an on-device knowledge base (text chunks + INT8 MiniLM embeddings) for the ESP32-P4 vector-search demo.
+Keywords: esp32,esp32-p4,embeddings,vector-search,tinyml,onnx,minilm
+Author: Simone Salerno
+Author-email: Simone Salerno <info@salernosimone.com>
+License: MIT
+Classifier: Development Status :: 3 - Alpha
+Classifier: Environment :: Console
+Classifier: Intended Audience :: Developers
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3 :: Only
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Software Development :: Embedded Systems
+Requires-Dist: numpy
+Requires-Dist: onnxruntime
+Requires-Dist: transformers
+Requires-Dist: typer
+Requires-Python: >=3.10
+Project-URL: Homepage, https://github.com/salernosimone/python-esp32kb
+Project-URL: Repository, https://github.com/salernosimone/python-esp32kb
+Project-URL: Issues, https://github.com/salernosimone/python-esp32kb/issues
+Description-Content-Type: text/markdown
+
+# esp32kb
+
+Build an on-device knowledge base for the KB ESP32-P4 demo — a MiniLM-L3
+embedding search that runs entirely on the board, with no PC attached at
+runtime.
+
+Point the tool at a folder of markdown/text documents and it emits everything
+the firmware flashes: sentence-aware chunks, INT8 embeddings, the packed
+`kb.img` container (vectors + text + file index + WordPiece vocabulary) and
+the tiny generated headers the sketch compiles against.
+
+## Install
+
+```sh
+uvx esp32kb --help            # or: pip install esp32kb
+# from a checkout:
+uv run --project esp32kb esp32kb --help
+```
+
+## Usage
+
+```sh
+# 1. documents -> kb.img + generated headers
+esp32kb embed --root path/to/documents --sketch-dir path/to/Arduino/project
+
+# 2. print the arduino-cli command that flashes app + model + kb.img
+esp32kb flash --sketch-dir path/to/Arduino/project   # then answer "y" to run it
+```
+
+Everything defaults relative to `--sketch-dir` (the current folder by default),
+so the demo ships with its MiniLM-L3 model and one flag is enough.
+
+### `esp32kb embed`
+
+| flag | default | meaning |
+|------|---------|---------|
+| `--root` / `-r` | required | folder with the documents to index (recursively) |
+| `--output-folder` / `-o` | `<sketch-dir>/model/kb` | where the `.bin` files and `kb.img` are written |
+| `--sketch-dir` / `-s` | current folder | holds `KB.ino`/`VectorSearch.h`; `kb_config.h` and `kb_queries.h` go here |
+| `--onnx` | `<sketch-dir>/model/maskseq.onnx` | float32 MiniLM-L3 the PC embeds with |
+| `--queries` / `-q` | sampled from the corpus | text file, one demo query per line |
+
+### `esp32kb flash`
+
+| flag | default | meaning |
+|------|---------|---------|
+| `--sketch-dir` / `-s` | current folder | the sketch to compile |
+| `--model` | `<sketch-dir>/model/minilm_l3.espdl` | INT8 model to flash |
+| `--kb` | `<sketch-dir>/model/kb/kb.img` | knowledge base to flash |
+| `--port` / `-p` | `$PORT` or `/dev/ttyUSB0` | serial port |
+| `--fqbn` | P4 board spec | arduino-cli board spec |
+| `--library` | `$ARDUINO_LIBS` | extra Arduino library folder |
+| `--yes` / `-y` | off | run without asking |
+| `--print-only` | off | print the command and exit |
+| `--port` | `$PORT` or `/dev/ttyUSB0` | serial port for the printed flash command |
+| `--fqbn` | esp32p4 FQBN | board spec for the printed flash command |
+
+The sequence length and embedding width are read straight from the ONNX model,
+so the chunks, the INT8 vectors **and** the generated `kb_config.h` the
+firmware compiles against can never disagree with the model. No flags to keep
+in sync.
+
+## Pipeline
+
+```
+documents
+  -> chunks (<= seq-2 WordPiece tokens each)
+  -> MiniLM-L3 embeddings (masked mean pool, L2 normalized)
+  -> symmetric INT8 quantization, one global scale
+  -> vectors.bin / meta.bin / chunks.bin / files.bin / vocab.txt
+  -> kb.img        single little-endian container with a section table
+  -> kb_config.h   KB_SEQ / KB_DIM / KB_TOPK + container magics (C++ constants)
+  -> kb_queries.h  KB_EXAMPLES[] example query texts for the sketch
+```
+
+The vocabulary travels inside `kb.img`, so the device tokenizes queries itself.
+
+## Additive attention mask
+
+If the ONNX model exposes a second `attention_mask` input (the variant built
+by `model/export_onnx_mask.py`), `esp32kb` feeds it the correct per-query
+padding bias (0 on real tokens, negative on `[PAD]`). This preserves retrieval
+quality at longer sequence lengths, where the frozen all-ones mask used by the
+plain export otherwise collapses query embeddings. With the stock single-input
+ONNX the behavior is unchanged.
+
+## Outputs
+
+| file | contents |
+|------|----------|
+| `vectors.bin` | 32 B header + N × dim INT8 (16-byte aligned rows) |
+| `meta.bin` | 16 B header + N × 16 B records (file id, text offset/length) |
+| `chunks.bin` | 16 B header + concatenated UTF-8 chunk text |
+| `files.bin` | 16 B header + N × padded filename records |
+| `vocab.txt` | WordPiece vocabulary, one token per line |
+| `kb.img` | container: magic/version/section table + the payloads above |
+| `kb_config.h` | `KB_SEQ`, `KB_DIM`, `KB_TOPK` and the container magics |
+| `kb_queries.h` | `KB_EXAMPLES[]` example query texts for the sketch |

esp32kb-0.1.0/README.md

@@ -0,0 +1,98 @@
+# esp32kb
+
+Build an on-device knowledge base for the KB ESP32-P4 demo — a MiniLM-L3
+embedding search that runs entirely on the board, with no PC attached at
+runtime.
+
+Point the tool at a folder of markdown/text documents and it emits everything
+the firmware flashes: sentence-aware chunks, INT8 embeddings, the packed
+`kb.img` container (vectors + text + file index + WordPiece vocabulary) and
+the tiny generated headers the sketch compiles against.
+
+## Install
+
+```sh
+uvx esp32kb --help            # or: pip install esp32kb
+# from a checkout:
+uv run --project esp32kb esp32kb --help
+```
+
+## Usage
+
+```sh
+# 1. documents -> kb.img + generated headers
+esp32kb embed --root path/to/documents --sketch-dir path/to/Arduino/project
+
+# 2. print the arduino-cli command that flashes app + model + kb.img
+esp32kb flash --sketch-dir path/to/Arduino/project   # then answer "y" to run it
+```
+
+Everything defaults relative to `--sketch-dir` (the current folder by default),
+so the demo ships with its MiniLM-L3 model and one flag is enough.
+
+### `esp32kb embed`
+
+| flag | default | meaning |
+|------|---------|---------|
+| `--root` / `-r` | required | folder with the documents to index (recursively) |
+| `--output-folder` / `-o` | `<sketch-dir>/model/kb` | where the `.bin` files and `kb.img` are written |
+| `--sketch-dir` / `-s` | current folder | holds `KB.ino`/`VectorSearch.h`; `kb_config.h` and `kb_queries.h` go here |
+| `--onnx` | `<sketch-dir>/model/maskseq.onnx` | float32 MiniLM-L3 the PC embeds with |
+| `--queries` / `-q` | sampled from the corpus | text file, one demo query per line |
+
+### `esp32kb flash`
+
+| flag | default | meaning |
+|------|---------|---------|
+| `--sketch-dir` / `-s` | current folder | the sketch to compile |
+| `--model` | `<sketch-dir>/model/minilm_l3.espdl` | INT8 model to flash |
+| `--kb` | `<sketch-dir>/model/kb/kb.img` | knowledge base to flash |
+| `--port` / `-p` | `$PORT` or `/dev/ttyUSB0` | serial port |
+| `--fqbn` | P4 board spec | arduino-cli board spec |
+| `--library` | `$ARDUINO_LIBS` | extra Arduino library folder |
+| `--yes` / `-y` | off | run without asking |
+| `--print-only` | off | print the command and exit |
+| `--port` | `$PORT` or `/dev/ttyUSB0` | serial port for the printed flash command |
+| `--fqbn` | esp32p4 FQBN | board spec for the printed flash command |
+
+The sequence length and embedding width are read straight from the ONNX model,
+so the chunks, the INT8 vectors **and** the generated `kb_config.h` the
+firmware compiles against can never disagree with the model. No flags to keep
+in sync.
+
+## Pipeline
+
+```
+documents
+  -> chunks (<= seq-2 WordPiece tokens each)
+  -> MiniLM-L3 embeddings (masked mean pool, L2 normalized)
+  -> symmetric INT8 quantization, one global scale
+  -> vectors.bin / meta.bin / chunks.bin / files.bin / vocab.txt
+  -> kb.img        single little-endian container with a section table
+  -> kb_config.h   KB_SEQ / KB_DIM / KB_TOPK + container magics (C++ constants)
+  -> kb_queries.h  KB_EXAMPLES[] example query texts for the sketch
+```
+
+The vocabulary travels inside `kb.img`, so the device tokenizes queries itself.
+
+## Additive attention mask
+
+If the ONNX model exposes a second `attention_mask` input (the variant built
+by `model/export_onnx_mask.py`), `esp32kb` feeds it the correct per-query
+padding bias (0 on real tokens, negative on `[PAD]`). This preserves retrieval
+quality at longer sequence lengths, where the frozen all-ones mask used by the
+plain export otherwise collapses query embeddings. With the stock single-input
+ONNX the behavior is unchanged.
+
+## Outputs
+
+| file | contents |
+|------|----------|
+| `vectors.bin` | 32 B header + N × dim INT8 (16-byte aligned rows) |
+| `meta.bin` | 16 B header + N × 16 B records (file id, text offset/length) |
+| `chunks.bin` | 16 B header + concatenated UTF-8 chunk text |
+| `files.bin` | 16 B header + N × padded filename records |
+| `vocab.txt` | WordPiece vocabulary, one token per line |
+| `kb.img` | container: magic/version/section table + the payloads above |
+| `kb_config.h` | `KB_SEQ`, `KB_DIM`, `KB_TOPK` and the container magics |
+| `kb_queries.h` | `KB_EXAMPLES[]` example query texts for the sketch |

esp32kb-0.1.0/pyproject.toml

@@ -0,0 +1,49 @@
+[project]
+name = "esp32kb"
+version = "0.1.0"
+description = "Build an on-device knowledge base (text chunks + INT8 MiniLM embeddings) for the ESP32-P4 vector-search demo."
+readme = "README.md"
+requires-python = ">=3.10"
+keywords = [
+    "esp32",
+    "esp32-p4",
+    "embeddings",
+    "vector-search",
+    "tinyml",
+    "onnx",
+    "minilm",
+]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Environment :: Console",
+    "Intended Audience :: Developers",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3 :: Only",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    "Topic :: Software Development :: Embedded Systems",
+]
+dependencies = [
+    "numpy",
+    "onnxruntime",
+    "transformers",
+    "typer",
+]
+
+[project.license]
+text = "MIT"
+
+[[project.authors]]
+name = "Simone Salerno"
+email = "info@salernosimone.com"
+
+[project.scripts]
+esp32kb = "esp32kb.cli:main"
+
+[project.urls]
+Homepage = "https://github.com/salernosimone/python-esp32kb"
+Repository = "https://github.com/salernosimone/python-esp32kb"
+Issues = "https://github.com/salernosimone/python-esp32kb/issues"
+
+[build-system]
+requires = ["uv_build>=0.12.17,<0.13.0"]
+build-backend = "uv_build"

esp32kb-0.1.0/pyproject.toml.orig

@@ -0,0 +1,40 @@
+[project]
+name = "esp32kb"
+version = "0.1.0"
+description = "Build an on-device knowledge base (text chunks + INT8 MiniLM embeddings) for the ESP32-P4 vector-search demo."
+readme = "README.md"
+requires-python = ">=3.10"
+license = { text = "MIT" }
+authors = [
+    { name = "Simone Salerno", email = "info@salernosimone.com" },
+]
+keywords = ["esp32", "esp32-p4", "embeddings", "vector-search", "tinyml", "onnx", "minilm"]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Environment :: Console",
+    "Intended Audience :: Developers",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3 :: Only",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    "Topic :: Software Development :: Embedded Systems",
+]
+dependencies = [
+    "numpy",
+    # CPU wheel; onnxruntime-gpu is a separate package and is not used.
+    "onnxruntime",
+    # Only AutoTokenizer is used, so PyTorch is never pulled in.
+    "transformers",
+    "typer",
+]
+
+[project.scripts]
+esp32kb = "esp32kb.cli:main"
+
+[project.urls]
+Homepage = "https://github.com/salernosimone/python-esp32kb"
+Repository = "https://github.com/salernosimone/python-esp32kb"
+Issues = "https://github.com/salernosimone/python-esp32kb/issues"
+
+[build-system]
+requires = ["uv_build>=0.12.17,<0.13.0"]
+build-backend = "uv_build"

esp32kb-0.1.0/src/esp32kb/__init__.py

@@ -0,0 +1,14 @@
+"""esp32kb: build an on-device knowledge base for the KB ESP32-P4 demo.
+
+The command line interface has two subcommands:
+
+    esp32kb embed --root DOCS [--output-folder DIR] [--sketch-dir DIR]
+    esp32kb flash [--sketch-dir DIR] [-y]
+
+`embed` turns a folder of documents into a packed `kb.img` plus the generated
+`kb_config.h` / `kb_queries.h` the firmware compiles against. `flash` prints
+the arduino-cli command that flashes the sketch, the model and `kb.img`, and
+offers to run it.
+"""
+
+__version__ = "0.1.0"

esp32kb-0.1.0/src/esp32kb/__main__.py

@@ -0,0 +1,4 @@
+from .cli import main
+
+if __name__ == "__main__":
+    main()

esp32kb-0.1.0/src/esp32kb/build.py

@@ -0,0 +1,476 @@
+"""Embed documents and pack them into the firmware's knowledge base image.
+
+The whole pipeline runs on the PC; the ESP32 only reads the result:
+
+    documents
+      -> chunks (<= seq-2 WordPiece tokens each)
+      -> MiniLM-L3 embeddings (masked mean pool, L2 normalized)
+      -> symmetric INT8 quantization with one global scale
+      -> vectors.bin / meta.bin / chunks.bin / files.bin / vocab.txt
+      -> kb.img        all payloads in one little-endian container
+      -> kb_config.h   build constants for VectorSearch.h (seq, dim, magics)
+      -> kb_queries.h  example query texts for the sketch
+
+The sequence length and embedding width are read straight from the ONNX model,
+so `kb.img` and `kb_config.h` can never disagree with the model the firmware
+loads. The vocabulary travels inside `kb.img`, so the device tokenizes queries
+itself and no PC is needed at runtime.
+"""
+
+from __future__ import annotations
+
+import glob
+import os
+import re
+import struct
+from dataclasses import dataclass
+from pathlib import Path
+
+import numpy as np
+import onnxruntime as ort
+from transformers import AutoTokenizer
+
+MODEL_ID = "sentence-transformers/paraphrase-MiniLM-L3-v2"
+FALLBACK_SEQ = 64          # used only if the ONNX has dynamic axes
+FALLBACK_DIM = 384
+KB_TOPK = 5                # hits the firmware keeps
+DOCUMENT_SUFFIXES = {".md", ".markdown", ".txt"}
+
+# on-disk container constants; kb_config.h is generated from these so packer
+# and firmware can never disagree.
+KB_MAGIC = 0x3130424B
+VEC_MAGIC = 0x3156424B
+META_MAGIC = 0x314D424B
+CHUNK_MAGIC = 0x3143424B
+FILE_MAGIC = 0x3146424B
+VOCAB_MAGIC = 0x3142424B
+SEC_VECTORS, SEC_META, SEC_CHUNKS, SEC_FILES, SEC_VOCAB = 1, 2, 3, 4, 5
+
+# flash layout: must match partitions.csv
+MODEL_OFFSET = 0x510000
+KB_OFFSET = 0x1910000
+KB_PARTITION_SIZE = 0x600000
+DEFAULT_FQBN = (
+    "esp32:esp32:esp32p4:ChipVariant=postv3,FlashSize=32M,"
+    "PSRAM=enabled,PartitionScheme=custom"
+)
+
+
+# ---------------------------------------------------------------- data model
+
+@dataclass
+class Document:
+    filename: str
+    text: str
+
+
+@dataclass
+class Chunk:
+    text: str
+    filename: str
+    index: int
+    embedding: np.ndarray | None = None
+    embedding_int8: np.ndarray | None = None
+
+
+@dataclass
+class BuildResult:
+    kb_image: Path
+    config_header: Path
+    queries_header: Path
+    chunks: int
+    files: int
+    seq: int
+    dim: int
+    scale: float
+
+
+def _fixed_dim(shape, axis: int, fallback: int) -> int:
+    """A concrete value from an ONNX tensor shape, or `fallback` if dynamic."""
+    try:
+        dim = shape[axis]
+        return int(dim) if dim is not None and int(dim) > 0 else fallback
+    except (IndexError, TypeError, ValueError):
+        return fallback
+
+
+def stride_of(dim: int) -> int:
+    """16-byte aligned vector row length (esp-nn wants aligned rows)."""
+    return (dim + 15) // 16 * 16
+
+
+# ---------------------------------------------------------------- embedding
+
+class Embedder:
+    """MiniLM-L3: text -> L2-normalized vector, mirroring the firmware.
+
+    The ONNX model returns the full token sequence [1, seq, dim]; pooling is
+    done here (masked mean) exactly like VectorSearch.h does on the device.
+    """
+
+    # Additive padding bias fed to the "attention_mask" input of the
+    # additive-mask ONNX variant (see model/export_onnx_mask.py). Large enough
+    # that the Softmax over [PAD] keys saturates to ~0.
+    MASK_NEG = -30.0
+
+    def __init__(self, onnx_path: str | os.PathLike):
+        if not os.path.isfile(onnx_path):
+            raise FileNotFoundError(onnx_path)
+        self.tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+        self.session = ort.InferenceSession(
+            str(onnx_path), providers=["CPUExecutionProvider"]
+        )
+        self.input_names = {i.name for i in self.session.get_inputs()}
+        # The ONNX is the single source of truth for the geometry.
+        self.seq = _fixed_dim(self.session.get_inputs()[0].shape, 1, FALLBACK_SEQ)
+        self.dim = _fixed_dim(self.session.get_outputs()[0].shape, -1, FALLBACK_DIM)
+
+    def tokenize(self, text: str) -> tuple[np.ndarray, int]:
+        """Return (padded int64 token ids, number of real tokens)."""
+        encoded = self.tokenizer(
+            text,
+            padding="max_length",
+            truncation=True,
+            max_length=self.seq,
+            return_tensors="np",
+        )
+        ids = encoded["input_ids"][0].astype(np.int64)
+        length = int(encoded["attention_mask"][0].sum())
+        return ids, length
+
+    def embed(self, text: str) -> np.ndarray:
+        """Return the normalized embedding of one piece of text."""
+        ids, length = self.tokenize(text)
+        feeds = {"input_ids": ids[None]}
+        if "attention_mask" in self.input_names:
+            bias = np.zeros((1, 1, 1, self.seq), dtype=np.float32)
+            bias[0, 0, 0, length:] = self.MASK_NEG
+            feeds["attention_mask"] = bias
+        hidden = self.session.run(None, feeds)[0][0]  # seq x dim
+        vector = hidden[:length].mean(axis=0)
+        return vector / (np.linalg.norm(vector) + 1e-9)
+
+    def vocabulary(self) -> bytes:
+        """The WordPiece vocab as vocab.txt: one token per line, id = line."""
+        vocab = self.tokenizer.get_vocab()
+        ordered = sorted(vocab.items(), key=lambda item: item[1])
+        return ("\n".join(token for token, _ in ordered) + "\n").encode("utf-8")
+
+
+# ---------------------------------------------------------------- chunking
+
+def count_tokens(tokenizer, text: str) -> int:
+    return len(tokenizer(text, add_special_tokens=False)["input_ids"])
+
+
+def split_markdown_into_fragments(text: str, tokenizer, max_tokens: int) -> list[str]:
+    """Cut text into sentence/line fragments no longer than max_tokens."""
+    fragments = []
+    for raw in re.split(r"(?<=[.!?])\s+|\n+", text):
+        fragment = raw.strip()
+        if not fragment:
+            continue
+        if count_tokens(tokenizer, fragment) <= max_tokens:
+            fragments.append(fragment)
+        else:
+            fragments.extend(split_oversize_fragment(fragment, tokenizer, max_tokens))
+    return fragments
+
+
+def split_oversize_fragment(fragment: str, tokenizer, max_tokens: int) -> list[str]:
+    """An oversize fragment is packed word by word to the token budget."""
+    pieces, current = [], []
+    for word in fragment.split():
+        if current and count_tokens(tokenizer, " ".join(current + [word])) > max_tokens:
+            pieces.append(" ".join(current))
+            current = [word]
+        else:
+            current.append(word)
+    if current:
+        pieces.append(" ".join(current))
+    return pieces
+
+
+def pack_fragments_into_chunks(fragments: list[str], tokenizer, max_tokens: int) -> list[str]:
+    """Greedily pack fragments into chunks that fit the model sequence."""
+    chunks, current, length = [], [], 0
+    for fragment in fragments:
+        tokens = count_tokens(tokenizer, fragment)
+        if tokens == 0:
+            continue
+        if current and length + tokens > max_tokens:
+            chunks.append(" ".join(current))
+            current, length = [], 0
+        current.append(fragment)
+        length += tokens
+    if current:
+        chunks.append(" ".join(current))
+    return chunks
+
+
+def load_documents_from_folder(root: str | os.PathLike) -> list[Document]:
+    """Read every markdown/text file under root (recursively)."""
+    documents = []
+    for path in sorted(glob.glob(os.path.join(str(root), "**", "*"), recursive=True)):
+        if not os.path.isfile(path):
+            continue
+        if os.path.splitext(path)[1].lower() not in DOCUMENT_SUFFIXES:
+            continue
+        text = Path(path).read_text(encoding="utf-8", errors="replace")
+        documents.append(Document(filename=os.path.relpath(path, root), text=text))
+    return documents
+
+
+def split_into_chunks(documents: list[Document], tokenizer, max_tokens: int) -> list[Chunk]:
+    """Chunk every document; chunks never cross a file boundary."""
+    chunks = []
+    for document in documents:
+        fragments = split_markdown_into_fragments(document.text, tokenizer, max_tokens)
+        packed = pack_fragments_into_chunks(fragments, tokenizer, max_tokens)
+        for index, text in enumerate(packed):
+            chunks.append(Chunk(text=text, filename=document.filename, index=index))
+    return chunks
+
+
+def embed_chunks(chunks: list[Chunk], embedder: Embedder, progress=None) -> None:
+    for i, chunk in enumerate(chunks):
+        chunk.embedding = embedder.embed(chunk.text)
+        if progress and (i + 1) % 100 == 0:
+            progress(i + 1, len(chunks))
+
+
+def quantize_chunks(chunks: list[Chunk]) -> float:
+    """Symmetric INT8 with one global scale shared by corpus and queries."""
+    max_abs = max(float(np.max(np.abs(chunk.embedding))) for chunk in chunks)
+    scale = 127.0 / max_abs
+    for chunk in chunks:
+        chunk.embedding_int8 = np.clip(
+            np.rint(chunk.embedding * scale), -127, 127
+        ).astype(np.int8)
+    return scale
+
+
+# ---------------------------------------------------------------- binary layout
+
+def pack_vector_section(chunks: list[Chunk], scale: float, dim: int) -> bytes:
+    """32 B header + N x stride INT8 (16-byte aligned rows)."""
+    stride = stride_of(dim)
+    header = struct.pack(
+        "<5If2I", VEC_MAGIC, 1, dim, len(chunks), stride, float(scale), 0, 0
+    )
+    rows = b"".join(chunk.embedding_int8.tobytes() for chunk in chunks)
+    return header + rows
+
+
+def pack_chunks_and_meta(chunks: list[Chunk]) -> tuple[bytes, bytes]:
+    """chunks.bin (header + UTF-8 text) and meta.bin (header + 16 B records)."""
+    file_ids = {name: i for i, name in enumerate(dict.fromkeys(c.filename for c in chunks))}
+
+    text = bytearray()
+    records = bytearray()
+    for chunk in chunks:
+        data = chunk.text.encode("utf-8")
+        records += struct.pack("<2H3I", file_ids[chunk.filename], 0, len(text), len(data), 0)
+        text += data
+
+    chunks_bin = struct.pack("<4I", CHUNK_MAGIC, 1, len(text), 0) + bytes(text)
+    meta_bin = struct.pack("<4I", META_MAGIC, 1, len(chunks), 0) + bytes(records)
+    return chunks_bin, meta_bin
+
+
+def pack_files_section(chunks: list[Chunk]) -> bytes:
+    """16 B header + N zero-padded file-name records."""
+    names = list(dict.fromkeys(chunk.filename for chunk in chunks))
+    name_len = max((len(name) + 1 for name in names), default=1)
+    name_len = (name_len + 15) // 16 * 16
+
+    records = bytearray()
+    for name in names:
+        encoded = name.encode("utf-8")[: name_len - 1]
+        records += encoded + b"\x00" * (name_len - len(encoded))
+
+    return struct.pack("<4I", FILE_MAGIC, 1, len(names), name_len) + bytes(records)
+
+
+def pack_vocab_section(vocab: bytes) -> bytes:
+    """16 B header + the raw vocab.txt bytes (one token per line)."""
+    return struct.pack("<4I", VOCAB_MAGIC, 1, len(vocab), 0) + vocab
+
+
+def pack_container(sections: list[tuple[int, bytes]]) -> bytes:
+    """Concatenate (type, payload) sections into one little-endian image."""
+    image = bytearray(struct.pack("<4I", KB_MAGIC, 1, len(sections), 0))
+    image += b"\x00" * (16 * len(sections))
+
+    table = []
+    for section_type, payload in sections:
+        while len(image) % 16:
+            image.append(0)
+        table.append((section_type, len(image), len(payload), 0))
+        image += payload
+
+    for i, entry in enumerate(table):
+        struct.pack_into("<4I", image, 16 + 16 * i, *entry)
+    return bytes(image)
+
+
+# ---------------------------------------------------------------- writers
+
+def store_embeddings(chunks: list[Chunk], scale: float, dim: int, output_folder: Path) -> None:
+    Path(output_folder).mkdir(parents=True, exist_ok=True)
+    Path(output_folder, "vectors.bin").write_bytes(pack_vector_section(chunks, scale, dim))
+
+
+def store_files_index(chunks: list[Chunk], output_folder: Path) -> None:
+    Path(output_folder, "files.bin").write_bytes(pack_files_section(chunks))
+
+
+def store_texts(chunks: list[Chunk], output_folder: Path) -> None:
+    chunks_bin, meta_bin = pack_chunks_and_meta(chunks)
+    Path(output_folder, "chunks.bin").write_bytes(chunks_bin)
+    Path(output_folder, "meta.bin").write_bytes(meta_bin)
+
+
+def store_vocab(vocab: bytes, output_folder: Path) -> None:
+    Path(output_folder, "vocab.txt").write_bytes(vocab)
+
+
+def store_container(
+    chunks: list[Chunk], scale: float, vocab: bytes, dim: int, output_folder: Path
+) -> None:
+    """Repack every payload into kb.img for the 'kb' partition."""
+    chunks_bin, meta_bin = pack_chunks_and_meta(chunks)
+    sections = [
+        (SEC_VECTORS, pack_vector_section(chunks, scale, dim)),
+        (SEC_META, meta_bin),
+        (SEC_CHUNKS, chunks_bin),
+        (SEC_FILES, pack_files_section(chunks)),
+        (SEC_VOCAB, pack_vocab_section(vocab)),
+    ]
+    Path(output_folder, "kb.img").write_bytes(pack_container(sections))
+
+
+def store_config_header(seq: int, dim: int, sketch_dir: Path) -> Path:
+    """Write kb_config.h: the build constants VectorSearch.h compiles against."""
+    path = Path(sketch_dir, "kb_config.h")
+    path.write_text(
+        "#pragma once\n"
+        "// Generated by esp32kb. Do not edit; re-run esp32kb instead.\n"
+        f"#define KB_SEQ {seq}\n"
+        f"#define KB_DIM {dim}\n"
+        f"#define KB_TOPK {KB_TOPK}\n"
+        "\n"
+        f"#define KB_MAGIC 0x{KB_MAGIC:08X}\n"
+        f"#define KB_VEC_MAGIC 0x{VEC_MAGIC:08X}\n"
+        f"#define KB_META_MAGIC 0x{META_MAGIC:08X}\n"
+        f"#define KB_CHUNK_MAGIC 0x{CHUNK_MAGIC:08X}\n"
+        f"#define KB_FILE_MAGIC 0x{FILE_MAGIC:08X}\n"
+        f"#define KB_VOCAB_MAGIC 0x{VOCAB_MAGIC:08X}\n"
+        "\n"
+        f"#define KB_SEC_VECTORS {SEC_VECTORS}\n"
+        f"#define KB_SEC_META {SEC_META}\n"
+        f"#define KB_SEC_CHUNKS {SEC_CHUNKS}\n"
+        f"#define KB_SEC_FILES {SEC_FILES}\n"
+        f"#define KB_SEC_VOCAB {SEC_VOCAB}\n"
+    )
+    return path
+
+
+def cstr(text: str) -> str:
+    """A C string literal (double quoted, escaped) — never Python repr()."""
+    escaped = text.replace("\\", "\\\\").replace('"', '\\"').replace("\n", "\\n")
+    return f'"{escaped}"'
+
+
+def store_queries(queries: list[str], sketch_dir: Path) -> Path:
+    """Write kb_queries.h: example query texts, tokenized on device."""
+    examples = ",".join(cstr(query) for query in queries)
+    size = len(queries) or 1  # C arrays can't have length 0
+    lines = [
+        "#pragma once",
+        f"#define KB_NEXAMPLE {len(queries)}",
+        f"static const char* KB_EXAMPLES[{size}] = {{{examples}}};",
+    ]
+    path = Path(sketch_dir, "kb_queries.h")
+    path.write_text("\n".join(lines) + "\n")
+    return path
+
+
+# ---------------------------------------------------------------- queries
+
+def load_queries_from_file(path: str | os.PathLike) -> list[str]:
+    return [
+        line.strip()
+        for line in Path(path).read_text(encoding="utf-8").splitlines()
+        if line.strip() and not line.lstrip().startswith("#")
+    ]
+
+
+def first_sentence(text: str) -> str:
+    sentence = re.split(r"(?<=[.!?])\s+|\n+", text.strip())[0].strip(" #*->`")
+    return sentence or text.strip()[:120]
+
+
+def default_queries(chunks: list[Chunk], count: int = 5) -> list[str]:
+    """Sample a few chunks and use their opening sentence as demo queries."""
+    if not chunks:
+        return []
+    step = max(1, len(chunks) // count)
+    return [first_sentence(chunk.text) for chunk in chunks[::step][:count]]
+
+
+# ---------------------------------------------------------------- pipeline
+
+def build(
+    root: str | os.PathLike,
+    output_folder: str | os.PathLike,
+    sketch_dir: str | os.PathLike,
+    onnx_path: str | os.PathLike,
+    queries_path: str | os.PathLike | None = None,
+    log=print,
+) -> BuildResult:
+    """Run the whole PC pipeline and return the paths that were written."""
+    output_folder = Path(output_folder)
+    sketch_dir = Path(sketch_dir)
+
+    documents = load_documents_from_folder(root)
+    if not documents:
+        raise ValueError(
+            f"no {', '.join(sorted(DOCUMENT_SUFFIXES))} files found under {root}"
+        )
+
+    embedder = Embedder(onnx_path)
+    chunks = split_into_chunks(documents, embedder.tokenizer, embedder.seq - 2)
+    log(f"indexed {len(documents)} files -> {len(chunks)} chunks "
+        f"(seq={embedder.seq}, dim={embedder.dim})")
+    embed_chunks(chunks, embedder, progress=lambda i, n: log(f"  embedded {i}/{n}"))
+    scale = quantize_chunks(chunks)
+
+    vocab = embedder.vocabulary()
+    store_embeddings(chunks, scale, embedder.dim, output_folder)
+    store_files_index(chunks, output_folder)
+    store_texts(chunks, output_folder)
+    store_vocab(vocab, output_folder)
+    store_container(chunks, scale, vocab, embedder.dim, output_folder)
+
+    config_header = store_config_header(embedder.seq, embedder.dim, sketch_dir)
+    queries = (
+        load_queries_from_file(queries_path) if queries_path else default_queries(chunks)
+    )
+    queries_header = store_queries(queries, sketch_dir)
+
+    kb_image = output_folder / "kb.img"
+    if kb_image.stat().st_size > KB_PARTITION_SIZE:
+        log(f"warning: kb.img is {kb_image.stat().st_size} B, larger than the "
+            f"{KB_PARTITION_SIZE} B 'kb' partition (see partitions.csv)")
+
+    return BuildResult(
+        kb_image=kb_image,
+        config_header=config_header,
+        queries_header=queries_header,
+        chunks=len(chunks),
+        files=len(documents),
+        seq=embedder.seq,
+        dim=embedder.dim,
+        scale=scale,
+    )

esp32kb-0.1.0/src/esp32kb/cli.py

@@ -0,0 +1,144 @@
+"""`esp32kb embed` / `esp32kb flash` — the command line interface."""
+
+from __future__ import annotations
+
+import os
+import subprocess
+from pathlib import Path
+
+import typer
+
+from . import build as _build
+from . import flash as _flash
+
+app = typer.Typer(
+    add_completion=False,
+    no_args_is_help=True,
+    help="Build an on-device knowledge base for the KB ESP32-P4 demo.",
+)
+
+
+def _resolve_sketch(sketch_dir: Path | None) -> Path:
+    return (sketch_dir or Path.cwd()).expanduser().resolve()
+
+
+@app.command()
+def embed(
+    root: Path = typer.Option(
+        ..., "--root", "-r", exists=True, file_okay=False, dir_okay=True,
+        help="folder with the documents to index (recursively)",
+    ),
+    output_folder: Path | None = typer.Option(
+        None, "--output-folder", "-o",
+        help="where the .bin files and kb.img are written "
+             "(default: <sketch-dir>/model/kb)",
+    ),
+    sketch_dir: Path | None = typer.Option(
+        None, "--sketch-dir", "-s",
+        help="folder holding KB.ino/VectorSearch.h; kb_config.h and kb_queries.h "
+             "are written here (default: current folder)",
+    ),
+    onnx: Path | None = typer.Option(
+        None, "--onnx",
+        help="MiniLM-L3 ONNX the PC embeds with "
+             "(default: <sketch-dir>/model/maskseq.onnx)",
+    ),
+    queries: Path | None = typer.Option(
+        None, "--queries", "-q", exists=True, dir_okay=False,
+        help="text file, one demo query per line (default: sampled from the corpus)",
+    ),
+) -> None:
+    """Embed documents and write kb.img + the generated C++ headers."""
+    sketch = _resolve_sketch(sketch_dir)
+    onnx_path = (onnx or sketch / "model" / "maskseq.onnx").expanduser().resolve()
+    out = (output_folder or sketch / "model" / "kb").expanduser().resolve()
+
+    if not onnx_path.is_file():
+        typer.secho(f"error: ONNX model not found: {onnx_path}", fg=typer.colors.RED, err=True)
+        typer.echo("Build it with model/export_onnx_mask.py, or pass --onnx.", err=True)
+        raise typer.Exit(1)
+
+    try:
+        result = _build.build(
+            root=root,
+            output_folder=out,
+            sketch_dir=sketch,
+            onnx_path=onnx_path,
+            queries_path=queries,
+        )
+    except ValueError as error:
+        typer.secho(f"error: {error}", fg=typer.colors.RED, err=True)
+        raise typer.Exit(1)
+
+    typer.echo(f"\nwrote {result.kb_image.parent}/ "
+               f"(vectors.bin, meta.bin, chunks.bin, files.bin, vocab.txt, kb.img)")
+    typer.echo(f"wrote {result.config_header} "
+               f"(seq={result.seq}, dim={result.dim}, topk={_build.KB_TOPK})")
+    typer.echo(f"wrote {result.queries_header} (scale={result.scale:.3f})")
+    typer.echo("\n# flash it with:  esp32kb flash --sketch-dir " + str(sketch))
+
+
+@app.command()
+def flash(
+    sketch_dir: Path | None = typer.Option(
+        None, "--sketch-dir", "-s", help="folder holding KB.ino/VectorSearch.h "
+                                         "(default: current folder)",
+    ),
+    model: Path | None = typer.Option(
+        None, "--model", help="quantized .espdl to flash "
+                              "(default: <sketch-dir>/model/minilm_l3.espdl)",
+    ),
+    kb: Path | None = typer.Option(
+        None, "--kb", help="kb.img to flash (default: <sketch-dir>/model/kb/kb.img)",
+    ),
+    port: str | None = typer.Option(
+        None, "--port", "-p", help="serial port (default: $PORT or /dev/ttyUSB0)",
+    ),
+    fqbn: str | None = typer.Option(
+        None, "--fqbn", help="arduino-cli board spec (default: the P4 board)",
+    ),
+    library: str | None = typer.Option(
+        None, "--library", help="extra Arduino library folder (default: $ARDUINO_LIBS)",
+    ),
+    yes: bool = typer.Option(False, "--yes", "-y", help="run without asking"),
+    print_only: bool = typer.Option(False, "--print-only", help="print the command and exit"),
+) -> None:
+    """Print the command that flashes app + model + kb.img, and offer to run it."""
+    sketch = _resolve_sketch(sketch_dir)
+    model_path = (model or sketch / "model" / "minilm_l3.espdl").expanduser().resolve()
+    kb_path = (kb or sketch / "model" / "kb" / "kb.img").expanduser().resolve()
+    port = port or os.environ.get("PORT", "/dev/ttyUSB0")
+    fqbn = fqbn or os.environ.get("FQBN", _flash.DEFAULT_FQBN)
+    library = library or os.environ.get("ARDUINO_LIBS")
+
+    for label, path in (("model", model_path), ("kb.img", kb_path)):
+        if not path.is_file():
+            typer.secho(f"error: {label} not found: {path}", fg=typer.colors.RED, err=True)
+            typer.echo("Run `esp32kb embed` and build the model first.", err=True)
+            raise typer.Exit(1)
+
+    compile_cmd, upload_cmd = _flash.commands(
+        sketch, model_path, kb_path, port, fqbn, library
+    )
+    typer.echo(_flash.as_shell(compile_cmd, upload_cmd))
+
+    if print_only:
+        return
+    if not (yes or typer.confirm("Run it now?", default=False)):
+        return
+
+    try:
+        _flash.run(compile_cmd, upload_cmd)
+    except FileNotFoundError:
+        typer.secho("error: arduino-cli not found on PATH", fg=typer.colors.RED, err=True)
+        raise typer.Exit(1)
+    except subprocess.CalledProcessError as error:
+        raise typer.Exit(error.returncode or 1)
+
+
+def main() -> None:
+    app()
+
+
+if __name__ == "__main__":
+    main()

esp32kb-0.1.0/src/esp32kb/flash.py

@@ -0,0 +1,56 @@
+"""Build the arduino-cli command that flashes app + model + kb.img.
+
+arduino-cli's upload recipe is hardcoded to the four sketch binaries, so the
+model and the knowledge base are appended to the same esptool write-flash via
+`--upload-property upload.extra_flags`. No Arduino system files are touched.
+"""
+
+from __future__ import annotations
+
+import shlex
+import subprocess
+from pathlib import Path
+
+MODEL_OFFSET = 0x510000    # must match the 'model' partition in partitions.csv
+KB_OFFSET = 0x1910000      # must match the 'kb' partition in partitions.csv
+
+DEFAULT_FQBN = (
+    "esp32:esp32:esp32p4:ChipVariant=postv3,FlashSize=32M,"
+    "PSRAM=enabled,PartitionScheme=custom"
+)
+
+
+def commands(
+    sketch_dir: Path,
+    model_path: Path,
+    kb_image: Path,
+    port: str,
+    fqbn: str,
+    library: str | None = None,
+) -> tuple[list[str], list[str]]:
+    """Return (compile_cmd, upload_cmd) as argv lists (no shell involved)."""
+    compile_cmd = ["arduino-cli", "compile", "-b", fqbn]
+    if library:
+        compile_cmd += ["--library", library]
+    compile_cmd.append(str(sketch_dir))
+
+    extra = f"{hex(MODEL_OFFSET)} {model_path} {hex(KB_OFFSET)} {kb_image}"
+    upload_cmd = [
+        "arduino-cli", "upload",
+        "-p", port,
+        "-b", fqbn,
+        "--upload-property", f"upload.extra_flags={extra}",
+        str(sketch_dir),
+    ]
+    return compile_cmd, upload_cmd
+
+
+def as_shell(compile_cmd: list[str], upload_cmd: list[str]) -> str:
+    """A copy-pasteable one-liner: compile, then upload app + model + kb.img."""
+    return f"{shlex.join(compile_cmd)} && \\\n  {shlex.join(upload_cmd)}"
+
+
+def run(compile_cmd: list[str], upload_cmd: list[str]) -> None:
+    """Run both commands, streaming their output."""
+    subprocess.run(compile_cmd, check=True)
+    subprocess.run(upload_cmd, check=True)