diff --git a/sd_files/BruceLM/models/Chat3.3M/chat3.3M.bin b/sd_files/BruceLM/models/Chat3.3M/chat3.3M.bin new file mode 100644 index 0000000000..5108091b0b Binary files /dev/null and b/sd_files/BruceLM/models/Chat3.3M/chat3.3M.bin differ diff --git a/sd_files/BruceLM/models/Chat3.3M/tok1024.bin b/sd_files/BruceLM/models/Chat3.3M/tok1024.bin new file mode 100644 index 0000000000..006fff71a6 Binary files /dev/null and b/sd_files/BruceLM/models/Chat3.3M/tok1024.bin differ diff --git a/sd_files/BruceLM/models/Stories260K/stories260K.bin b/sd_files/BruceLM/models/Stories260K/stories260K.bin new file mode 100644 index 0000000000..cfebf5a594 Binary files /dev/null and b/sd_files/BruceLM/models/Stories260K/stories260K.bin differ diff --git a/sd_files/BruceLM/models/Stories260K/tok512.bin b/sd_files/BruceLM/models/Stories260K/tok512.bin new file mode 100644 index 0000000000..be175e5ee3 Binary files /dev/null and b/sd_files/BruceLM/models/Stories260K/tok512.bin differ diff --git a/src/core/menu_items/OthersMenu.cpp b/src/core/menu_items/OthersMenu.cpp index 53e78ac7dd..1beb397dc7 100644 --- a/src/core/menu_items/OthersMenu.cpp +++ b/src/core/menu_items/OthersMenu.cpp @@ -6,6 +6,7 @@ #include "modules/bjs_interpreter/interpreter.h" #include "modules/others/clicker.h" #include "modules/others/ibutton.h" +#include "modules/others/BruceLM/bruce_lm.h" #include "modules/others/mic.h" #include "modules/others/qrcode_menu.h" #include "modules/others/tururururu.h" @@ -16,6 +17,7 @@ void OthersMenu::optionsMenu() { options = { {"QRCodes", qrcode_menu }, {"Megalodon", shark_setup }, + {"BruceLM", bruceLM_setup }, #if defined(MIC_SPM1423) || defined(MIC_INMP441) {"Microphone", [this]() { micMenu(); } }, //@deveclipse diff --git a/src/modules/others/BruceLM/bruce_lm.cpp b/src/modules/others/BruceLM/bruce_lm.cpp new file mode 100644 index 0000000000..fb4f90482e --- /dev/null +++ b/src/modules/others/BruceLM/bruce_lm.cpp @@ -0,0 +1,1178 @@ +/** + * @file bruce_lm.cpp + * @brief "Others -> BruceLM" module: on-device LLM. + * + * ---- Made by @Doominator1 on GitHub, Jul 2026 + * ---- Much thanks to @karpathy for the base of this project; llama2.c! + * + * Models + tokenizers live on SD at /BruceLM/models/ (auto-created here, + * matching the same idiom used by ir_read.cpp's /BruceIR and + * wardriving.cpp's /BruceWardriving). Model selection reuses Bruce's + * existing loopSD() file picker; text entry reuses the existing keyboard() + * overlay. Everything is drawn with the shared theme colors/border helpers + * so it matches the rest of the firmware. + * + * Layout (top to bottom, inside the standard Bruce border+title): + * - scrolling chat transcript, word-wrapped, "You:" lines in the primary + * theme color and assistant text in the secondary color + * - a horizontal separator line + * - a one-line "input bar" row showing the current/last prompt + * - a footer hint row ("OK: new prompt Esc: exit") + * + * The whole frame is redrawn from scratch on every change (new token, + * scroll tick, new turn) instead of incrementally patching the screen - + * slightly more pixels pushed, but it removes an entire class of stale-text/ + * overlap bugs that incremental cursor-based drawing is prone to. + * + * Scrolling uses this board's rotary encoder, which the board's InputHandler + * (boards/lilygo-t-embed-cc1101/interface.cpp) maps to NextPress/PrevPress - + * a separate signal from the physical Select/Back buttons (SelPress/EscPress), + * so scrolling and cancelling never conflict. NextPress scrolls down toward + * newer content, PrevPress scrolls up toward older content, matching how + * Next/Prev already mean forward/backward everywhere else in this firmware. + */ +#include "bruce_lm.h" + +#include "core/display.h" +#include "core/mykeyboard.h" +#include "core/sd_functions.h" +#include "llm_engine.h" +#include +#include +#include + +using namespace bruce_llm; + +namespace { + +constexpr const char *kRootDir = "/BruceLM"; +constexpr const char *kModelsDir = "/BruceLM/models"; +constexpr const char *kConfigPath = "/BruceLM/config.json"; + +// Seed presets cycled through with Next/Prev on the Seed row, rather than +// free-form numeric entry - "Random" reseeds from the hardware RNG each +// reply (value 0 is the sentinel llm_engine.cpp treats as "reseed"); the +// rest give byte-for-byte reproducible output for the same prompt. +struct SeedPreset { + const char *label; + uint32_t value; +}; +constexpr SeedPreset kSeedPresets[] = { + {"1", 1 }, + {"2", 2 }, + {"3", 3 }, + {"4", 4 }, + {"5", 5 }, + {"6", 6 }, + {"7", 7 }, + {"8", 8 }, + {"9", 9 }, + {"10", 10 }, + {"42", 42 }, + {"67", 67 }, + {"100", 100 }, + {"1000", 1000}, + {"Random", 0 }, +}; +constexpr int kNumSeedPresets = sizeof(kSeedPresets) / sizeof(kSeedPresets[0]); +constexpr int kDefaultSeedPresetIndex = kNumSeedPresets - 1; // "Random" + +// Every field here is a real, functioning knob on the engine +// (llm_engine.cpp's generate() / GenerationParams) - nothing decorative. +// Defaults are tuned for a model this small (3.3M params) rather than copied +// from typical LLM defaults - see git history/PR for the prompt-comparison +// data behind the specific values. +// +// Wraps prompts as `userTag + prompt + "\n" + botTag` before encoding, so a +// chat-finetuned model (trained on that exact ": ...\n: ..." pair +// format) is cued to answer immediately instead of hallucinating a whole +// fake turn. On by default to match Chat3.3M.bin, the default recommended +// model - plain story models (e.g. stories260K.bin) expect raw text instead. +struct BruceLMSettings { + float temperature = 0.5f; + float topP = 0.6f; + float repetitionPenalty = 1.15f; + int maxTokens = 128; + int seedPresetIndex = kDefaultSeedPresetIndex; + bool chatTemplateEnabled = true; + String userTag = ": "; + String botTag = ":"; + + uint32_t seed() const { return kSeedPresets[seedPresetIndex].value; } +}; + +BruceLMSettings loadSettings(FS &fs) { + BruceLMSettings s; + if (!fs.exists(kConfigPath)) return s; + File f = fs.open(kConfigPath, FILE_READ); + if (!f) return s; + JsonDocument doc; + if (!deserializeJson(doc, f)) { + if (!doc["temperature"].isNull()) s.temperature = doc["temperature"]; + if (!doc["topP"].isNull()) s.topP = doc["topP"]; + if (!doc["repetitionPenalty"].isNull()) s.repetitionPenalty = doc["repetitionPenalty"]; + if (!doc["maxTokens"].isNull()) s.maxTokens = doc["maxTokens"]; + if (!doc["seedPresetIndex"].isNull()) { + int idx = doc["seedPresetIndex"]; + if (idx >= 0 && idx < kNumSeedPresets) s.seedPresetIndex = idx; + } + if (!doc["chatTemplateEnabled"].isNull()) s.chatTemplateEnabled = doc["chatTemplateEnabled"]; + if (!doc["userTag"].isNull()) s.userTag = doc["userTag"].as(); + if (!doc["botTag"].isNull()) s.botTag = doc["botTag"].as(); + } + f.close(); + return s; +} + +void saveSettings(FS &fs, const BruceLMSettings &s) { + JsonDocument doc; + doc["temperature"] = s.temperature; + doc["topP"] = s.topP; + doc["repetitionPenalty"] = s.repetitionPenalty; + doc["maxTokens"] = s.maxTokens; + doc["seedPresetIndex"] = s.seedPresetIndex; + doc["chatTemplateEnabled"] = s.chatTemplateEnabled; + doc["userTag"] = s.userTag; + doc["botTag"] = s.botTag; + File f = fs.open(kConfigPath, FILE_WRITE); + if (!f) return; + serializeJsonPretty(doc, f); + f.close(); +} + +// All of this is deliberately expressed in terms of existing global layout +// constants (BORDER_PAD_X/Y, FP/FM, LW/LH) rather than magic numbers, and the +// outer rounded-rect border is drawn at (5,5,tftWidth-10,tftHeight-10) by +// drawStatusBar() - kBottomMargin/kTopGap keep our own content clear of it. +constexpr int kTopGap = 6; +constexpr int kBottomMargin = 6; +constexpr int kFooterH = 16; +constexpr int kInputBarH = 16; +constexpr int kSepGap = 4; + +struct ChatGeometry { + int outputTop; + int sepY; + int inputBarY; + int footerY; + int lineH; + int maxChars; + int visibleLines; +}; + +ChatGeometry computeGeometry() { + ChatGeometry g; + g.outputTop = BORDER_PAD_Y + FM * LH + kTopGap; + g.footerY = tftHeight - kBottomMargin - kFooterH; + g.inputBarY = g.footerY - kInputBarH; + g.sepY = g.inputBarY - kSepGap; + g.lineH = FP * LH; + g.maxChars = (tftWidth - 2 * BORDER_PAD_X) / (FP * LW); + if (g.maxChars < 1) g.maxChars = 1; + g.visibleLines = (g.sepY - g.outputTop) / g.lineH; + if (g.visibleLines < 1) g.visibleLines = 1; + return g; +} + +// Greedy word-wrap of a single line (no embedded '\n'). Always emits at +// least one (possibly empty) line, so blank paragraphs still take up a row. +std::vector wrapLine(const String &text, int maxChars) { + std::vector out; + int n = text.length(); + int start = 0; + String currentLine; + while (start < n) { + int spaceIdx = text.indexOf(' ', start); + int wordEnd = (spaceIdx == -1) ? n : spaceIdx; + String word = text.substring(start, wordEnd); + while ((int)word.length() > maxChars) { + if (currentLine.length() > 0) { + out.push_back(currentLine); + currentLine = ""; + } + out.push_back(word.substring(0, maxChars)); + word = word.substring(maxChars); + } + if (currentLine.length() == 0) currentLine = word; + else if ((int)(currentLine.length() + 1 + word.length()) <= maxChars) currentLine += " " + word; + else { + out.push_back(currentLine); + currentLine = word; + } + start = (spaceIdx == -1) ? n : spaceIdx + 1; + } + out.push_back(currentLine); + return out; +} + +// Splits on explicit '\n' first (the model can emit these), then word-wraps +// each resulting segment independently. +std::vector wrapText(const String &text, int maxChars) { + std::vector out; + int start = 0; + int n = text.length(); + for (;;) { + int nl = text.indexOf('\n', start); + String segment = (nl == -1) ? text.substring(start) : text.substring(start, nl); + for (auto &l : wrapLine(segment, maxChars)) out.push_back(l); + if (nl == -1) break; + start = nl + 1; + } + return out; +} + +struct ChatLine { + String text; + bool isUser; +}; + +struct ChatSession { + // Each entry is one message ("You: ..." or the assistant's reply text, + // appended to as tokens stream in). true = user message. + std::vector> messages; + int scrollOffset = 0; + bool followBottom = true; + String inputBarPreview; + BruceLMSettings settings; +}; + +std::vector buildLines(const ChatSession &s, int maxChars) { + std::vector out; + for (size_t i = 0; i < s.messages.size(); i++) { + bool isUser = s.messages[i].first; + for (auto &l : wrapText(s.messages[i].second, maxChars)) out.push_back({l, isUser}); + if (i + 1 < s.messages.size()) out.push_back({"", false}); // blank line between messages + } + return out; +} + +// Applies rotary-encoder scroll input. Returns true if the visible content +// actually needs to move as a result. +bool handleScrollInput(ChatSession &s, int totalLines, int visibleLines) { + int maxScroll = totalLines > visibleLines ? totalLines - visibleLines : 0; + bool changed = false; + if (check(NextPress)) { // rotate toward newer/bottom content + if (s.scrollOffset < maxScroll) { + s.scrollOffset++; + changed = true; + } + s.followBottom = (s.scrollOffset >= maxScroll); + } + if (check(PrevPress)) { // rotate toward older/top content + if (s.scrollOffset > 0) { + s.scrollOffset--; + changed = true; + } + s.followBottom = false; + } + return changed; +} + +void clampScroll(ChatSession &s, int totalLines, int visibleLines) { + int maxScroll = totalLines > visibleLines ? totalLines - visibleLines : 0; + if (s.followBottom) s.scrollOffset = maxScroll; + if (s.scrollOffset > maxScroll) s.scrollOffset = maxScroll; + if (s.scrollOffset < 0) s.scrollOffset = 0; +} + +String truncateToWidth(const String &text, int maxChars) { + if ((int)text.length() <= maxChars) return text; + if (maxChars <= 3) return text.substring(0, maxChars); + return text.substring(0, maxChars - 3) + "..."; +} + +void drawChatFrame(const ChatGeometry &g, const String &inputBarPreview) { + drawMainBorderWithTitle("BruceLM"); + + tft.drawFastHLine(BORDER_PAD_X, g.sepY, tftWidth - 2 * BORDER_PAD_X, bruceConfig.priColor); + + tft.setTextColor(bruceConfig.secColor, bruceConfig.bgColor); + int inputTextY = g.inputBarY + (kInputBarH - g.lineH) / 2; + tft.setCursor(BORDER_PAD_X, inputTextY); + String preview = inputBarPreview.length() ? ("> " + inputBarPreview) : "> Tap OK to type a prompt"; + tft.print(truncateToWidth(preview, g.maxChars)); + + tft.setTextColor(bruceConfig.priColor, bruceConfig.bgColor); + String hint = "OK: new prompt Esc: exit"; + int hintTextY = g.footerY + (kFooterH - g.lineH) / 2; + int hx = (tftWidth - (int)(hint.length() * FP * LW)) / 2; + if (hx < BORDER_PAD_X) hx = BORDER_PAD_X; + tft.setCursor(hx, hintTextY); + tft.print(hint); +} + +void drawVisibleLines(const std::vector &lines, int scrollOffset, const ChatGeometry &g) { + for (int i = 0; i < g.visibleLines; i++) { + int idx = scrollOffset + i; + if (idx >= (int)lines.size()) break; + tft.setTextColor( + lines[idx].isUser ? bruceConfig.priColor : bruceConfig.secColor, bruceConfig.bgColor + ); + tft.setCursor(BORDER_PAD_X, g.outputTop + i * g.lineH); + tft.print(lines[idx].text); + } +} + +// Single entry point for both drawing and encoder-scroll handling, so the +// idle loop and the token-streaming loop can never disagree about how +// scrolling behaves. Returns true if it actually redrew the screen. +bool renderChat(ChatSession &s, bool force) { + ChatGeometry g = computeGeometry(); + std::vector lines = buildLines(s, g.maxChars); + int total = (int)lines.size(); + + bool scrolled = handleScrollInput(s, total, g.visibleLines); + clampScroll(s, total, g.visibleLines); + + if (!force && !scrolled) return false; + + drawChatFrame(g, s.inputBarPreview); + drawVisibleLines(lines, s.scrollOffset, g); + return true; +} + +void ensureFolders(FS &fs) { + // sd_functions.cpp's createFolder() is the interactive file-manager + // "New Folder" action - it always pops a keyboard prompt for a name and + // creates it *under* the given path. We want these two fixed paths to + // just silently exist, so call fs.mkdir() directly instead. Without an + // SD card (LittleFS fallback, where these folders don't already exist + // from a prior run) createFolder() here was popping the folder-name + // keyboard twice on every launch before ever reaching the start screen. + if (!folderExists(fs, kRootDir)) fs.mkdir(kRootDir); + if (!folderExists(fs, kModelsDir)) fs.mkdir(kModelsDir); +} + +// Draws a message starting at the top of the content area, an optional +// standout note centered in the gap between that text and the footer (for +// calling out the required action - e.g. "Press OK to..." - so it isn't +// just a small line buried in the bottom hint), and a hint row pinned to +// the bottom (same geometry the chat frame uses) so button mappings always +// land in the same place instead of wherever the message text happened to end. +void drawMessageScreen(const String &msg, const String &hint, const String ¢erNote = "") { + ChatGeometry g = computeGeometry(); + drawMainBorderWithTitle("BruceLM"); + + tft.setTextColor(bruceConfig.priColor, bruceConfig.bgColor); + int y = g.outputTop; + for (auto &line : wrapText(msg, g.maxChars)) { + tft.setCursor(BORDER_PAD_X, y); + tft.print(line); + y += g.lineH; + } + + if (centerNote.length() > 0) { + int aboveFooterY = g.footerY - 4; + // Biased toward the footer (rather than dead-centered in the gap) so + // the note sits a bit further down the screen. + int gapSpace = max(0, aboveFooterY - y - g.lineH); + int noteY = y + (gapSpace * 2) / 3; + tft.setTextColor(TFT_YELLOW, bruceConfig.bgColor); + for (auto &line : wrapText(centerNote, g.maxChars)) { + int nx = (tftWidth - (int)(line.length() * FP * LW)) / 2; + if (nx < BORDER_PAD_X) nx = BORDER_PAD_X; + tft.setCursor(nx, noteY); + tft.print(line); + noteY += g.lineH; + } + } + + tft.setTextColor(bruceConfig.priColor, bruceConfig.bgColor); + int hintTextY = g.footerY + (kFooterH - g.lineH) / 2; + int hx = (tftWidth - (int)(hint.length() * FP * LW)) / 2; + if (hx < BORDER_PAD_X) hx = BORDER_PAD_X; + tft.setCursor(hx, hintTextY); + tft.print(hint); +} + +void showMessageAndWait(const String &msg) { + drawMessageScreen(msg, "Esc: return"); + while (!check(EscPress)) delay(30); +} + +// Blocks until the user picks Select (true) or Esc (false). +bool showConfirm(const String &msg, const String &hint, const String ¢erNote = "") { + drawMessageScreen(msg, hint, centerNote); + for (;;) { + if (check(SelPress)) return true; + if (check(EscPress)) return false; + delay(30); + } +} + +float adjustFloat(float v, float step, float lo, float hi, int dir) { + v += step * dir; + if (v < lo) v = lo; + if (v > hi) v = hi; + return v; +} + +int adjustInt(int v, int step, int lo, int hi, int dir) { + v += step * dir; + if (v < lo) v = lo; + if (v > hi) v = hi; + return v; +} + +int cycleIndex(int idx, int count, int dir) { return (idx + dir + count) % count; } + +constexpr int kSettingsRowCount = 7; +constexpr int kSettingsSaveRow = kSettingsRowCount; +constexpr int kSettingsNumItems = kSettingsRowCount + 1; + +void ensureListScroll(int cursor, int itemCount, int visibleRows, int &scroll) { + if (itemCount <= visibleRows) { + scroll = 0; + return; + } + if (cursor < scroll) scroll = cursor; + if (cursor >= scroll + visibleRows) scroll = cursor - visibleRows + 1; +} + +struct SettingsEditState { + bool editing = false; + int editRow = -1; + float tempBackup = 0, topPBackup = 0, repPenBackup = 0; + int maxTokBackup = 0, seedIdxBackup = 0; +}; + +void renderSettingsRows( + const BruceLMSettings &s, int cursor, const SettingsEditState &edit, int startY, int rowH, + int scrollOffset, int visibleRows +) { + struct Row { + const char *label; + String value; + } rows[kSettingsRowCount] = { + {"Temperature", String(s.temperature, 1)}, + {"Top-p", String(s.topP, 2)}, + {"Rep. Penalty", String(s.repetitionPenalty, 2)}, + {"Max Tokens", String(s.maxTokens)}, + {"Seed", String(kSeedPresets[s.seedPresetIndex].label)}, + {"Chat Template", s.chatTemplateEnabled ? "On >" : "Off >"}, + {"Reset to Default", "Reset >"}, + }; + + tft.setTextSize(FP); + for (int i = 0; i < visibleRows; i++) { + int idx = scrollOffset + i; + int rowY = startY + i * rowH; + tft.fillRect(BORDER_PAD_X, rowY, tftWidth - 2 * BORDER_PAD_X, rowH, bruceConfig.bgColor); + if (idx >= kSettingsNumItems) continue; + + bool selected = (cursor == idx); + + if (idx == kSettingsSaveRow) { + tft.setTextColor(selected ? TFT_YELLOW : bruceConfig.priColor, bruceConfig.bgColor); + String saveLabel = "[ Save ]"; + int sx = (tftWidth - (int)(saveLabel.length() * FP * LW)) / 2; + tft.setCursor(sx, rowY + 2); + tft.print(saveLabel); + continue; + } + + bool editingThis = (edit.editing && edit.editRow == idx); + + tft.setTextColor(selected ? TFT_YELLOW : bruceConfig.priColor, bruceConfig.bgColor); + tft.setCursor(BORDER_PAD_X, rowY + 2); + tft.print(rows[idx].label); + + String valueText = editingThis ? ("[ " + rows[idx].value + " ]") : rows[idx].value; + int vx = tftWidth - BORDER_PAD_X - (int)(valueText.length() * FP * LW); + tft.setCursor(vx, rowY + 2); + tft.print(valueText); + } +} + +// Sub-menu opened from the "Chat Template" row: toggles the template on/off +// and lets the user re-type the tags with the same keyboard() overlay used +// for prompt entry. Mutates `s` in place - the parent settings screen owns +// persisting it (via its own Save row / Esc-to-save), so nothing is written +// to disk here. +void showChatTemplateScreen(BruceLMSettings &s) { + constexpr int kNumRows = 3; + constexpr int kSaveRow = 3; // bottom action row, same spot/style as the parent screen's Save + constexpr int kNumItems = 4; + int cursor = 0; + int scroll = 0; + ChatGeometry g = computeGeometry(); + int rowH = g.lineH + 4; + int startY = g.outputTop; + int visibleRows = max(1, (g.footerY - startY) / rowH); + bool redraw = true; + + for (;;) { + if (redraw) { + drawMainBorderWithTitle("Chat Template"); + ensureListScroll(cursor, kNumItems, visibleRows, scroll); + + struct Row { + const char *label; + String value; + } rows[kNumRows] = { + {"Enabled", s.chatTemplateEnabled ? "On" : "Off"}, + {"User tag", s.userTag }, + {"Bot tag", s.botTag }, + }; + + tft.setTextSize(FP); + for (int i = 0; i < visibleRows; i++) { + int idx = scroll + i; + int rowY = startY + i * rowH; + tft.fillRect(BORDER_PAD_X, rowY, tftWidth - 2 * BORDER_PAD_X, rowH, bruceConfig.bgColor); + if (idx >= kNumItems) continue; + + bool selected = (cursor == idx); + if (idx == kSaveRow) { + tft.setTextColor(selected ? TFT_YELLOW : bruceConfig.priColor, bruceConfig.bgColor); + String saveLabel = "[ Save ]"; + int sx = (tftWidth - (int)(saveLabel.length() * FP * LW)) / 2; + tft.setCursor(sx, rowY + 2); + tft.print(saveLabel); + continue; + } + + tft.setTextColor(selected ? TFT_YELLOW : bruceConfig.priColor, bruceConfig.bgColor); + tft.setCursor(BORDER_PAD_X, rowY + 2); + tft.print(rows[idx].label); + + int labelW = (int)String(rows[idx].label).length() + 1; + String valueText = truncateToWidth(rows[idx].value, max(1, g.maxChars - labelW)); + int vx = tftWidth - BORDER_PAD_X - (int)(valueText.length() * FP * LW); + tft.setCursor(vx, rowY + 2); + tft.print(valueText); + } + + tft.setTextColor(bruceConfig.priColor, bruceConfig.bgColor); + String hint = "OK: edit Esc: save+exit"; + int hintTextY = g.footerY + (kFooterH - g.lineH) / 2; + tft.fillRect( + BORDER_PAD_X, hintTextY - 2, tftWidth - 2 * BORDER_PAD_X, g.lineH + 4, bruceConfig.bgColor + ); + int hx = (tftWidth - (int)(hint.length() * FP * LW)) / 2; + if (hx < BORDER_PAD_X) hx = BORDER_PAD_X; + tft.setCursor(hx, hintTextY); + tft.print(hint); + redraw = false; + } + + // Every field here mutates `s` immediately (no separate edit-mode + // buffer to revert) - Save and Esc are therefore equivalent, both + // just returning to the parent screen, which owns actually + // persisting `s` to disk. Kept as two controls anyway to match the + // parent settings screen's look and feel. + if (check(EscPress)) return; + if (check(SelPress)) { + if (cursor == kSaveRow) { + return; + } else if (cursor == 0) { + s.chatTemplateEnabled = !s.chatTemplateEnabled; + } else if (cursor == 1) { + String v = keyboard(s.userTag, 32, "User tag:"); + if (v.length() > 0) s.userTag = v; + } else { + String v = keyboard(s.botTag, 32, "Bot tag:"); + if (v.length() > 0) s.botTag = v; + } + redraw = true; + } else if (check(NextPress)) { + cursor = (cursor + 1) % kNumItems; + redraw = true; + } else if (check(PrevPress)) { + cursor = (cursor + kNumItems - 1) % kNumItems; + redraw = true; + } + + delay(30); + } +} + +// Sub-menu opened from the "Reset to Default" row: explains what's about to +// happen, then a single centered "Confirm" action (same highlighted-row look +// as showStartScreen's Start/Settings rows) rather than a bare Select/Esc +// dialog, so it can't be triggered by one accidental button press. Returns +// true if the user confirmed (config file deleted); the caller is +// responsible for treating that as "abandon whatever was in progress". +bool showResetToDefaultScreen(FS &fs) { + ChatGeometry g = computeGeometry(); + String msg = "Reset all BruceLM settings to\n" + "default values? This deletes\n" + "the saved config file.\n" + "\n" + "Models on the SD card are not\n" + "affected."; + std::vector lines = wrapText(msg, g.maxChars); + + int rowH = g.lineH + 6; + int afterTextY = g.outputTop + (int)lines.size() * g.lineH + 4; + int aboveFooterY = g.footerY - 4; + int confirmY = afterTextY + max(0, (aboveFooterY - afterTextY - rowH) / 2); + + bool redraw = true; + for (;;) { + if (redraw) { + drawMainBorderWithTitle("Reset to Default"); + tft.setTextColor(bruceConfig.priColor, bruceConfig.bgColor); + int y = g.outputTop; + for (auto &line : lines) { + tft.setCursor(BORDER_PAD_X, y); + tft.print(line); + y += g.lineH; + } + + tft.fillRect(BORDER_PAD_X, confirmY, tftWidth - 2 * BORDER_PAD_X, rowH, bruceConfig.bgColor); + tft.setTextColor(TFT_YELLOW, bruceConfig.bgColor); + String label = "> Confirm"; + int lx = (tftWidth - (int)(label.length() * FP * LW)) / 2; + tft.setCursor(lx, confirmY + 2); + tft.print(label); + + tft.setTextColor(bruceConfig.priColor, bruceConfig.bgColor); + String hint = "OK: confirm Esc: cancel"; + int hintTextY = g.footerY + (kFooterH - g.lineH) / 2; + int hx = (tftWidth - (int)(hint.length() * FP * LW)) / 2; + if (hx < BORDER_PAD_X) hx = BORDER_PAD_X; + tft.setCursor(hx, hintTextY); + tft.print(hint); + redraw = false; + } + + if (check(EscPress)) return false; + if (check(SelPress)) { + if (fs.exists(kConfigPath)) fs.remove(kConfigPath); + return true; + } + delay(30); + } +} + +// ble_spam.cpp-style row editor: Next/Prev moves the cursor between rows, +// Select enters/exits edit mode on the highlighted row, Next/Prev while +// editing adjusts that row's value, Esc while editing reverts to the value +// it had before editing started, Esc while browsing saves and exits. The +// "Chat Template" and "Reset to Default" rows are the exceptions - Select on +// either always opens its own sub-screen instead of inline-editing, since +// neither is a single adjustable number. +// +// Returns true if the user reset to defaults (config file deleted) - the +// caller should treat this as "abandon whatever was in progress and go back +// to the top-level BruceLM menu" rather than resuming with the just-deleted +// settings. +bool showSettingsScreen(FS &fs) { + BruceLMSettings s = loadSettings(fs); + SettingsEditState edit; + int cursor = 0; + int scroll = 0; + ChatGeometry g = computeGeometry(); + int rowH = g.lineH + 4; + int startY = g.outputTop; + int visibleRows = max(1, (g.footerY - startY) / rowH); + + bool redraw = true; + drawMainBorderWithTitle("BruceLM Settings"); + + constexpr int kChatTemplateRow = 5; // opens showChatTemplateScreen() on Select + constexpr int kResetRow = 6; // opens showResetToDefaultScreen() on Select + constexpr int kSaveRow = kSettingsSaveRow; // bottom action row, ble_spam "[ Start ]"-style + constexpr int kNumItems = kSettingsNumItems; // 7 rows + save + + for (;;) { + if (redraw) { + ensureListScroll(cursor, kNumItems, visibleRows, scroll); + renderSettingsRows(s, cursor, edit, startY, rowH, scroll, visibleRows); + + tft.setTextColor(bruceConfig.priColor, bruceConfig.bgColor); + String hint = edit.editing ? "OK: confirm Esc: revert" : "OK: edit Esc: save+exit"; + int hintTextY = g.footerY + (kFooterH - g.lineH) / 2; + tft.fillRect( + BORDER_PAD_X, hintTextY - 2, tftWidth - 2 * BORDER_PAD_X, g.lineH + 4, bruceConfig.bgColor + ); + int hx = (tftWidth - (int)(hint.length() * FP * LW)) / 2; + if (hx < BORDER_PAD_X) hx = BORDER_PAD_X; + tft.setCursor(hx, hintTextY); + tft.print(hint); + redraw = false; + } + + if (check(EscPress)) { + // Esc always means "save and exit" from the row list (or revert + // while editing) - unchanged regardless of the Save row below. + if (edit.editing) { + switch (edit.editRow) { + case 0: s.temperature = edit.tempBackup; break; + case 1: s.topP = edit.topPBackup; break; + case 2: s.repetitionPenalty = edit.repPenBackup; break; + case 3: s.maxTokens = edit.maxTokBackup; break; + case 4: s.seedPresetIndex = edit.seedIdxBackup; break; + } + edit.editing = false; + redraw = true; + } else { + saveSettings(fs, s); + return false; + } + } else if (check(SelPress)) { + if (edit.editing) { + edit.editing = false; + } else if (cursor == kSaveRow) { + saveSettings(fs, s); + return false; + } else if (cursor == kResetRow) { + if (showResetToDefaultScreen(fs)) return true; + drawMainBorderWithTitle("BruceLM Settings"); + } else if (cursor == kChatTemplateRow) { + showChatTemplateScreen(s); + drawMainBorderWithTitle("BruceLM Settings"); + } else { + edit.editing = true; + edit.editRow = cursor; + edit.tempBackup = s.temperature; + edit.topPBackup = s.topP; + edit.repPenBackup = s.repetitionPenalty; + edit.maxTokBackup = s.maxTokens; + edit.seedIdxBackup = s.seedPresetIndex; + } + redraw = true; + } else if (edit.editing) { + int dir = check(NextPress) ? 1 : (check(PrevPress) ? -1 : 0); + if (dir != 0) { + switch (edit.editRow) { + case 0: s.temperature = adjustFloat(s.temperature, 0.1f, 0.0f, 2.0f, dir); break; + case 1: s.topP = adjustFloat(s.topP, 0.05f, 0.0f, 1.0f, dir); break; + case 2: + s.repetitionPenalty = adjustFloat(s.repetitionPenalty, 0.05f, 1.0f, 2.0f, dir); + break; + case 3: s.maxTokens = adjustInt(s.maxTokens, 16, 16, 2048, dir); break; + case 4: s.seedPresetIndex = cycleIndex(s.seedPresetIndex, kNumSeedPresets, dir); break; + } + redraw = true; + } + } else if (check(NextPress)) { + cursor = (cursor + 1) % kNumItems; + redraw = true; + } else if (check(PrevPress)) { + cursor = (cursor + kNumItems - 1) % kNumItems; + redraw = true; + } + + delay(30); + } +} + +enum class StartAction { Start, HowToRun, Settings, Exit }; + +// First screen: a one-line pitch above three selectable action rows +// ("Start" / "How to Run" / "Settings"), same cursor+Next/Prev+Select style +// as the settings screen above, with the button hints pinned to the bottom. +StartAction showStartScreen() { + int cursor = 0; + const char *labels[3] = {"Start", "How to Run", "Settings"}; + constexpr int kNumRows = 3; + ChatGeometry g = computeGeometry(); + std::vector pitchLines = wrapText("Run small LLMs locally on your Bruce device!", g.maxChars); + + int rowH = g.lineH + 6; + // A fixed-size gap under the header for the pitch text, so the button + // block's position doesn't shift around based on how long the pitch is. + int headerGap = g.lineH * 3; + int afterTextY = g.outputTop + headerGap; + int aboveFooterY = g.footerY - 4; + int rowsBlockH = rowH * kNumRows; + // Biased toward the footer (rather than dead-centered in the gap) so the + // buttons sit a bit further down the screen. + int gapSpace = max(0, aboveFooterY - afterTextY - rowsBlockH); + int firstRowY = afterTextY + (gapSpace * 2) / 3; + // Pitch text biased toward the bottom of its header gap (rather than + // dead-centered) so it sits a bit further down, closer to the Start button. + int pitchBlockH = (int)pitchLines.size() * g.lineH; + int pitchY = g.outputTop + max(0, ((headerGap - pitchBlockH) * 2) / 3); + + bool redraw = true; + for (;;) { + if (redraw) { + drawMainBorderWithTitle("BruceLM"); + tft.setTextColor(bruceConfig.priColor, bruceConfig.bgColor); + int y = pitchY; + for (auto &line : pitchLines) { + int lx = (tftWidth - (int)(line.length() * FP * LW)) / 2; + if (lx < BORDER_PAD_X) lx = BORDER_PAD_X; + tft.setCursor(lx, y); + tft.print(line); + y += g.lineH; + } + + for (int i = 0; i < kNumRows; i++) { + int rowY = firstRowY + i * rowH; + tft.fillRect(BORDER_PAD_X, rowY, tftWidth - 2 * BORDER_PAD_X, rowH, bruceConfig.bgColor); + tft.setTextColor(cursor == i ? TFT_YELLOW : bruceConfig.priColor, bruceConfig.bgColor); + String label = String("> ") + labels[i]; + int lx = (tftWidth - (int)(label.length() * FP * LW)) / 2; + tft.setCursor(lx, rowY + 2); + tft.print(label); + } + + tft.setTextColor(bruceConfig.priColor, bruceConfig.bgColor); + String hint = "OK: select Esc: exit"; + int hintTextY = g.footerY + (kFooterH - g.lineH) / 2; + int hx = (tftWidth - (int)(hint.length() * FP * LW)) / 2; + if (hx < BORDER_PAD_X) hx = BORDER_PAD_X; + tft.setCursor(hx, hintTextY); + tft.print(hint); + redraw = false; + } + + if (check(EscPress)) return StartAction::Exit; + if (check(SelPress)) { + if (cursor == 0) return StartAction::Start; + if (cursor == 1) return StartAction::HowToRun; + return StartAction::Settings; + } + if (check(NextPress)) { + cursor = (cursor + 1) % kNumRows; + redraw = true; + } else if (check(PrevPress)) { + cursor = (cursor + kNumRows - 1) % kNumRows; + redraw = true; + } + delay(30); + } +} + +// Opened from the "How to Run" row: where to get models and where to put +// them, ending in a single centered "Okay" action (same look as +// showResetToDefaultScreen's "Confirm") that just returns to the caller, +// which redisplays the start screen. The "Okay" row is pinned just above the +// footer hint regardless of content length (sized to just fit its own text, +// not a full row like the settings list), and the text above it scrolls with +// the encoder (same NextPress/PrevPress convention as the chat view) - on +// smaller screens the full text otherwise overflows past the border and +// overlaps the fixed-position button/hint row. When there's more content +// than fits, the last visible line is deliberately clipped to half height +// (rather than simply omitted) as a visual cue that there's more below. +void showHowToRunScreen() { + ChatGeometry g = computeGeometry(); + // Built as explicit (text, indented, dimmed) rows rather than a single + // wrapped string - wrapLine() treats runs of leading spaces as word + // breaks and silently eats them, so indentation can't survive a plain + // "\n"-joined string round-tripped through wrapText(). + struct InfoRow { + String text; + bool indented; + bool dimmed; + }; + std::vector lines = { + {"1. Download zipped model and", false, false}, + {"tokenizer files from:", true, false}, + {"https://archive.org/details/BruceLM", true, true }, + {"2. Unzip both files to SD card", false, false}, + {"BruceLM/models/:", true, false}, + {"models/Chat3.3M/chat3.3M.bin", true, true }, + {"models/Chat3.3M/tok1024.bin", true, true }, + {"models/Stories260K/stories260K.bin", true, true }, + {"models/Stories260K/tok512.bin", true, true }, + {"3. Turn off chat templates", false, false}, + {"for stories260K", true, false}, + {"4. You can create your own models", false, false}, + {"using llama2.c by Karpathy", true, false}, + }; + int indentPx = FP * LW * 3; + int lineH = g.lineH; + + int okayRowH = lineH + 4; // just tall enough for the label, not a full list row + int okayY = g.footerY - 4 - okayRowH; // pinned above the footer, independent of content length + + // Reserve one row at the top for the scroll-direction indicators, kept + // out of the text flow itself. + int textTop = g.outputTop; + int contentBottom = okayY - 4; + // Always leave room for a half-height clipped line at the bottom, so a + // full line's worth of space isn't silently reserved-but-unused when + // there's nothing left to scroll to. + int visibleFullLines = max(1, (contentBottom - textTop - lineH / 2) / lineH); + int scrollOffset = 0; + int maxScroll = max(0, (int)lines.size() - visibleFullLines); + + bool redraw = true; + for (;;) { + if (redraw) { + drawMainBorderWithTitle("How to Run"); + + int y = textTop; + int i = 0; + for (; i < visibleFullLines; i++) { + int idx = scrollOffset + i; + if (idx >= (int)lines.size()) break; + InfoRow &row = lines[idx]; + tft.setTextColor( + row.dimmed ? bruceConfig.secColor : bruceConfig.priColor, bruceConfig.bgColor + ); + tft.setCursor(BORDER_PAD_X + (row.indented ? indentPx : 0), y); + tft.print(row.text); + y += lineH; + } + // One more line, then mask its bottom half back out with the + // background color - Bruce's tft wrapper doesn't expose + // TFT_eSPI's setViewport/clipping, so "print then paint over" + // is the only way to fake a half-height clip here. Visual cue + // that there's more content below to scroll to. + int cutIdx = scrollOffset + i; + if (cutIdx < (int)lines.size()) { + InfoRow &row = lines[cutIdx]; + tft.setTextColor( + row.dimmed ? bruceConfig.secColor : bruceConfig.priColor, bruceConfig.bgColor + ); + tft.setCursor(BORDER_PAD_X + (row.indented ? indentPx : 0), y); + tft.print(row.text); + tft.fillRect( + BORDER_PAD_X, + y + lineH / 2, + tftWidth - 2 * BORDER_PAD_X, + lineH - lineH / 2, + bruceConfig.bgColor + ); + } + + tft.fillRect(BORDER_PAD_X, okayY, tftWidth - 2 * BORDER_PAD_X, okayRowH, bruceConfig.bgColor); + tft.setTextColor(TFT_YELLOW, bruceConfig.bgColor); + String label = "> Okay"; + int lx = (tftWidth - (int)(label.length() * FP * LW)) / 2; + tft.setCursor(lx, okayY + 2); + tft.print(label); + + tft.setTextColor(bruceConfig.priColor, bruceConfig.bgColor); + String hint = "OK: back to menu Esc: back"; + int hintTextY = g.footerY + (kFooterH - lineH) / 2; + int hx = (tftWidth - (int)(hint.length() * FP * LW)) / 2; + if (hx < BORDER_PAD_X) hx = BORDER_PAD_X; + tft.setCursor(hx, hintTextY); + tft.print(hint); + redraw = false; + } + + if (check(EscPress) || check(SelPress)) return; + if (check(NextPress)) { + if (scrollOffset < maxScroll) { + scrollOffset++; + redraw = true; + } + } else if (check(PrevPress)) { + if (scrollOffset > 0) { + scrollOffset--; + redraw = true; + } + } + delay(30); + } +} + +void showLoadingScreen(const String &modelFileName, const String &tokenizerFileName) { + ChatGeometry g = computeGeometry(); + drawMainBorderWithTitle("BruceLM"); + tft.setTextColor(bruceConfig.priColor, bruceConfig.bgColor); + tft.setCursor(BORDER_PAD_X, g.outputTop); + tft.print("Loading model + tokenizer..."); + tft.setTextColor(bruceConfig.secColor, bruceConfig.bgColor); + tft.setCursor(BORDER_PAD_X, g.outputTop + g.lineH * 2); + tft.print(truncateToWidth(modelFileName, g.maxChars)); + tft.setCursor(BORDER_PAD_X, g.outputTop + g.lineH * 3); + tft.print(truncateToWidth(tokenizerFileName, g.maxChars)); +} + +// Streams the transcript into the top output area; returns false if the user +// cancelled generation (Esc) partway through. +bool runChatTurn(ChatSession &session, LLMEngine &engine, const String &userPrompt) { + session.messages.push_back({true, "You: " + userPrompt}); + session.messages.push_back({false, "> "}); // "> " prefix, same style as the input bar preview + session.inputBarPreview = userPrompt; + session.followBottom = true; + renderChat(session, /*force=*/true); + + GenerationParams params; + params.temperature = session.settings.temperature; + params.topP = session.settings.topP; + params.repetitionPenalty = session.settings.repetitionPenalty; + params.seed = session.settings.seed(); + params.chatTemplateEnabled = session.settings.chatTemplateEnabled; + params.userTag = session.settings.userTag; + params.botTag = session.settings.botTag; + + bool cancelled = false; + engine.generate(userPrompt, session.settings.maxTokens, params, [&](const String &piece) -> bool { + // Generation can run longer than Bruce's normal screen-timeout + // window. wakeUpScreen() just resets its inactivity clock (the + // same thing any real button press does) - calling it here from + // our own module, not touching core/display.cpp itself, keeps + // the screen from dimming/sleeping mid-reply. + wakeUpScreen(); + if (check(EscPress)) { + cancelled = true; + return false; + } + session.messages.back().second += piece; + renderChat(session, /*force=*/true); + return true; + }); + + if (cancelled) { + session.messages.back().second += " [cancelled]"; + renderChat(session, /*force=*/true); + } + return !cancelled; +} + +// Exit: user backed out of the chat (Esc) - bruceLM_setup() should exit the +// whole module. ResetToMenu: settings were reset to defaults from within the +// chat - bruceLM_setup() should abandon this session and redisplay the +// top-level start screen instead. +enum class ChatLoopResult { Exit, ResetToMenu }; + +ChatLoopResult chatLoop(LLMEngine &engine, FS &fs, const BruceLMSettings &settings) { + ChatSession session; + session.settings = settings; + renderChat(session, /*force=*/true); + + for (;;) { + if (check(EscPress)) return ChatLoopResult::Exit; + if (check(SelPress)) { + String prompt = keyboard("", 200, "Prompt:"); + String trimmed = prompt; + trimmed.trim(); + trimmed.toLowerCase(); + if (trimmed == "/settings") { + // No spare physical input to dedicate to "open settings" while + // chatting (Next/Prev already scroll, Select starts a prompt, + // Esc exits) - reusing the existing prompt entry as a command + // avoids adding a new input mapping. + if (showSettingsScreen(fs)) return ChatLoopResult::ResetToMenu; + session.settings = loadSettings(fs); + renderChat(session, /*force=*/true); + } else if (prompt.length() > 0) { + runChatTurn(session, engine, prompt); + } + } else { + renderChat(session, /*force=*/false); // only redraws if the encoder scrolled + } + delay(30); + } +} + +// IncompatibleGroupSize and ConfigTooLarge are both "this will probably go +// badly, but try anyway?" situations - one confirm+retry path for both +// instead of duplicating the same dialog/retry logic per error code. +LLMLoadError confirmRiskyLoad( + FS &fs, const String &checkpointPath, const String &tokenizerPath, LLMEngine &engine, LLMLoadError err, + const String &modelFileName, const String &tokenizerFileName +) { + String warning; + if (err == LLMLoadError::IncompatibleGroupSize) { + warning = "This model's quantization group\n" + "size doesn't evenly divide its\n" + "layer sizes - output will likely\n" + "come out garbled."; + } else if (err == LLMLoadError::ConfigTooLarge) { + warning = "This model looks too large for\n" + "available memory - it may fail to\n" + "load or crash mid-reply."; + } else { + return err; + } + + if (!showConfirm(warning, "OK: run anyway Esc: cancel")) return err; + showLoadingScreen(modelFileName, tokenizerFileName); + return engine.load(fs, checkpointPath, tokenizerPath, /*overrideSafetyChecks=*/true); +} + +} // namespace + +void bruceLM_setup() { + FS *fs; + if (!getFsStorage(fs)) { + showMessageAndWait("No storage available (SD/LittleFS)."); + return; + } + + ensureFolders(*fs); + + for (;;) { + StartAction action = showStartScreen(); + if (action == StartAction::Exit) return; + if (action == StartAction::HowToRun) { + showHowToRunScreen(); + continue; + } + if (action == StartAction::Settings) { + // Return value (reset-to-default) is irrelevant here - we're + // already about to redisplay the start screen either way. + showSettingsScreen(*fs); + continue; + } + // action == StartAction::Start + + BruceLMSettings settings = loadSettings(*fs); + + if (!showConfirm( + "Select a model checkpoint file.\n" + "\n" + "(e.g. \"chat3.3M.bin\")", + "OK: continue Esc: cancel", + "Press OK to open the file picker" + )) + return; + String checkpointPath = loopSD(*fs, true, "bin", kModelsDir); + if (checkpointPath.length() == 0) return; // user backed out of the picker + + if (!showConfirm( + "Select the tokenizer file that\n" + "matches your model.\n" + "\n" + "(e.g. \"tok1024.bin\")", + "OK: continue Esc: cancel", + "Press OK to open the file picker" + )) + return; + String tokenizerPath = loopSD(*fs, true, "bin", kModelsDir); + if (tokenizerPath.length() == 0) return; // user backed out of the picker + + String modelFileName = checkpointPath.substring(checkpointPath.lastIndexOf('/') + 1); + String tokenizerFileName = tokenizerPath.substring(tokenizerPath.lastIndexOf('/') + 1); + showLoadingScreen(modelFileName, tokenizerFileName); + + LLMEngine engine; + LLMLoadError err = engine.load(*fs, checkpointPath, tokenizerPath); + + if (err == LLMLoadError::IncompatibleGroupSize || err == LLMLoadError::ConfigTooLarge) + err = confirmRiskyLoad( + *fs, checkpointPath, tokenizerPath, engine, err, modelFileName, tokenizerFileName + ); + + switch (err) { + case LLMLoadError::None: break; + case LLMLoadError::CheckpointNotFound: + case LLMLoadError::TokenizerNotFound: + showMessageAndWait("Model or tokenizer file went missing."); + return; + case LLMLoadError::BadMagicOrVersion: + showMessageAndWait( + "Unrecognized checkpoint format.\n" + "Expected a llama2.c export.py v1/v2\n" + "file, or the original header-less\n" + "format karpathy/tinyllamas ships." + ); + return; + case LLMLoadError::ConfigTooLarge: + showMessageAndWait("Could not load: model too large for available memory."); + return; + case LLMLoadError::OutOfMemory: + showMessageAndWait("Out of memory while allocating model buffers."); + return; + case LLMLoadError::IncompatibleGroupSize: + showMessageAndWait("Could not load: incompatible quantized model."); + return; + } + + // ResetToMenu (settings reset to default from within the chat) + // loops back to the top-level start screen instead of exiting. + if (chatLoop(engine, *fs, settings) == ChatLoopResult::ResetToMenu) continue; + return; + } +} diff --git a/src/modules/others/BruceLM/bruce_lm.h b/src/modules/others/BruceLM/bruce_lm.h new file mode 100644 index 0000000000..5cced53e48 --- /dev/null +++ b/src/modules/others/BruceLM/bruce_lm.h @@ -0,0 +1,6 @@ +#pragma once + +// Entry point wired into OthersMenu. Settings are reached from inside the +// module itself (its first screen, or typing "/settings" while chatting), +// not as a separate top-level menu entry. +void bruceLM_setup(); diff --git a/src/modules/others/BruceLM/llm_engine.cpp b/src/modules/others/BruceLM/llm_engine.cpp new file mode 100644 index 0000000000..bcc7e6c6e1 --- /dev/null +++ b/src/modules/others/BruceLM/llm_engine.cpp @@ -0,0 +1,803 @@ +/** + * @file llm_engine.cpp + * @brief Forward pass for llama2.c-format checkpoints, ported for ESP32-S3. + * + * ---- Made by @Doominator1 on GitHub, Jul 2026 + * ---- Much thanks to @karpathy for the base of this project; llama2.c! + * + * File-backed (no mmap) port of karpathy/llama2.c's run.c. Reads the "v2" + * export format written by export.py's `version()` path: + * uint32 magic ("ak42" as little-endian int, i.e. 0x616b3432) + * int32 version (0 = fp32, 2 = int8/Q8_0 quantized) + * Config (7 x int32) + * uint8 shared_classifier + * int32 group_size (quantized only) + * ... weights ... + * Legacy (no-header) checkpoints are not supported - re-export with export.py. + */ +#include "llm_engine.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace bruce_llm { + +namespace { + +constexpr uint32_t kMagic = 0x616b3432; // "ak42" + +void *psram_alloc(size_t n) { + void *p = heap_caps_malloc(n, MALLOC_CAP_SPIRAM | MALLOC_CAP_8BIT); + return p; +} + +struct QuantizedTensor { + int8_t *q = nullptr; // quantized values + float *s = nullptr; // per-group scales +}; + +// Vocab entries sorted by string, for run.c-style binary-search token lookup +// during BPE encoding (str_lookup()/sorted_vocab in run.c's encode()). +struct TokenIndex { + const char *str; + int id; +}; + +int strLookup(const char *str, const std::vector &sorted) { + auto it = std::lower_bound(sorted.begin(), sorted.end(), str, [](const TokenIndex &a, const char *s) { + return strcmp(a.str, s) < 0; + }); + if (it != sorted.end() && strcmp(it->str, str) == 0) return it->id; + return -1; +} + +float *readF32Vec(File &f, size_t n) { + float *buf = (float *)psram_alloc(n * sizeof(float)); + if (!buf) return nullptr; + f.read((uint8_t *)buf, n * sizeof(float)); + return buf; +} + +// Reads `units` back-to-back quantized sub-tensors of `sizePerUnit` elements +// each (e.g. one per transformer layer), matching the on-disk layout written +// by export.py's version2_export and read by upstream runq.c's +// init_quantized_tensors(): every unit's int8 values are immediately +// followed by that same unit's fp32 scale factors, unit by unit - +// [q_0][s_0][q_1][s_1]...[q_{units-1}][s_{units-1}] - NOT one contiguous +// block of all units' q values followed by one block of all their s values. +// The in-memory result is still a single contiguous QuantizedTensor (q and s +// each span all units), so layerView()/dequantize()/matmulQ() are unchanged; +// only the read order from disk differs, to line up with the real format. +bool readQuantized(File &f, QuantizedTensor &t, size_t sizePerUnit, int units, int groupSize) { + size_t groupsPerUnit = sizePerUnit / groupSize; + size_t totalSize = sizePerUnit * (size_t)units; + size_t totalGroups = groupsPerUnit * (size_t)units; + t.q = (int8_t *)psram_alloc(totalSize); + t.s = (float *)psram_alloc(totalGroups * sizeof(float)); + if (!t.q || !t.s) return false; + for (int i = 0; i < units; i++) { + f.read((uint8_t *)(t.q + (size_t)i * sizePerUnit), sizePerUnit); + f.read((uint8_t *)(t.s + (size_t)i * groupsPerUnit), groupsPerUnit * sizeof(float)); + } + return true; +} + +void dequantize(const QuantizedTensor &t, float *out, size_t n, int groupSize) { + for (size_t i = 0; i < n; i++) out[i] = t.q[i] * t.s[i / groupSize]; +} + +// View into one layer's slice of a multi-layer quantized tensor (no copy). +QuantizedTensor layerView(const QuantizedTensor &t, size_t elemOffset, int groupSize) { + QuantizedTensor v; + v.q = t.q + elemOffset; + v.s = t.s + elemOffset / groupSize; + return v; +} + +void quantize(QuantizedTensor &t, const float *x, size_t n, int groupSize) { + size_t nGroups = n / groupSize; + for (size_t g = 0; g < nGroups; g++) { + float wmax = 0.0f; + for (int i = 0; i < groupSize; i++) { + float v = fabsf(x[g * groupSize + i]); + if (v > wmax) wmax = v; + } + float scale = wmax / 127.0f; + t.s[g] = scale; + for (int i = 0; i < groupSize; i++) { + float v = x[g * groupSize + i] / (scale == 0 ? 1.0f : scale); + t.q[g * groupSize + i] = (int8_t)roundf(v); + } + } +} + +// out(d) = W(d,n) @ x(n), quantized weights, fp32 activations (quantizes x on the fly). +void matmulQ(float *out, const float *x, const QuantizedTensor &w, int n, int d, int groupSize) { + QuantizedTensor xq; + xq.q = (int8_t *)alloca(n); + xq.s = (float *)alloca((n / groupSize) * sizeof(float)); + quantize(xq, x, n, groupSize); + + for (int i = 0; i < d; i++) { + float val = 0.0f; + int32_t ival = 0; + int in = i * n; + for (int j = 0; j <= n - groupSize; j += groupSize) { + for (int k = 0; k < groupSize; k++) ival += (int32_t)xq.q[j + k] * (int32_t)w.q[in + j + k]; + val += ((float)ival) * w.s[(in + j) / groupSize] * xq.s[j / groupSize]; + ival = 0; + } + out[i] = val; + } +} + +void matmulF(float *out, const float *x, const float *w, int n, int d) { + for (int i = 0; i < d; i++) { + float val = 0.0f; + const float *row = w + i * n; + for (int j = 0; j < n; j++) val += row[j] * x[j]; + out[i] = val; + } +} + +void rmsnorm(float *out, const float *x, const float *weight, int size) { + float ss = 0.0f; + for (int i = 0; i < size; i++) ss += x[i] * x[i]; + ss = 1.0f / sqrtf(ss / size + 1e-5f); + for (int i = 0; i < size; i++) out[i] = weight[i] * (x[i] * ss); +} + +void softmax(float *x, int size) { + float maxv = x[0]; + for (int i = 1; i < size; i++) + if (x[i] > maxv) maxv = x[i]; + float sum = 0.0f; + for (int i = 0; i < size; i++) { + x[i] = expf(x[i] - maxv); + sum += x[i]; + } + for (int i = 0; i < size; i++) x[i] /= sum; +} + +} // namespace + +struct LLMEngine::Impl { + bool quantized = false; + bool legacy = false; + int groupSize = 0; + bool sharedClassifier = true; + + // fp32 weights (used when !quantized) + float *tokEmbF = nullptr; + float *rmsAttW = nullptr, *rmsFfnW = nullptr, *rmsFinalW = nullptr; + float *wqF = nullptr, *wkF = nullptr, *wvF = nullptr, *woF = nullptr; + float *w1F = nullptr, *w2F = nullptr, *w3F = nullptr; + float *wclsF = nullptr; + + // quantized weights (used when quantized) + QuantizedTensor tokEmbQ; + QuantizedTensor wqQ, wkQ, wvQ, woQ; + QuantizedTensor w1Q, w2Q, w3Q; + QuantizedTensor wclsQ; + float *tokEmbDequant = nullptr; // dequantized embedding table (small, kept fp32 for lookup) + + // run state + float *x = nullptr, *xb = nullptr, *xb2 = nullptr; + float *hb = nullptr, *hb2 = nullptr; + float *q = nullptr, *att = nullptr, *logits = nullptr; + float *keyCache = nullptr, *valCache = nullptr; + + // tokenizer + std::unique_ptr vocab; + std::unique_ptr vocabScores; + int vocabSize = 0; + std::vector sortedVocab; // built once after vocab loads, for BPE merge lookups + + ~Impl() { freeAll(); } + + void freeAll() { + auto f = [](void *p) { + if (p) heap_caps_free(p); + }; + f(tokEmbF); + f(rmsAttW); + f(rmsFfnW); + f(rmsFinalW); + f(wqF); + f(wkF); + f(wvF); + f(woF); + f(w1F); + f(w2F); + f(w3F); + // wclsF aliases tokEmbF when the checkpoint has a shared classifier + // (see load()) - freeing both would double-free the same allocation. + if (wclsF != tokEmbF) f(wclsF); + f(tokEmbQ.q); + f(tokEmbQ.s); + f(wqQ.q); + f(wqQ.s); + f(wkQ.q); + f(wkQ.s); + f(wvQ.q); + f(wvQ.s); + f(woQ.q); + f(woQ.s); + f(w1Q.q); + f(w1Q.s); + f(w2Q.q); + f(w2Q.s); + f(w3Q.q); + f(w3Q.s); + f(wclsQ.q); + f(wclsQ.s); + f(tokEmbDequant); + f(x); + f(xb); + f(xb2); + f(hb); + f(hb2); + f(q); + f(att); + f(logits); + f(keyCache); + f(valCache); + if (vocab) { + for (int i = 0; i < vocabSize; i++) + if (vocab[i]) free(vocab[i]); + } + } +}; + +LLMEngine::LLMEngine() : impl(std::make_unique()) {} +LLMEngine::~LLMEngine() = default; + +LLMLoadError LLMEngine::load( + FS &fs, const String &checkpointPath, const String &tokenizerPath, bool overrideSafetyChecks +) { + unload(); + impl = std::make_unique(); + + if (!fs.exists(checkpointPath)) return LLMLoadError::CheckpointNotFound; + if (!fs.exists(tokenizerPath)) return LLMLoadError::TokenizerNotFound; + + File cf = fs.open(checkpointPath, FILE_READ); + if (!cf) return LLMLoadError::CheckpointNotFound; + + // export.py's legacy_export() (--version 0, the original karpathy/tinyllamas + // .bin files) writes NO magic/header at all - the file starts directly + // with the 7-int Config struct. version1_export()/version2_export() both + // start with a 4-byte "ak42" magic. We tell them apart by trying the + // magic first and falling back to legacy parsing if it doesn't match, + // rather than rejecting anything header-less outright. + uint32_t magic = 0; + cf.read((uint8_t *)&magic, 4); + bool legacy = (magic != kMagic); + + int32_t version = 0; + uint8_t sharedClassifierByte = 1; + int32_t groupSize = 0; + + if (legacy) { + // Legacy format signals shared_classifier via the sign of vocab_size + // instead of an explicit byte, and has no version field - the Config + // struct starts at byte 0, so back up over the 4 bytes we already + // spec-ulatively read as "magic". + cf.seek(0, SeekSet); + LLMConfig c{}; + cf.read((uint8_t *)&c, sizeof(LLMConfig)); + sharedClassifierByte = (c.vocab_size > 0) ? 1 : 0; + c.vocab_size = c.vocab_size > 0 ? c.vocab_size : -c.vocab_size; + cfg = c; + } else { + // export.py's model_export(): version 1 is the header-having fp32 + // format, version 2 is header-having int8/Q8_0. (There is no + // magic-having version 0 - that version number is the legacy format + // above, which never writes a magic at all.) + cf.read((uint8_t *)&version, 4); + if (version != 1 && version != 2) { + cf.close(); + return LLMLoadError::BadMagicOrVersion; + } + + // Unlike the legacy format, v1/v2 write vocab_size as a plain + // positive int and signal shared_classifier via an explicit byte + // below - no sign trick to undo here. + LLMConfig c{}; + cf.read((uint8_t *)&c, sizeof(LLMConfig)); + cfg = c; + + // Both v1 and v2 write a shared_classifier byte here; only v2 + // additionally writes a group_size int before the header is + // zero-padded out to 256 bytes. + cf.read((uint8_t *)&sharedClassifierByte, 1); + if (version == 2) cf.read((uint8_t *)&groupSize, 4); + cf.seek(256, SeekSet); + } + + impl->legacy = legacy; + impl->quantized = (version == 2); + impl->groupSize = groupSize; + impl->sharedClassifier = sharedClassifierByte != 0; + + int dim = cfg.dim, hidden = cfg.hidden_dim, layers = cfg.n_layers; + int heads = cfg.n_heads, kvHeads = cfg.n_kv_heads, vocab = cfg.vocab_size, seqLen = cfg.seq_len; + int headSize = dim / heads; + int kvDim = (dim * kvHeads) / heads; + + // quantize_q80 flattens each whole weight tensor into consecutive + // group_size-sized chunks with no regard for row boundaries. Our matmul + // (like reference run.c's matmul()) only sums whole groups within each + // row's own local range, so it silently drops/misaligns data whenever a + // row length isn't itself a multiple of group_size. dim, kvDim, and + // hidden are all used as row lengths (n) somewhere in the forward pass, + // so all three must divide evenly or generation will be garbage. + if (!overrideSafetyChecks && impl->quantized && + (dim % groupSize != 0 || kvDim % groupSize != 0 || hidden % groupSize != 0)) { + cf.close(); + return LLMLoadError::IncompatibleGroupSize; + } + + // Rough PSRAM budget guard - refuse loads that clearly won't fit, unless + // the caller has already warned the user and wants to try anyway. + size_t psramFree = heap_caps_get_free_size(MALLOC_CAP_SPIRAM); + size_t approxBytes = (size_t)vocab * dim * (impl->quantized ? 1 : 4) + + (size_t)layers * dim * dim * 4 * (impl->quantized ? 1 : 4) + + (size_t)layers * dim * hidden * 3 * (impl->quantized ? 1 : 4); + if (!overrideSafetyChecks && approxBytes > psramFree * 0.85) { + cf.close(); + return LLMLoadError::ConfigTooLarge; + } + + if (!impl->quantized && impl->legacy) { + // Matches legacy_export()'s exact write order: embedding table + // first, then every layer's attention weights/norms, then FFN + // weights/norm, then the final norm, then two freq_cis tables + // (precomputed RoPE cos/sin - unused here since we compute RoPE on + // the fly) that must be skipped over, then an optional classifier. + impl->tokEmbF = readF32Vec(cf, (size_t)vocab * dim); + impl->rmsAttW = readF32Vec(cf, (size_t)layers * dim); + impl->wqF = readF32Vec(cf, (size_t)layers * dim * (heads * headSize)); + impl->wkF = readF32Vec(cf, (size_t)layers * dim * (kvHeads * headSize)); + impl->wvF = readF32Vec(cf, (size_t)layers * dim * (kvHeads * headSize)); + impl->woF = readF32Vec(cf, (size_t)layers * (heads * headSize) * dim); + impl->rmsFfnW = readF32Vec(cf, (size_t)layers * dim); + impl->w1F = readF32Vec(cf, (size_t)layers * dim * hidden); + impl->w2F = readF32Vec(cf, (size_t)layers * hidden * dim); + impl->w3F = readF32Vec(cf, (size_t)layers * dim * hidden); + impl->rmsFinalW = readF32Vec(cf, dim); + cf.seek((size_t)seqLen * (headSize / 2) * sizeof(float) * 2, SeekCur); // freqs_cos + freqs_sin + impl->wclsF = impl->sharedClassifier ? impl->tokEmbF : readF32Vec(cf, (size_t)vocab * dim); + } else if (!impl->quantized) { + // Matches version1_export()'s exact write order: norms first, then + // the embedding table, then every layer's attention/FFN weights, then + // optionally an unshared classifier. There is no freq_cis table in + // this header-having format - RoPE is computed on the fly instead. + impl->rmsAttW = readF32Vec(cf, (size_t)layers * dim); + impl->rmsFfnW = readF32Vec(cf, (size_t)layers * dim); + impl->rmsFinalW = readF32Vec(cf, dim); + impl->tokEmbF = readF32Vec(cf, (size_t)vocab * dim); + impl->wqF = readF32Vec(cf, (size_t)layers * dim * (heads * headSize)); + impl->wkF = readF32Vec(cf, (size_t)layers * dim * (kvHeads * headSize)); + impl->wvF = readF32Vec(cf, (size_t)layers * dim * (kvHeads * headSize)); + impl->woF = readF32Vec(cf, (size_t)layers * (heads * headSize) * dim); + impl->w1F = readF32Vec(cf, (size_t)layers * dim * hidden); + impl->w2F = readF32Vec(cf, (size_t)layers * hidden * dim); + impl->w3F = readF32Vec(cf, (size_t)layers * dim * hidden); + impl->wclsF = impl->sharedClassifier ? impl->tokEmbF : readF32Vec(cf, (size_t)vocab * dim); + } else { + int gs = groupSize; + impl->rmsAttW = readF32Vec(cf, (size_t)layers * dim); + impl->rmsFfnW = readF32Vec(cf, (size_t)layers * dim); + impl->rmsFinalW = readF32Vec(cf, dim); + + readQuantized(cf, impl->tokEmbQ, (size_t)vocab * dim, 1, gs); + readQuantized(cf, impl->wqQ, (size_t)dim * (heads * headSize), layers, gs); + readQuantized(cf, impl->wkQ, (size_t)dim * (kvHeads * headSize), layers, gs); + readQuantized(cf, impl->wvQ, (size_t)dim * (kvHeads * headSize), layers, gs); + readQuantized(cf, impl->woQ, (size_t)(heads * headSize) * dim, layers, gs); + readQuantized(cf, impl->w1Q, (size_t)dim * hidden, layers, gs); + readQuantized(cf, impl->w2Q, (size_t)hidden * dim, layers, gs); + readQuantized(cf, impl->w3Q, (size_t)dim * hidden, layers, gs); + if (!impl->sharedClassifier) readQuantized(cf, impl->wclsQ, (size_t)vocab * dim, 1, gs); + + impl->tokEmbDequant = (float *)psram_alloc((size_t)vocab * dim * sizeof(float)); + if (impl->tokEmbDequant) dequantize(impl->tokEmbQ, impl->tokEmbDequant, (size_t)vocab * dim, gs); + } + cf.close(); + + // Tokenizer: llama2.c tokenizer.bin format = max_token_length(int32) then + // per-token: score(float) len(int32) bytes[len] + File tf = fs.open(tokenizerPath, FILE_READ); + if (!tf) return LLMLoadError::TokenizerNotFound; + int32_t maxTokLen = 0; + tf.read((uint8_t *)&maxTokLen, 4); + impl->vocabSize = vocab; + impl->vocab = std::make_unique(vocab); + impl->vocabScores = std::make_unique(vocab); + for (int i = 0; i < vocab; i++) { + tf.read((uint8_t *)&impl->vocabScores[i], 4); + int32_t len = 0; + tf.read((uint8_t *)&len, 4); + char *s = (char *)malloc(len + 1); + tf.read((uint8_t *)s, len); + s[len] = '\0'; + impl->vocab[i] = s; + } + tf.close(); + + impl->sortedVocab.resize(vocab); + for (int i = 0; i < vocab; i++) impl->sortedVocab[i] = {impl->vocab[i], i}; + std::sort( + impl->sortedVocab.begin(), impl->sortedVocab.end(), [](const TokenIndex &a, const TokenIndex &b) { + return strcmp(a.str, b.str) < 0; + } + ); + + // run-state buffers (small, fp32, PSRAM) + impl->x = (float *)psram_alloc(dim * sizeof(float)); + impl->xb = (float *)psram_alloc(dim * sizeof(float)); + impl->xb2 = (float *)psram_alloc(dim * sizeof(float)); + impl->hb = (float *)psram_alloc(hidden * sizeof(float)); + impl->hb2 = (float *)psram_alloc(hidden * sizeof(float)); + impl->q = (float *)psram_alloc(dim * sizeof(float)); + impl->att = (float *)psram_alloc(heads * seqLen * sizeof(float)); + impl->logits = (float *)psram_alloc(vocab * sizeof(float)); + impl->keyCache = (float *)psram_alloc((size_t)layers * seqLen * kvDim * sizeof(float)); + impl->valCache = (float *)psram_alloc((size_t)layers * seqLen * kvDim * sizeof(float)); + + if (!impl->x || !impl->xb || !impl->xb2 || !impl->hb || !impl->hb2 || !impl->q || !impl->att || + !impl->logits || !impl->keyCache || !impl->valCache) { + return LLMLoadError::OutOfMemory; + } + + loaded = true; + return LLMLoadError::None; +} + +void LLMEngine::unload() { + impl = std::make_unique(); + loaded = false; +} + +namespace { + +int sampleArgmax(const float *probabilities, int n) { + int best = 0; + for (int i = 1; i < n; i++) + if (probabilities[i] > probabilities[best]) best = i; + return best; +} + +int sampleMult(const float *probabilities, int n, float coin) { + float cdf = 0.0f; + for (int i = 0; i < n; i++) { + cdf += probabilities[i]; + if (coin < cdf) return i; + } + return n - 1; // rounding fallback +} + +// Faithful port of run.c's sample_topp(): nucleus sampling over the smallest +// set of tokens whose cumulative probability exceeds topP, so the tail of +// very-low-probability tokens can never be picked. +int sampleTopp(const float *probabilities, int n, float topP, float coin) { + std::vector> probIndex; + probIndex.reserve(n); + float cutoff = (1.0f - topP) / (n - 1); + for (int i = 0; i < n; i++) + if (probabilities[i] >= cutoff) probIndex.push_back({probabilities[i], i}); + std::sort(probIndex.begin(), probIndex.end(), [](auto &a, auto &b) { return a.first > b.first; }); + + float cumulative = 0.0f; + int lastIdx = (int)probIndex.size() - 1; + for (size_t i = 0; i < probIndex.size(); i++) { + cumulative += probIndex[i].first; + if (cumulative > topP) { + lastIdx = (int)i; + break; + } + } + + float r = coin * cumulative; + float cdf = 0.0f; + for (int i = 0; i <= lastIdx; i++) { + cdf += probIndex[i].first; + if (r < cdf) return probIndex[i].second; + } + return probIndex[lastIdx].second; +} + +// xorshift RNG, identical to run.c's random_u32/random_f32 - lets a fixed +// seed reproduce the same output for the same prompt, which esp_random() +// (a hardware TRNG) can't do. +uint32_t randomU32(uint64_t &state) { + state ^= state >> 12; + state ^= state << 25; + state ^= state >> 27; + return (uint32_t)((state * 0x2545F4914F6CDD1Dull) >> 32); +} +float randomF32(uint64_t &state) { return (randomU32(state) >> 8) / 16777216.0f; } + +int sampleToken(float *logits, int n, float temperature, float topP, uint64_t &rngState) { + if (temperature <= 0.0f) return sampleArgmax(logits, n); + for (int i = 0; i < n; i++) logits[i] /= temperature; + softmax(logits, n); + float coin = randomF32(rngState); + if (topP <= 0.0f || topP >= 1.0f) return sampleMult(logits, n, coin); + return sampleTopp(logits, n, topP, coin); +} + +// Not part of upstream llama2.c. Downweights logits of recently-used tokens +// so tiny models are less likely to loop on the same phrase - the standard +// HF-transformers-style repetition penalty (divide positive logits, multiply +// negative ones, by the penalty). +void applyRepetitionPenalty(float *logits, const std::vector &history, float penalty) { + if (penalty == 1.0f) return; + for (int id : history) { + if (logits[id] > 0) logits[id] /= penalty; + else logits[id] *= penalty; + } +} + +// llama2.c's vocab stores raw bytes that have no printable token as +// "<0xXX>" fallback pieces (see run.c's decode()). Those must be turned back +// into the literal byte, or streamed output looks like hex garbage. +String decodePiece(const char *piece) { + unsigned int byteVal; + if (piece[0] == '<' && sscanf(piece, "<0x%02X>", &byteVal) == 1) { + char c = (char)byteVal; + return String(c); + } + return String(piece); +} + +// Faithful port of run.c's encode(): seed one token per UTF-8 codepoint +// (falling back to individual bytes, offset +3 for the // +// slots, when a codepoint has no standalone vocab entry) - prefixed with BOS +// and a "dummy prefix" space token, same as run.c - then repeatedly merge +// whichever adjacent pair scores highest in vocab_scores until no merge is +// left. Critically, BOS and the dummy prefix must be seeded *before* the +// merge loop runs (as run.c does) so they can themselves take part in a +// merge, e.g. the dummy-prefix space merging with a leading "<" into one +// token - appending them afterwards, outside the merge loop, silently +// produces a different, longer token sequence than run.c for the same text. +// This is the actual BPE algorithm the tokenizer/model were trained with - +// a naive longest-match-at-each-position tokenizer (what an earlier version +// of this function did) produces a different token sequence for the same +// text, which is out-of-distribution for the model and comes out as +// unrelated, garbled output even though the weights and prompt text are +// both correct. +void encodeBpe( + const String &text, char **vocab, const float *vocabScores, const std::vector &sortedVocab, + int *outIds, int &outCount +) { + outCount = 0; + outIds[outCount++] = 1; // BOS + if (text.length() == 0) return; + outIds[outCount++] = strLookup(" ", sortedVocab); // dummy prefix + + char strBuffer[5]; // up to 4 UTF-8 bytes + null terminator + size_t strLen = 0; + const char *text_c = text.c_str(); + + for (const char *c = text_c; *c != '\0'; c++) { + // 0x80 continuation bytes start a fresh codepoint's byte run when + // absent - see run.c's encode() for the same bit-trick explanation. + if ((*c & 0xC0) != 0x80) strLen = 0; + strBuffer[strLen++] = *c; + strBuffer[strLen] = '\0'; + + if ((unsigned char)(*(c + 1)) != 0 && (*(c + 1) & 0xC0) == 0x80 && strLen < 4) continue; + + int id = strLookup(strBuffer, sortedVocab); + if (id != -1) { + outIds[outCount++] = id; + } else { + for (size_t i = 0; i < strLen; i++) outIds[outCount++] = (unsigned char)strBuffer[i] + 3; + } + strLen = 0; + } + + char mergeBuf[256]; + for (;;) { + float bestScore = -1e10f; + int bestId = -1, bestIdx = -1; + for (int i = 0; i < outCount - 1; i++) { + snprintf(mergeBuf, sizeof(mergeBuf), "%s%s", vocab[outIds[i]], vocab[outIds[i + 1]]); + int id = strLookup(mergeBuf, sortedVocab); + if (id != -1 && vocabScores[id] > bestScore) { + bestScore = vocabScores[id]; + bestId = id; + bestIdx = i; + } + } + if (bestIdx == -1) break; + outIds[bestIdx] = bestId; + for (int i = bestIdx + 1; i < outCount - 1; i++) outIds[i] = outIds[i + 1]; + outCount--; + } +} +} // namespace + +void LLMEngine::generate( + const String &prompt, int maxTokens, const GenerationParams ¶ms, const TokenCallback &onToken +) { + if (!loaded) return; + Impl *m = impl.get(); + uint64_t rngState = + params.seed != 0 ? (uint64_t)params.seed : (((uint64_t)esp_random() << 32) | esp_random()); + std::vector history; + int dim = cfg.dim, hidden = cfg.hidden_dim, layers = cfg.n_layers; + int heads = cfg.n_heads, kvHeads = cfg.n_kv_heads, vocab = cfg.vocab_size, seqLen = cfg.seq_len; + int headSize = dim / heads; + int kvDim = (dim * kvHeads) / heads; + int kvMul = heads / kvHeads; + int gs = m->groupSize; + + // Chat-finetuned checkpoints expect the exact ": ...\n: ..." + // shape they were trained on; wrap the raw prompt in it here rather than + // pushing that formatting onto every caller. + String effectivePrompt = + params.chatTemplateEnabled ? (params.userTag + prompt + "\n" + params.botTag) : prompt; + + // encodeBpe() seeds BOS + the dummy-prefix space token itself, ahead of + // the merge loop, so they can take part in merges exactly like run.c's + // encode() does. + int *promptIds = (int *)alloca(sizeof(int) * (effectivePrompt.length() + 2)); + int nPrompt = 0; + encodeBpe(effectivePrompt, m->vocab.get(), m->vocabScores.get(), m->sortedVocab, promptIds, nPrompt); + + int steps = maxTokens < seqLen ? maxTokens : seqLen; + int token = promptIds[0]; + // Rolling tail of recently-emitted text, only used when chatTemplateEnabled + // - lets us notice the model re-emitting userTag (drifting into a fake new + // turn) and stop right there instead of streaming it to the UI. + String tailBuffer; + + for (int pos = 0; pos < steps; pos++) { + // --- forward pass for `token` at position `pos` --- + memcpy( + m->x, (m->quantized ? m->tokEmbDequant : m->tokEmbF) + (size_t)token * dim, dim * sizeof(float) + ); + + for (int l = 0; l < layers; l++) { + rmsnorm(m->xb, m->x, m->rmsAttW + l * dim, dim); + + float *kRow = m->keyCache + ((size_t)l * seqLen + pos) * kvDim; + float *vRow = m->valCache + ((size_t)l * seqLen + pos) * kvDim; + + if (!m->quantized) { + matmulF(m->q, m->xb, m->wqF + (size_t)l * dim * dim, dim, dim); + matmulF(kRow, m->xb, m->wkF + (size_t)l * dim * kvDim, dim, kvDim); + matmulF(vRow, m->xb, m->wvF + (size_t)l * dim * kvDim, dim, kvDim); + } else { + matmulQ(m->q, m->xb, layerView(m->wqQ, (size_t)l * dim * dim, gs), dim, dim, gs); + matmulQ(kRow, m->xb, layerView(m->wkQ, (size_t)l * dim * kvDim, gs), dim, kvDim, gs); + matmulQ(vRow, m->xb, layerView(m->wvQ, (size_t)l * dim * kvDim, gs), dim, kvDim, gs); + } + + // RoPE rotation + for (int h = 0; h < heads; h++) { + for (int i = 0; i < headSize; i += 2) { + float freq = 1.0f / powf(10000.0f, (float)i / headSize); + float val = pos * freq; + float fcr = cosf(val), fci = sinf(val); + int base = h * headSize; + if (h < kvHeads) { + float v0 = kRow[base + i], v1 = kRow[base + i + 1]; + kRow[base + i] = v0 * fcr - v1 * fci; + kRow[base + i + 1] = v0 * fci + v1 * fcr; + } + float v0 = m->q[base + i], v1 = m->q[base + i + 1]; + m->q[base + i] = v0 * fcr - v1 * fci; + m->q[base + i + 1] = v0 * fci + v1 * fcr; + } + } + + for (int h = 0; h < heads; h++) { + float *qh = m->q + h * headSize; + float *attRow = m->att + h * seqLen; + for (int t = 0; t <= pos; t++) { + float *kt = m->keyCache + ((size_t)l * seqLen + t) * kvDim + (h / kvMul) * headSize; + float score = 0.0f; + for (int i = 0; i < headSize; i++) score += qh[i] * kt[i]; + attRow[t] = score / sqrtf((float)headSize); + } + softmax(attRow, pos + 1); + float *out = m->xb2 + h * headSize; + memset(out, 0, headSize * sizeof(float)); + for (int t = 0; t <= pos; t++) { + float *vt = m->valCache + ((size_t)l * seqLen + t) * kvDim + (h / kvMul) * headSize; + float a = attRow[t]; + for (int i = 0; i < headSize; i++) out[i] += a * vt[i]; + } + } + + if (!m->quantized) matmulF(m->xb, m->xb2, m->woF + (size_t)l * dim * dim, dim, dim); + else matmulQ(m->xb, m->xb2, layerView(m->woQ, (size_t)l * dim * dim, gs), dim, dim, gs); + + for (int i = 0; i < dim; i++) m->x[i] += m->xb[i]; + + rmsnorm(m->xb, m->x, m->rmsFfnW + l * dim, dim); + if (!m->quantized) { + matmulF(m->hb, m->xb, m->w1F + (size_t)l * dim * hidden, dim, hidden); + matmulF(m->hb2, m->xb, m->w3F + (size_t)l * dim * hidden, dim, hidden); + } else { + matmulQ(m->hb, m->xb, layerView(m->w1Q, (size_t)l * dim * hidden, gs), dim, hidden, gs); + matmulQ(m->hb2, m->xb, layerView(m->w3Q, (size_t)l * dim * hidden, gs), dim, hidden, gs); + } + for (int i = 0; i < hidden; i++) { + float v = m->hb[i]; + v *= 1.0f / (1.0f + expf(-v)); // SiLU + m->hb[i] = v * m->hb2[i]; + } + if (!m->quantized) matmulF(m->xb, m->hb, m->w2F + (size_t)l * hidden * dim, hidden, dim); + else matmulQ(m->xb, m->hb, layerView(m->w2Q, (size_t)l * hidden * dim, gs), hidden, dim, gs); + + for (int i = 0; i < dim; i++) m->x[i] += m->xb[i]; + } + + rmsnorm(m->x, m->x, m->rmsFinalW, dim); + if (!m->quantized) matmulF(m->logits, m->x, m->wclsF, dim, vocab); + else matmulQ(m->logits, m->x, m->sharedClassifier ? m->tokEmbQ : m->wclsQ, dim, vocab, gs); + + int nextToken; + if (pos + 1 < nPrompt) { + nextToken = promptIds[pos + 1]; + } else { + applyRepetitionPenalty(m->logits, history, params.repetitionPenalty); + nextToken = sampleToken(m->logits, vocab, params.temperature, params.topP, rngState); + } + + // run.c's generate() loop prints decode(token, next) every iteration + // - i.e. it always prints the *next* token's text, not the current + // one - so its prompt-echo naturally ends and real generation begins + // the moment `next` stops being forced from the prompt (pos >= + // nPrompt - 1). We print `token` (already-advanced from the previous + // iteration) instead of `next`, which is the same stream one loop + // tick later - so our matching cutover is pos >= nPrompt, not + // nPrompt - 1. Gating on nPrompt - 1 here would print `token` while + // it's still the final prompt token itself (e.g. the last subword of + // the user's own prompt text), leaking one leftover fragment of the + // prompt onto the front of the displayed reply. + if (pos >= nPrompt) { + String piece = decodePiece(m->vocab[token]); + if (params.chatTemplateEnabled && params.userTag.length() > 0) { + String combined = tailBuffer + piece; + int tagPos = combined.indexOf(params.userTag); + if (tagPos != -1) { + // The model started echoing userTag - it's hallucinating a + // new turn. Emit only what came before the tag, then stop. + int emitLen = tagPos - (int)tailBuffer.length(); + if (emitLen > 0) onToken(piece.substring(0, emitLen)); + return; + } + tailBuffer = combined; + int maxKeep = (int)params.userTag.length() * 2 + 8; + if ((int)tailBuffer.length() > maxKeep) + tailBuffer = tailBuffer.substring(tailBuffer.length() - maxKeep); + } + if (!onToken(piece)) return; // cancelled + } + // EOS (id 2, by llama2.c/sentencepiece convention) marks a natural + // end of generation - stop instead of rambling on past it. Only + // meaningful once we're actually generating, not echoing the prompt. + if (pos >= nPrompt - 1 && nextToken == 2) return; + // BOS (id 1) delimits sequences, matching run.c's generate() loop - a model that + // starts predicting a fresh BOS mid-generation is drifting into a new/unrelated + // sequence, so stop here too rather than rambling into a hallucinated new turn. + if (pos >= nPrompt - 1 && nextToken == 1) return; + history.push_back(token); + if (history.size() > 64) history.erase(history.begin()); + token = nextToken; + } +} + +} // namespace bruce_llm diff --git a/src/modules/others/BruceLM/llm_engine.h b/src/modules/others/BruceLM/llm_engine.h new file mode 100644 index 0000000000..8ca2edbfb1 --- /dev/null +++ b/src/modules/others/BruceLM/llm_engine.h @@ -0,0 +1,109 @@ +/** + * @file llm_engine.h + * @brief Minimal on-device inference engine for llama2.c-format checkpoints. + * + * Supports any model exported by karpathy/llama2.c's export.py in either + * version 1 (fp32) or version 2 (int8 / Q8_0 symmetric quantized) format. + * (version 0 is the legacy header-less format and is not supported - re-export.) + * Model dimensions (dim, n_layers, n_heads, n_kv_heads, vocab_size, seq_len) + * are read from the checkpoint header at load time - nothing is hardcoded, + * so any model that fits in available PSRAM can be loaded. + */ +#pragma once + +#include +#include +#include + +namespace bruce_llm { + +// Mirrors llama2.c's Config struct layout (7 x int32, little-endian). +struct LLMConfig { + int32_t dim; + int32_t hidden_dim; + int32_t n_layers; + int32_t n_heads; + int32_t n_kv_heads; + int32_t vocab_size; // negative => unshared classifier weights (fp32 export quirk) + int32_t seq_len; +}; + +enum class LLMLoadError { + None, + CheckpointNotFound, + TokenizerNotFound, + BadMagicOrVersion, + ConfigTooLarge, // wouldn't fit in available PSRAM + OutOfMemory, + // Quantized (v2) export flattens the whole tensor into group_size chunks + // without regard to row boundaries; our matmul (like reference run.c's) + // only handles this correctly when every row length is itself a whole + // number of groups. Rejected here rather than silently corrupting output. + IncompatibleGroupSize, +}; + +// Called once per generated token. Return false to cancel generation. +using TokenCallback = std::function; + +// Sampling knobs. temperature/topP/seed are faithful ports of run.c's own +// Sampler (temperature, nucleus/top-p sampling, seeded xorshift RNG). +// repetitionPenalty is NOT part of upstream llama2.c - it's a common, +// cheap addition (downweights recently-used tokens' logits) that measurably +// helps small models avoid repetition loops. +struct GenerationParams { + float temperature = 0.8f; // 0 = greedy/deterministic, higher = more random + float topP = 0.9f; // nucleus sampling threshold; >=1.0 or <=0 disables it + float repetitionPenalty = 1.0f; // 1.0 = disabled, >1.0 discourages repeating recent tokens + uint32_t seed = 0; // 0 = reseed from hardware RNG each call (non-reproducible); + // nonzero = deterministic output for the same prompt + + // Chat-finetuned checkpoints (unlike plain story models) are trained on + // ": ...\n: ..." pairs - fed raw text with no cue for whose + // turn it is, they'll hallucinate a whole fake exchange (including their + // own userTag) before ever answering. When enabled, the prompt is + // wrapped as userTag + prompt + "\n" + botTag before encoding, and + // generation stops the moment userTag reappears in the output (the model + // drifting into a new fake turn) instead of rambling past it. + bool chatTemplateEnabled = true; + String userTag = ": "; + String botTag = ":"; +}; + +class LLMEngine { +public: + LLMEngine(); + ~LLMEngine(); + + // Loads a checkpoint (.bin, llama2.c format) + matching tokenizer (.bin) from fs. + // If overrideSafetyChecks is true, loads anyway even when the model looks + // too large for available PSRAM or the quantized group_size doesn't + // evenly divide every row length - the caller has warned the user that + // this may fail or produce garbled output, instead of refusing outright. + LLMLoadError load( + FS &fs, + const String &checkpointPath, + const String &tokenizerPath, + bool overrideSafetyChecks = false + ); + + void unload(); + bool isLoaded() const { return loaded; } + + const LLMConfig &config() const { return cfg; } + + // Runs generation from `prompt`, streaming pieces to `onToken`, up to maxTokens + // (or the model's seq_len, whichever is smaller). Blocking call - designed + // to be invoked from the UI loop with onToken also servicing UI/cancel checks. + void generate( + const String &prompt, int maxTokens, const GenerationParams ¶ms, + const TokenCallback &onToken + ); + +private: + struct Impl; + std::unique_ptr impl; + LLMConfig cfg{}; + bool loaded = false; +}; + +} // namespace bruce_llm