feat(dictate): use local socket-activated Parakeet

2026-08-26 12:38:45 -04:00
parent ea74dbc6a7
commit e234468a73
34 changed files with 6702 additions and 527 deletions
@@ -0,0 +1,183 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+MODEL_REVISION="541d1f99c6b0c3cd0b11a95167540bb8edefd82b"
+MODEL_SHA256="3cbdc85877e668ca7b82d0d56770eb1fac76691f55d6b97545e8d61ca588d10d"
+MODEL_FILENAME="parakeet-tdt-0.6b-v3.nemo"
+
+SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
+REPOSITORY_ROOT="$(cd -- "$SCRIPT_DIR/../.." && pwd)"
+INSTALL_ROOT="$HOME/.local/share/pi-dictate"
+UNIT_DIR="${XDG_CONFIG_HOME:-$HOME/.config}/systemd/user"
+MODEL_DIR="$INSTALL_ROOT/models/parakeet-tdt-0.6b-v3"
+MODEL_PATH="$MODEL_DIR/$MODEL_FILENAME"
+DICTIONARY_SOURCE="$SCRIPT_DIR/coding-terms.txt"
+DICTIONARY_PATH="$INSTALL_ROOT/config/coding-terms.txt"
+DRY_RUN=false
+MODEL_SOURCE=""
+ACTION=install
+
+usage() {
+  cat <<'EOF'
+Usage: install.sh [--dry-run] [--model-source PATH]
+       install.sh uninstall [--dry-run]
+
+Installs the pinned service environment, model snapshot, and systemd user units.
+Without --model-source, installation downloads the exact pinned model revision.
+EOF
+}
+
+run() {
+  printf '+ '
+  printf '%q ' "$@"
+  printf '\n'
+  if ! $DRY_RUN; then
+    "$@"
+  fi
+}
+
+is_owned_unit() {
+  local destination="$1"
+  local source="$2"
+  [[ -L "$destination" && "$(readlink -f -- "$destination")" == "$(readlink -f -- "$source")" ]] ||
+    { [[ -f "$destination" ]] && grep -qxF '# Managed by pi-dictate install.sh' "$destination"; }
+}
+
+install_owned_unit() {
+  local source="$1"
+  local destination="$2"
+  if [[ -L "$destination" && "$(readlink -f -- "$destination")" == "$(readlink -f -- "$source")" ]]; then
+    run rm "$destination"
+  elif [[ -e "$destination" || -L "$destination" ]]; then
+    if [[ ! -f "$destination" ]] || ! grep -qxF '# Managed by pi-dictate install.sh' "$destination"; then
+      printf 'Refusing to replace unmanaged systemd unit: %s\n' "$destination" >&2
+      return 1
+    fi
+  fi
+  run install -m 0644 "$source" "$destination"
+}
+
+remove_owned_unit() {
+  local destination="$1"
+  local source="$2"
+  if [[ -L "$destination" && "$(readlink -f -- "$destination")" == "$(readlink -f -- "$source")" ]]; then
+    run rm "$destination"
+  elif [[ -f "$destination" ]] && grep -qxF '# Managed by pi-dictate install.sh' "$destination"; then
+    run rm "$destination"
+  elif [[ -e "$destination" || -L "$destination" ]]; then
+    printf 'Leaving unmanaged unit in place: %s\n' "$destination" >&2
+  fi
+}
+
+while (($#)); do
+  case "$1" in
+    uninstall) ACTION=uninstall ;;
+    --dry-run) DRY_RUN=true ;;
+    --model-source)
+      shift
+      (($#)) || { printf '%s\n' '--model-source requires a path' >&2; exit 2; }
+      MODEL_SOURCE="$1"
+      ;;
+    -h|--help) usage; exit 0 ;;
+    *) printf 'Unknown argument: %s\n' "$1" >&2; usage >&2; exit 2 ;;
+  esac
+  shift
+done
+
+socket_source="$REPOSITORY_ROOT/systemd/user/pi-dictate.socket"
+service_source="$REPOSITORY_ROOT/systemd/user/pi-dictate.service"
+socket_link="$UNIT_DIR/pi-dictate.socket"
+service_link="$UNIT_DIR/pi-dictate.service"
+
+if [[ "$ACTION" == uninstall ]]; then
+  for pair in "$socket_link:$socket_source" "$service_link:$service_source"; do
+    destination="${pair%%:*}"
+    source="${pair#*:}"
+    if [[ ( -e "$destination" || -L "$destination" ) ]] && ! is_owned_unit "$destination" "$source"; then
+      printf 'Refusing to disable or remove unmanaged systemd unit: %s\n' "$destination" >&2
+      exit 1
+    fi
+  done
+  run systemctl --user disable --now pi-dictate.socket
+  remove_owned_unit "$socket_link" "$socket_source"
+  remove_owned_unit "$service_link" "$service_source"
+  run systemctl --user daemon-reload
+  run rm -rf "$INSTALL_ROOT"
+  exit 0
+fi
+
+command -v uv >/dev/null || { printf '%s\n' 'uv is required' >&2; exit 1; }
+command -v systemctl >/dev/null || { printf '%s\n' 'systemctl is required' >&2; exit 1; }
+
+run mkdir -p "$INSTALL_ROOT" "$MODEL_DIR" "$(dirname "$DICTIONARY_PATH")" "$UNIT_DIR"
+run chmod 0700 "$INSTALL_ROOT" "$(dirname "$MODEL_DIR")" "$MODEL_DIR" "$(dirname "$DICTIONARY_PATH")"
+if $DRY_RUN; then
+  printf '+ UV_PROJECT_ENVIRONMENT=%q uv sync --frozen --no-dev --no-editable --project %q\n' "$INSTALL_ROOT/venv" "$SCRIPT_DIR"
+else
+  UV_PROJECT_ENVIRONMENT="$INSTALL_ROOT/venv" uv sync --frozen --no-dev --no-editable --project "$SCRIPT_DIR"
+fi
+
+MODEL_TEMP="$MODEL_DIR/.${MODEL_FILENAME}.install.$$"
+cleanup_model_temp() {
+  if ! $DRY_RUN; then
+    rm -f -- "$MODEL_TEMP"
+  fi
+}
+trap cleanup_model_temp EXIT
+
+if [[ -n "$MODEL_SOURCE" ]]; then
+  [[ -f "$MODEL_SOURCE" ]] || { printf 'Model source not found: %s\n' "$MODEL_SOURCE" >&2; exit 1; }
+  run cp --reflink=auto "$MODEL_SOURCE" "$MODEL_TEMP"
+else
+  if $DRY_RUN; then
+    printf '+ download nvidia/parakeet-tdt-0.6b-v3@%s to %q\n' "$MODEL_REVISION" "$MODEL_TEMP"
+  else
+    MODEL_REVISION="$MODEL_REVISION" MODEL_PATH="$MODEL_TEMP" \
+      "$INSTALL_ROOT/venv/bin/python" - <<'PY'
+import os
+import shutil
+from pathlib import Path
+from huggingface_hub import hf_hub_download
+
+source = hf_hub_download(
+    repo_id="nvidia/parakeet-tdt-0.6b-v3",
+    filename="parakeet-tdt-0.6b-v3.nemo",
+    revision=os.environ["MODEL_REVISION"],
+)
+destination = Path(os.environ["MODEL_PATH"])
+destination.parent.mkdir(parents=True, exist_ok=True)
+shutil.copyfile(source, destination)
+PY
+  fi
+fi
+
+if ! $DRY_RUN; then
+  actual_sha256="$(sha256sum "$MODEL_TEMP" | cut -d' ' -f1)"
+  if [[ "$actual_sha256" != "$MODEL_SHA256" ]]; then
+    rm -f "$MODEL_TEMP"
+    printf 'Model checksum mismatch: expected %s, got %s\n' "$MODEL_SHA256" "$actual_sha256" >&2
+    exit 1
+  fi
+fi
+run chmod 0400 "$MODEL_TEMP"
+run mv -f "$MODEL_TEMP" "$MODEL_PATH"
+trap - EXIT
+
+if [[ ! -e "$DICTIONARY_PATH" ]]; then
+  run cp "$DICTIONARY_SOURCE" "$DICTIONARY_PATH"
+  run chmod 0600 "$DICTIONARY_PATH"
+else
+  printf 'Keeping existing coding dictionary: %s\n' "$DICTIONARY_PATH"
+fi
+
+install_owned_unit "$socket_source" "$socket_link"
+install_owned_unit "$service_source" "$service_link"
+run systemctl --user daemon-reload
+run systemctl --user enable --now pi-dictate.socket
+
+if $DRY_RUN; then
+  printf 'Would install pi-dictate with model revision %s.\n' "$MODEL_REVISION"
+else
+  printf 'Installed pi-dictate with model revision %s.\n' "$MODEL_REVISION"
+  printf 'Check it with: systemctl --user status pi-dictate.socket pi-dictate.service\n'
+fi