Test a custom prompt against an already-trained, already-flashed ESP32-S3 model

run_prompt.sh

#!/usr/bin/env bash
# esp32-ai: test a custom prompt against an already-trained, already-flashed model.
# Does NOT retrain and does NOT rewrite the 15MB model partition.
#
# Usage:
#   ./run_custom_prompt.sh "Once there was a robot"
set -euo pipefail

if [ "$#" -lt 1 ]; then
  echo "Usage: $0 \"Your custom prompt\""
  exit 1
fi
PROMPT="$1"

REPO_DIR="esp32-ai"
if [ -f "$REPO_DIR/pyproject.toml" ]; then
  cd "$REPO_DIR"
elif [ ! -f "pyproject.toml" ]; then
  echo "Can't find the esp32-ai repo. Run this from the same folder as build_and_flash.sh, after it has cloned esp32-ai."
  exit 1
fi

if [ ! -f "data/bpe32768.json" ]; then
  echo "data/bpe32768.json not found. Run build_and_flash.sh at least once first (it trains the tokenizer)."
  exit 1
fi

# ---- Detect port ----
CANDIDATE_PORTS=(/dev/cu.usbmodem*)
if [ ! -e "${CANDIDATE_PORTS[0]}" ]; then
  echo "No /dev/cu.usbmodem* device found. Is the board plugged in?"
  exit 1
elif [ "${#CANDIDATE_PORTS[@]}" -eq 1 ]; then
  ESP32_PORT="${CANDIDATE_PORTS[0]}"
else
  select choice in "${CANDIDATE_PORTS[@]}"; do
    [ -n "$choice" ] && ESP32_PORT="$choice" && break
  done
fi
echo ">> Using port: $ESP32_PORT"

# ---- Tokenize the custom prompt with the same tokenizer gen_assets.py uses ----
echo ">> Tokenizing prompt: \"$PROMPT\""

TOKENIZE_SCRIPT="$(mktemp /tmp/tokenize_prompt_XXXX.py)"
cat > "$TOKENIZE_SCRIPT" << 'PYEOF'
import sys
from tokenizers import Tokenizer

tok = Tokenizer.from_file("data/bpe32768.json")
prompt = sys.argv[1]
ids = tok.encode(prompt).ids
print(f'prompt "{prompt}" -> {len(ids)} ids: {ids}')
print(f"round-trip: {tok.decode(ids)!r}")
print("static const int PROMPT_IDS[] = {" + ", ".join(map(str, ids)) + "};")
PYEOF

GEN_OUTPUT="$(uv run python "$TOKENIZE_SCRIPT" "$PROMPT")"
rm -f "$TOKENIZE_SCRIPT"
echo "$GEN_OUTPUT"

PROMPT_IDS_LINE="$(echo "$GEN_OUTPUT" | grep -o 'static const int PROMPT_IDS\[\] = {[^}]*};')"
if [ -z "$PROMPT_IDS_LINE" ]; then
  echo "!! Couldn't parse a PROMPT_IDS line out of the tokenizer output. Aborting."
  exit 1
fi
echo ">> New prompt line: $PROMPT_IDS_LINE"

# ---- Paste it into the sketch ----
SKETCH_FILE="firmware/esp32_llm/esp32_llm.ino"
sed -i.bak -E "s/static const int PROMPT_IDS\[\] = \{[^}]*\};/${PROMPT_IDS_LINE//\//\\/}/" "$SKETCH_FILE"

# ---- Rebuild and reflash just the app (model partition untouched) ----
FQBN='esp32:esp32:esp32s3:UploadSpeed=921600,USBMode=hwcdc,CDCOnBoot=default,UploadMode=default,CPUFreq=240,FlashMode=qio,FlashSize=16M,PartitionScheme=custom,PSRAM=opi,DebugLevel=info'

echo ">> Recompiling firmware with new prompt..."
arduino-cli compile \
  --fqbn "$FQBN" \
  --build-property compiler.optimization_flags=-O3 \
  --build-path /tmp/esp32-llm-uart \
  firmware/esp32_llm

echo ">> Flashing app (model partition is untouched)..."
arduino-cli upload -p "$ESP32_PORT" \
  --fqbn "$FQBN" \
  --input-dir /tmp/esp32-llm-uart firmware/esp32_llm

echo ">> Opening serial monitor. Press the board's RESET button now."
arduino-cli monitor -p "$ESP32_PORT" --config baudrate=115200