mirror of
https://github.com/mudler/LocalAI.git
synced 2026-08-06 21:32:59 -04:00
Compare commits
2 Commits
feat/nemo-
...
bot/issue-
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
70f1eb3e77 | ||
|
|
d2ed733ad0 |
@@ -304,9 +304,7 @@ React pages that want to filter the ModelSelector by capability import this symb
|
||||
|
||||
### 4. `docs/content/` (user-facing documentation)
|
||||
|
||||
A new capability deserves its own page under `docs/content/features/`, plus cross-links from related features. See the pattern used by `face-recognition.md` / `object-detection.md`.
|
||||
|
||||
Announcing it is the release's job, not this page's: the capability gets covered in the release blog post under `website/content/blog/`. See [preparing-a-release.md](preparing-a-release.md). `docs/content/whats-new.md` is only a pointer at the blog and GitHub Releases, so there is nothing to add there.
|
||||
A new capability deserves its own page under `docs/content/features/`, plus cross-links from related features and an entry in `docs/content/whats-new.md`. See the pattern used by `face-recognition.md` / `object-detection.md`.
|
||||
|
||||
## Path protection rules
|
||||
|
||||
@@ -336,7 +334,7 @@ When adding a new endpoint:
|
||||
- [ ] Swagger block on the handler: `@Summary`, `@Tags`, `@Param`, `@Success`, `@Router`
|
||||
- [ ] If new capability area (new swagger tag): entry in `instructionDefs` in `core/http/endpoints/localai/api_instructions.go` + test count bumped in `api_instructions_test.go`
|
||||
- [ ] If new `FLAG_*` usecase flag: matching `CAP_*` symbol exported from `core/http/react-ui/src/utils/capabilities.js`
|
||||
- [ ] `docs/content/features/<feature>.md` created; cross-links from related feature pages; capability covered in the release blog post (see [preparing-a-release.md](preparing-a-release.md))
|
||||
- [ ] `docs/content/features/<feature>.md` created; cross-links from related feature pages; entry in `docs/content/whats-new.md`
|
||||
|
||||
**Quality**
|
||||
- [ ] Error responses use `schema.ErrorResponse` format (or `echo.NewHTTPError` with a mapped gRPC status — see the `mapBackendError` helper in `core/http/endpoints/localai/images.go`)
|
||||
|
||||
@@ -16,7 +16,8 @@ side (`pkg/oci/cosignverify` plus the gallery YAML).
|
||||
per-arch manifest before checking signatures.
|
||||
- **Storage:** Signatures are written as OCI 1.1 referrers
|
||||
(`--registry-referrers-mode=oci-1-1`) in the new Sigstore bundle format
|
||||
(`--new-bundle-format`). No `:sha256-<hex>.sig` tag clutter.
|
||||
(current cosign releases do this by default; no `--new-bundle-format`
|
||||
flag). No `:sha256-<hex>.sig` tag clutter.
|
||||
- **Consumer:** `pkg/oci/cosignverify` discovers the bundle via the
|
||||
referrers API, hands it to `sigstore-go`, and verifies it against the
|
||||
policy declared in the gallery YAML (`Gallery.Verification`).
|
||||
@@ -33,15 +34,14 @@ to sign. The job needs:
|
||||
|
||||
- `permissions: { id-token: write, contents: read }` at the job level so
|
||||
the runner can exchange its GitHub OIDC token for a Fulcio cert.
|
||||
- `sigstore/cosign-installer@v3` step (the pinned cosign v2 release needs
|
||||
`--new-bundle-format` explicitly).
|
||||
- `sigstore/cosign-installer@v3` step (current cosign releases already
|
||||
default to the new bundle format).
|
||||
- After each `docker buildx imagetools create`, resolve the resulting
|
||||
list digest with `docker buildx imagetools inspect <tag> --format
|
||||
'{{.Manifest.Digest}}'` and sign:
|
||||
|
||||
```sh
|
||||
cosign sign --yes --recursive \
|
||||
--new-bundle-format \
|
||||
--registry-referrers-mode=oci-1-1 \
|
||||
"${REGISTRY_REPO}@${DIGEST}"
|
||||
```
|
||||
@@ -70,7 +70,7 @@ entry (`backend/index.yaml`):
|
||||
url: github:mudler/LocalAI/backend/index.yaml@master
|
||||
verification:
|
||||
issuer: "https://token.actions.githubusercontent.com"
|
||||
identity_regex: "^https://github\\.com/mudler/LocalAI/\\.github/workflows/backend_merge\\.yml@refs/(heads/master|tags/.+)$"
|
||||
identity_regex: "^https://github\\.com/mudler/LocalAI/\\.github/workflows/backend_merge\\.yml@refs/heads/master$"
|
||||
# Optional revocation cutoff; advance during incident response.
|
||||
# not_before: "2026-06-01T00:00:00Z"
|
||||
```
|
||||
|
||||
@@ -4,24 +4,6 @@ set -euo pipefail
|
||||
arch=${1:?target architecture is required}
|
||||
build_type=${2-}
|
||||
|
||||
# SYCL compiles the whole tree with icpx -fsycl, and icpx never finishes
|
||||
# ggml-cpu/arch/x86/repack.cpp at -march=sapphirerapids: the job sits on that one
|
||||
# translation unit until GitHub kills it at 6h. gcc builds the same file in
|
||||
# seconds, so only the SYCL images have to give up the CPU variant matrix.
|
||||
#
|
||||
# ROCm runs out of the same 6h budget for a different reason: volume, not a
|
||||
# stall. hipcc compiles ggml's HIP kernels once per entry in AMDGPU_TARGETS,
|
||||
# which is eleven architectures (gfx908 through gfx1201), and the CPU variant
|
||||
# matrix lands on top of that. The job built in 2h27m before it was added and
|
||||
# has been killed at exactly 6h00m on every run since, so no ROCm llama-cpp
|
||||
# image has been published since 2026-08-01.
|
||||
case "$build_type" in
|
||||
sycl*|hipblas*)
|
||||
echo llama-cpp-fallback
|
||||
exit 0
|
||||
;;
|
||||
esac
|
||||
|
||||
# GPU arm64 base images do not consistently provide the gcc-14 toolchain needed
|
||||
# to compile ggml's armv9.2 CPU variants. Keep their portable fallback until the
|
||||
# builder images can supply that compiler.
|
||||
|
||||
@@ -4,17 +4,6 @@ set -euo pipefail
|
||||
arch=${1:?target architecture is required}
|
||||
build_type=${2-}
|
||||
|
||||
# SYCL compiles the whole tree with icpx -fsycl, and icpx never finishes
|
||||
# ggml-cpu/arch/x86/repack.cpp at -march=sapphirerapids: the job sits on that one
|
||||
# translation unit until GitHub kills it at 6h. gcc builds the same file in
|
||||
# seconds, so only the SYCL images have to give up the CPU variant matrix.
|
||||
case "$build_type" in
|
||||
sycl*)
|
||||
echo turboquant-fallback
|
||||
exit 0
|
||||
;;
|
||||
esac
|
||||
|
||||
# GPU arm64 base images do not consistently provide the gcc-14 toolchain needed
|
||||
# to compile ggml's armv9.2 CPU variants. Keep their portable fallback until the
|
||||
# builder images can supply that compiler.
|
||||
|
||||
128
.github/backend-matrix.yml
vendored
128
.github/backend-matrix.yml
vendored
@@ -860,19 +860,6 @@ include:
|
||||
dockerfile: "./backend/Dockerfile.golang"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
- build-type: 'cublas'
|
||||
cuda-major-version: "12"
|
||||
cuda-minor-version: "8"
|
||||
platforms: 'linux/amd64'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-gpu-nvidia-cuda-12-nemo-speech-cpp'
|
||||
runs-on: 'ubuntu-latest'
|
||||
base-image: "ubuntu:24.04"
|
||||
skip-drivers: 'false'
|
||||
backend: "nemo-speech-cpp"
|
||||
dockerfile: "./backend/Dockerfile.golang"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
- build-type: 'cublas'
|
||||
cuda-major-version: "12"
|
||||
cuda-minor-version: "8"
|
||||
@@ -1924,19 +1911,6 @@ include:
|
||||
dockerfile: "./backend/Dockerfile.golang"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
- build-type: 'cublas'
|
||||
cuda-major-version: "13"
|
||||
cuda-minor-version: "0"
|
||||
platforms: 'linux/amd64'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-gpu-nvidia-cuda-13-nemo-speech-cpp'
|
||||
runs-on: 'ubuntu-latest'
|
||||
base-image: "ubuntu:24.04"
|
||||
skip-drivers: 'false'
|
||||
backend: "nemo-speech-cpp"
|
||||
dockerfile: "./backend/Dockerfile.golang"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
- build-type: 'cublas'
|
||||
cuda-major-version: "13"
|
||||
cuda-minor-version: "0"
|
||||
@@ -1989,24 +1963,6 @@ include:
|
||||
backend: "parakeet-cpp"
|
||||
dockerfile: "./backend/Dockerfile.golang"
|
||||
context: "./"
|
||||
# The CUDA-13 counterpart to the JetPack r36.4.0 row in the nemo-speech-cpp
|
||||
# block below. A Jetson whose CUDA 13 runtime is present reports the
|
||||
# nvidia-l4t-cuda-13 capability, and pointing that key at the JetPack image
|
||||
# would hand it a ggml linked against CUDA 12 whose libcudart.so.12 is not
|
||||
# there to dlopen. Same base and runner as the parakeet-cpp row above.
|
||||
- build-type: 'cublas'
|
||||
cuda-major-version: "13"
|
||||
cuda-minor-version: "0"
|
||||
platforms: 'linux/arm64'
|
||||
skip-drivers: 'false'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-nvidia-l4t-cuda-13-arm64-nemo-speech-cpp'
|
||||
base-image: "ubuntu:24.04"
|
||||
ubuntu-version: '2404'
|
||||
runs-on: 'ubuntu-24.04-arm'
|
||||
backend: "nemo-speech-cpp"
|
||||
dockerfile: "./backend/Dockerfile.golang"
|
||||
context: "./"
|
||||
- build-type: 'cublas'
|
||||
cuda-major-version: "13"
|
||||
cuda-minor-version: "0"
|
||||
@@ -4227,86 +4183,6 @@ include:
|
||||
dockerfile: "./backend/Dockerfile.golang"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
# nemo-speech-cpp
|
||||
#
|
||||
# No hipblas and no sycl rows, unlike the parakeet-cpp block above: upstream
|
||||
# NeMo-Speech.cpp builds ggml with CUDA, Vulkan or Metal only, so a ROCm or
|
||||
# SYCL image would be a CPU build wearing a GPU tag.
|
||||
#
|
||||
# cpu and vulkan are per-arch pairs sharing one tag-suffix, so
|
||||
# backend-merge-jobs assembles a multi-arch manifest from the two digests.
|
||||
# The arm64 legs are not redundant with the Jetson image below: an ARM server
|
||||
# with no NVIDIA GPU reports the "default" capability and would otherwise pull
|
||||
# an amd64-only manifest.
|
||||
- build-type: ''
|
||||
cuda-major-version: ""
|
||||
cuda-minor-version: ""
|
||||
platforms: 'linux/amd64'
|
||||
platform-tag: 'amd64'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-cpu-nemo-speech-cpp'
|
||||
runs-on: 'ubuntu-latest'
|
||||
base-image: "ubuntu:24.04"
|
||||
skip-drivers: 'false'
|
||||
backend: "nemo-speech-cpp"
|
||||
dockerfile: "./backend/Dockerfile.golang"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
- build-type: ''
|
||||
cuda-major-version: ""
|
||||
cuda-minor-version: ""
|
||||
platforms: 'linux/arm64'
|
||||
platform-tag: 'arm64'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-cpu-nemo-speech-cpp'
|
||||
runs-on: 'ubuntu-24.04-arm'
|
||||
base-image: "ubuntu:24.04"
|
||||
skip-drivers: 'false'
|
||||
backend: "nemo-speech-cpp"
|
||||
dockerfile: "./backend/Dockerfile.golang"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
- build-type: 'vulkan'
|
||||
cuda-major-version: ""
|
||||
cuda-minor-version: ""
|
||||
platforms: 'linux/amd64'
|
||||
platform-tag: 'amd64'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-gpu-vulkan-nemo-speech-cpp'
|
||||
runs-on: 'ubuntu-latest'
|
||||
base-image: "ubuntu:24.04"
|
||||
skip-drivers: 'false'
|
||||
backend: "nemo-speech-cpp"
|
||||
dockerfile: "./backend/Dockerfile.golang"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
- build-type: 'vulkan'
|
||||
cuda-major-version: ""
|
||||
cuda-minor-version: ""
|
||||
platforms: 'linux/arm64'
|
||||
platform-tag: 'arm64'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-gpu-vulkan-nemo-speech-cpp'
|
||||
runs-on: 'ubuntu-24.04-arm'
|
||||
base-image: "ubuntu:24.04"
|
||||
skip-drivers: 'false'
|
||||
backend: "nemo-speech-cpp"
|
||||
dockerfile: "./backend/Dockerfile.golang"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
- build-type: 'cublas'
|
||||
cuda-major-version: "12"
|
||||
cuda-minor-version: "0"
|
||||
platforms: 'linux/arm64'
|
||||
skip-drivers: 'false'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-nvidia-l4t-arm64-nemo-speech-cpp'
|
||||
base-image: "nvcr.io/nvidia/l4t-jetpack:r36.4.0"
|
||||
runs-on: 'ubuntu-24.04-arm'
|
||||
backend: "nemo-speech-cpp"
|
||||
dockerfile: "./backend/Dockerfile.golang"
|
||||
context: "./"
|
||||
ubuntu-version: '2204'
|
||||
# moss-transcribe-cpp
|
||||
- build-type: ''
|
||||
cuda-major-version: ""
|
||||
@@ -6350,10 +6226,6 @@ includeDarwin:
|
||||
tag-suffix: "-metal-darwin-arm64-moss-transcribe-cpp"
|
||||
build-type: "metal"
|
||||
lang: "go"
|
||||
- backend: "nemo-speech-cpp"
|
||||
tag-suffix: "-metal-darwin-arm64-nemo-speech-cpp"
|
||||
build-type: "metal"
|
||||
lang: "go"
|
||||
- backend: "ced"
|
||||
tag-suffix: "-metal-darwin-arm64-ced"
|
||||
build-type: "metal"
|
||||
|
||||
6
.github/workflows/backend_merge.yml
vendored
6
.github/workflows/backend_merge.yml
vendored
@@ -71,8 +71,8 @@ jobs:
|
||||
|
||||
# cosign signs each pushed manifest list with --recursive so the
|
||||
# index and every per-arch entry get an attached Sigstore bundle.
|
||||
# The pinned cosign v2 release needs --new-bundle-format explicitly;
|
||||
# the verifier only consumes OCI 1.1 Sigstore bundle referrers.
|
||||
# Recent cosign releases always emit the new bundle format, so
|
||||
# there's no extra CLI flag to opt into it.
|
||||
- name: Install cosign
|
||||
if: github.event_name != 'pull_request'
|
||||
uses: sigstore/cosign-installer@v3
|
||||
@@ -159,7 +159,6 @@ jobs:
|
||||
# manifest before checking signatures need the per-arch
|
||||
# signatures, not just the list-level one.
|
||||
cosign sign --yes --recursive \
|
||||
--new-bundle-format \
|
||||
--registry-referrers-mode=oci-1-1 \
|
||||
"quay.io/go-skynet/local-ai-backends@${digest}"
|
||||
|
||||
@@ -186,7 +185,6 @@ jobs:
|
||||
' <<< "$DOCKER_METADATA_OUTPUT_JSON")
|
||||
digest=$(docker buildx imagetools inspect "$first_tag" --format '{{.Manifest.Digest}}')
|
||||
cosign sign --yes --recursive \
|
||||
--new-bundle-format \
|
||||
--registry-referrers-mode=oci-1-1 \
|
||||
"localai/localai-backends@${digest}"
|
||||
|
||||
|
||||
16
.github/workflows/bump_deps.yaml
vendored
16
.github/workflows/bump_deps.yaml
vendored
@@ -62,10 +62,6 @@ jobs:
|
||||
variable: "MOSS_VERSION"
|
||||
branch: "master"
|
||||
file: "backend/go/moss-transcribe-cpp/Makefile"
|
||||
- repository: "NVIDIA/NeMo-Speech.cpp"
|
||||
variable: "NEMO_SPEECH_VERSION"
|
||||
branch: "main"
|
||||
file: "backend/go/nemo-speech-cpp/Makefile"
|
||||
- repository: "localai-org/ced.cpp"
|
||||
variable: "CED_VERSION"
|
||||
branch: "main"
|
||||
@@ -114,14 +110,10 @@ jobs:
|
||||
variable: "LOCATEANYTHING_VERSION"
|
||||
branch: "master"
|
||||
file: "backend/go/locate-anything-cpp/Makefile"
|
||||
# qwentts.cpp is held, not tracked: upstream master hangs in synthesis
|
||||
# (see the comment on QWEN3TTS_CPP_VERSION in the backend Makefile).
|
||||
# Leaving it here would re-bump the pin back onto the hang every night.
|
||||
# Restore this entry once the upstream fix lands.
|
||||
# - repository: "ServeurpersoCom/qwentts.cpp"
|
||||
# variable: "QWEN3TTS_CPP_VERSION"
|
||||
# branch: "master"
|
||||
# file: "backend/go/qwen3-tts-cpp/Makefile"
|
||||
- repository: "ServeurpersoCom/qwentts.cpp"
|
||||
variable: "QWEN3TTS_CPP_VERSION"
|
||||
branch: "master"
|
||||
file: "backend/go/qwen3-tts-cpp/Makefile"
|
||||
- repository: "ServeurpersoCom/omnivoice.cpp"
|
||||
variable: "OMNIVOICE_VERSION"
|
||||
branch: "master"
|
||||
|
||||
11
.github/workflows/gh-pages.yml
vendored
11
.github/workflows/gh-pages.yml
vendored
@@ -51,16 +51,7 @@ jobs:
|
||||
- name: Setup Go
|
||||
uses: actions/setup-go@v5
|
||||
with:
|
||||
# Track go.mod rather than a literal. Pinned at 1.22 this installed a
|
||||
# toolchain older than the module's `go 1.26.0`, so the `go run` below
|
||||
# downloaded the real one from proxy.golang.org on every run. That
|
||||
# fetch is not always reachable from the runner and the deploy failed
|
||||
# on five of eight consecutive master pushes with:
|
||||
# go: download go1.26.0: ... connect: network is unreachable
|
||||
# ##[error]Command failed: go env GOPATH
|
||||
# Installing the version the module asks for removes the download
|
||||
# instead of depending on it succeeding.
|
||||
go-version-file: go.mod
|
||||
go-version: '1.22'
|
||||
cache: false
|
||||
|
||||
- name: Setup Hugo
|
||||
|
||||
52
.github/workflows/test-extra.yml
vendored
52
.github/workflows/test-extra.yml
vendored
@@ -50,7 +50,6 @@ jobs:
|
||||
sherpa-onnx: ${{ steps.detect.outputs.sherpa-onnx }}
|
||||
whisper: ${{ steps.detect.outputs.whisper }}
|
||||
parakeet-cpp: ${{ steps.detect.outputs.parakeet-cpp }}
|
||||
nemo-speech-cpp: ${{ steps.detect.outputs.nemo-speech-cpp }}
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@v7
|
||||
@@ -901,57 +900,6 @@ jobs:
|
||||
- name: Test magpie-tts-cpp
|
||||
run: |
|
||||
make --jobs=5 --output-sync=target -C backend/go/magpie-tts-cpp test
|
||||
# Per-backend unit suite for nemo-speech-cpp. This job exists for one reason
|
||||
# above all: abi_test.go asserts the size and field offsets of every Go mirror
|
||||
# struct against the C ABI it is dlopened into. Those assertions are the only
|
||||
# thing standing between a purego symbol rename or an upstream header change
|
||||
# and silent memory corruption at run time, and they are worthless unless
|
||||
# something executes them. `make -C backend/go/nemo-speech-cpp test` sets
|
||||
# NEMO_SPEECH_REQUIRE_LIBS=1, which turns "library missing" from a skip into a
|
||||
# failure, so this job cannot report green having checked nothing.
|
||||
#
|
||||
# The backend Makefile's `test` target depends on `stage-libs`, so it clones
|
||||
# upstream at the pinned SHA and builds the native runtime itself. There is no
|
||||
# separate build step for that reason, and no model download: the specs are
|
||||
# ABI and pure-Go only.
|
||||
#
|
||||
# WITH_NORM=OFF skips the Sparrowhawk/OpenFST inverse-text-normalization
|
||||
# stack, which is the single most expensive leg of the build and needs a gcc-12
|
||||
# pin because OpenFST's templates ICE on gcc-13/14. It costs no coverage here:
|
||||
# nothing in include/nemo_speech/{asr,tts,diar,nmt}.h is conditional on it (the
|
||||
# only preprocessor conditionals in those headers are include guards,
|
||||
# __cplusplus and the _WIN32 export macros), so every struct layout this suite
|
||||
# checks is identical either way. The shipped images still build WITH_NORM=ON;
|
||||
# that path is covered by the backend image build in backend_pr.yml.
|
||||
tests-nemo-speech-cpp:
|
||||
needs: detect-changes
|
||||
if: needs.detect-changes.outputs.nemo-speech-cpp == 'true' || needs.detect-changes.outputs.run-all == 'true'
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 90
|
||||
steps:
|
||||
- name: Clone
|
||||
uses: actions/checkout@v7
|
||||
with:
|
||||
submodules: true
|
||||
- name: Dependencies
|
||||
run: |
|
||||
sudo apt-get update
|
||||
sudo apt-get install -y build-essential cmake ninja-build curl libopenblas-dev ffmpeg
|
||||
- name: Setup Go
|
||||
uses: actions/setup-go@v5
|
||||
- name: Display Go version
|
||||
run: go version
|
||||
- name: Proto Dependencies
|
||||
run: |
|
||||
curl -L -s https://github.com/protocolbuffers/protobuf/releases/download/v26.1/protoc-26.1-linux-x86_64.zip -o protoc.zip && \
|
||||
unzip -j -d /usr/local/bin protoc.zip bin/protoc && \
|
||||
rm protoc.zip
|
||||
go install google.golang.org/protobuf/cmd/protoc-gen-go@v1.34.2
|
||||
go install google.golang.org/grpc/cmd/protoc-gen-go-grpc@1958fcbe2ca8bd93af633f11e97d44e567e945af
|
||||
PATH="$PATH:$HOME/go/bin" make protogen-go
|
||||
- name: Test nemo-speech-cpp
|
||||
run: |
|
||||
make --jobs=5 --output-sync=target -C backend/go/nemo-speech-cpp WITH_NORM=OFF test
|
||||
# Per-backend smoke for rfdetr-cpp: builds the .so + Go binary and runs
|
||||
# `make -C backend/go/rfdetr-cpp test`. test.sh fetches the small (~20 MB)
|
||||
# rfdetr-nano-q8_0 GGUF from the published mudler/rfdetr-cpp-nano HF repo
|
||||
|
||||
5
.gitignore
vendored
5
.gitignore
vendored
@@ -124,8 +124,3 @@ formal-verification/out/
|
||||
# package directory itself and untrack the source.
|
||||
/apexentries
|
||||
/.github/ci/apexentries/apexentries
|
||||
|
||||
# Runtime state written by `local-ai run` when it is started from the repo
|
||||
# root, which is what a contributor testing a build does. Nothing under here is
|
||||
# source: it is the instance's own models, outputs, traces and identity.
|
||||
/data/
|
||||
|
||||
7
Makefile
7
Makefile
@@ -1,5 +1,5 @@
|
||||
# Disable parallel execution for backend builds
|
||||
.NOTPARALLEL: backends/diffusers backends/llama-cpp backends/turboquant backends/bonsai backends/outetts backends/piper backends/stablediffusion-ggml backends/trellis2cpp backends/trellis2cpp-darwin backends/whisper backends/crispasr backends/parakeet-cpp backends/moss-transcribe-cpp backends/nemo-speech-cpp backends/faster-whisper backends/silero-vad backends/local-store backends/valkey-store backends/cloud-proxy backends/huggingface backends/rfdetr backends/rfdetr-cpp backends/insightface backends/speaker-recognition backends/kitten-tts backends/kokoro backends/chatterbox backends/llama-cpp-darwin backends/neutts build-darwin-python-backend build-darwin-go-backend backends/mlx backends/diffuser-darwin backends/mlx-vlm backends/mlx-audio backends/mlx-distributed backends/stablediffusion-ggml-darwin backends/vllm backends/vllm-omni backends/longcat-video backends/sglang backends/moonshine backends/pocket-tts backends/qwen-tts backends/faster-qwen3-tts backends/qwen-asr backends/nemo backends/voxcpm backends/whisperx backends/ace-step backends/acestep-cpp backends/fish-speech backends/voxtral backends/opus backends/trl backends/llama-cpp-quantization backends/kokoros backends/sam3-cpp backends/qwen3-tts-cpp backends/moss-tts-cpp backends/magpie-tts-cpp backends/vllm-cpp backends/omnivoice-cpp backends/vibevoice-cpp backends/localvqe backends/tinygrad backends/sherpa-onnx backends/ds4 backends/ds4-darwin backends/liquid-audio backends/supertonic backends/depth-anything-cpp backends/privacy-filter backends/privacy-filter-darwin backends/audio-cpp backends/audio-cpp-darwin
|
||||
.NOTPARALLEL: backends/diffusers backends/llama-cpp backends/turboquant backends/bonsai backends/outetts backends/piper backends/stablediffusion-ggml backends/trellis2cpp backends/trellis2cpp-darwin backends/whisper backends/crispasr backends/parakeet-cpp backends/moss-transcribe-cpp backends/faster-whisper backends/silero-vad backends/local-store backends/valkey-store backends/cloud-proxy backends/huggingface backends/rfdetr backends/rfdetr-cpp backends/insightface backends/speaker-recognition backends/kitten-tts backends/kokoro backends/chatterbox backends/llama-cpp-darwin backends/neutts build-darwin-python-backend build-darwin-go-backend backends/mlx backends/diffuser-darwin backends/mlx-vlm backends/mlx-audio backends/mlx-distributed backends/stablediffusion-ggml-darwin backends/vllm backends/vllm-omni backends/longcat-video backends/sglang backends/moonshine backends/pocket-tts backends/qwen-tts backends/faster-qwen3-tts backends/qwen-asr backends/nemo backends/voxcpm backends/whisperx backends/ace-step backends/acestep-cpp backends/fish-speech backends/voxtral backends/opus backends/trl backends/llama-cpp-quantization backends/kokoros backends/sam3-cpp backends/qwen3-tts-cpp backends/moss-tts-cpp backends/magpie-tts-cpp backends/vllm-cpp backends/omnivoice-cpp backends/vibevoice-cpp backends/localvqe backends/tinygrad backends/sherpa-onnx backends/ds4 backends/ds4-darwin backends/liquid-audio backends/supertonic backends/depth-anything-cpp backends/privacy-filter backends/privacy-filter-darwin backends/audio-cpp backends/audio-cpp-darwin
|
||||
|
||||
GOCMD=go
|
||||
GOTEST=$(GOCMD) test
|
||||
@@ -654,7 +654,6 @@ test-extra: prepare-test-extra
|
||||
$(MAKE) -C backend/go/depth-anything-cpp test
|
||||
$(MAKE) -C backend/go/supertonic test
|
||||
$(MAKE) -C backend/go/vllm-cpp test
|
||||
$(MAKE) -C backend/go/nemo-speech-cpp test
|
||||
$(MAKE) -C backend/go/trellis2cpp test
|
||||
$(MAKE) -C backend/go/valkey-store test
|
||||
|
||||
@@ -1299,7 +1298,6 @@ BACKEND_WHISPER = whisper|golang|.|false|true
|
||||
BACKEND_CRISPASR = crispasr|golang|.|false|true
|
||||
BACKEND_PARAKEET_CPP = parakeet-cpp|golang|.|false|true
|
||||
BACKEND_MOSS_TRANSCRIBE_CPP = moss-transcribe-cpp|golang|.|false|true
|
||||
BACKEND_NEMO_SPEECH_CPP = nemo-speech-cpp|golang|.|false|true
|
||||
BACKEND_DEPTH_ANYTHING_CPP = depth-anything-cpp|golang|.|false|true
|
||||
BACKEND_VOXTRAL = voxtral|golang|.|false|true
|
||||
BACKEND_ACESTEP_CPP = acestep-cpp|golang|.|false|true
|
||||
@@ -1402,7 +1400,6 @@ $(eval $(call generate-docker-build-target,$(BACKEND_WHISPER)))
|
||||
$(eval $(call generate-docker-build-target,$(BACKEND_CRISPASR)))
|
||||
$(eval $(call generate-docker-build-target,$(BACKEND_PARAKEET_CPP)))
|
||||
$(eval $(call generate-docker-build-target,$(BACKEND_MOSS_TRANSCRIBE_CPP)))
|
||||
$(eval $(call generate-docker-build-target,$(BACKEND_NEMO_SPEECH_CPP)))
|
||||
$(eval $(call generate-docker-build-target,$(BACKEND_DEPTH_ANYTHING_CPP)))
|
||||
$(eval $(call generate-docker-build-target,$(BACKEND_VOXTRAL)))
|
||||
$(eval $(call generate-docker-build-target,$(BACKEND_OPUS)))
|
||||
@@ -1459,7 +1456,7 @@ $(eval $(call generate-docker-build-target,$(BACKEND_SUPERTONIC)))
|
||||
docker-save-%: backend-images
|
||||
docker save local-ai-backend:$* -o backend-images/$*.tar
|
||||
|
||||
docker-build-backends: docker-build-llama-cpp docker-build-ik-llama-cpp docker-build-turboquant docker-build-bonsai docker-build-ds4 docker-build-rerankers docker-build-vllm docker-build-vllm-omni docker-build-longcat-video docker-build-sglang docker-build-transformers docker-build-outetts docker-build-diffusers docker-build-kokoro docker-build-faster-whisper docker-build-crispasr docker-build-coqui docker-build-chatterbox docker-build-vibevoice docker-build-liquid-audio docker-build-moonshine docker-build-pocket-tts docker-build-qwen-tts docker-build-fish-speech docker-build-faster-qwen3-tts docker-build-qwen-asr docker-build-nemo docker-build-voxcpm docker-build-whisperx docker-build-ace-step docker-build-acestep-cpp docker-build-voxtral docker-build-mlx-distributed docker-build-trl docker-build-llama-cpp-quantization docker-build-tinygrad docker-build-kokoros docker-build-sam3-cpp docker-build-rfdetr-cpp docker-build-qwen3-tts-cpp docker-build-moss-tts-cpp docker-build-magpie-tts-cpp docker-build-vllm-cpp docker-build-omnivoice-cpp docker-build-vibevoice-cpp docker-build-localvqe docker-build-insightface docker-build-speaker-recognition docker-build-sherpa-onnx docker-build-cloud-proxy docker-build-supertonic docker-build-depth-anything-cpp docker-build-moss-transcribe-cpp docker-build-nemo-speech-cpp docker-build-privacy-filter docker-build-trellis2cpp docker-build-valkey-store docker-build-audio-cpp
|
||||
docker-build-backends: docker-build-llama-cpp docker-build-ik-llama-cpp docker-build-turboquant docker-build-bonsai docker-build-ds4 docker-build-rerankers docker-build-vllm docker-build-vllm-omni docker-build-longcat-video docker-build-sglang docker-build-transformers docker-build-outetts docker-build-diffusers docker-build-kokoro docker-build-faster-whisper docker-build-crispasr docker-build-coqui docker-build-chatterbox docker-build-vibevoice docker-build-liquid-audio docker-build-moonshine docker-build-pocket-tts docker-build-qwen-tts docker-build-fish-speech docker-build-faster-qwen3-tts docker-build-qwen-asr docker-build-nemo docker-build-voxcpm docker-build-whisperx docker-build-ace-step docker-build-acestep-cpp docker-build-voxtral docker-build-mlx-distributed docker-build-trl docker-build-llama-cpp-quantization docker-build-tinygrad docker-build-kokoros docker-build-sam3-cpp docker-build-rfdetr-cpp docker-build-qwen3-tts-cpp docker-build-moss-tts-cpp docker-build-magpie-tts-cpp docker-build-vllm-cpp docker-build-omnivoice-cpp docker-build-vibevoice-cpp docker-build-localvqe docker-build-insightface docker-build-speaker-recognition docker-build-sherpa-onnx docker-build-cloud-proxy docker-build-supertonic docker-build-depth-anything-cpp docker-build-moss-transcribe-cpp docker-build-privacy-filter docker-build-trellis2cpp docker-build-valkey-store docker-build-audio-cpp
|
||||
|
||||
########################################################
|
||||
### Mock Backend for E2E Tests
|
||||
|
||||
@@ -161,7 +161,7 @@ local-ai run https://gist.githubusercontent.com/.../phi-2.yaml
|
||||
local-ai run oci://localai/phi-2:latest
|
||||
```
|
||||
|
||||
To work with a running LocalAI server from the terminal, start the built-in agent from another shell. It answers questions, reads your files and runs commands on your machine, asking you to approve anything that changes state. Inside a session, `/models` lists installed models and `/model <name>` switches between them. See the [Terminal agent](https://localai.io/docs/features/terminal-agent/) docs.
|
||||
To test a running LocalAI server from the terminal, open an interactive chat session from another shell. Inside the prompt, `/models` lists installed models and `/model <name>` switches between them.
|
||||
|
||||
```bash
|
||||
# Terminal 1
|
||||
@@ -195,7 +195,7 @@ For more details, see the [Getting Started guide](https://localai.io/basics/gett
|
||||
- **August 2025**: MLX, MLX-VLM, Diffusers, llama.cpp now supported on Apple Silicon
|
||||
- **July 2025**: All backends migrated outside the main binary — [lightweight, modular architecture](https://github.com/mudler/LocalAI/releases/tag/v3.2.0)
|
||||
|
||||
For older news and full release notes, see [GitHub Releases](https://github.com/mudler/LocalAI/releases) and the [blog](https://localai.io/blog/).
|
||||
For older news and full release notes, see [GitHub Releases](https://github.com/mudler/LocalAI/releases) and the [News page](https://localai.io/basics/news/).
|
||||
|
||||
## Features
|
||||
|
||||
@@ -260,7 +260,7 @@ We also maintain [apex-quant](https://github.com/localai-org/apex-quant), a per-
|
||||
- [Kubernetes installation](https://localai.io/basics/getting_started/#run-localai-in-kubernetes)
|
||||
- [Integrations & community projects](https://localai.io/docs/integrations/)
|
||||
- [Installation video walkthrough](https://www.youtube.com/watch?v=cMVNnlqwfw4)
|
||||
- [Blog: release write-ups, benchmarks and engineering notes](https://localai.io/blog/)
|
||||
- [Media & blog posts](https://localai.io/basics/news/#media-blogs-social)
|
||||
- [Examples](https://github.com/mudler/LocalAI-examples) — including the [realtime voice assistant demo](https://github.com/localai-org/localai-realtime-demo) (Go client for the Realtime API with tool calling)
|
||||
|
||||
## Team
|
||||
|
||||
@@ -248,52 +248,6 @@ RUN <<EOT bash
|
||||
fi
|
||||
EOT
|
||||
|
||||
# nemo-speech-cpp builds NVIDIA NeMo-Speech.cpp with text normalization enabled,
|
||||
# which compiles the Sparrowhawk/OpenFST WFST stack from source via
|
||||
# scripts/build_itn_deps.sh. That step needs gcc-12 specifically: OpenFST's
|
||||
# template-heavy translation units ICE on gcc-13 and gcc-14 at -O2, so upstream
|
||||
# pins gcc-12 for it while the runtime itself builds with the image default.
|
||||
# No update-alternatives here, so the default compiler is untouched; the backend
|
||||
# Makefile reaches gcc-12 by name for that one step.
|
||||
#
|
||||
# The rest is what build_itn_deps.sh and the WITH_NORM cmake block expect:
|
||||
# protobuf (headers plus protoc, which must come from the same apt set so the
|
||||
# generated stubs match the headers they compile against) and re2 for
|
||||
# Sparrowhawk, and autotools because OpenFST and Sparrowhawk ship autoconf
|
||||
# builds. ninja is not in the common apt list because this is the only Go
|
||||
# backend that configures with -G Ninja, and that list is a layer shared by
|
||||
# every backend image in the matrix.
|
||||
#
|
||||
# No libabsl-dev, despite upstream's Dockerfile installing it: upstream builds
|
||||
# against protobuf 25, which splits its runtime across libabsl_*, whereas every
|
||||
# base image in this matrix carries protobuf 3.21 (noble) or 3.12 (jammy), which
|
||||
# has no absl dependency. The cmake block's file(GLOB ... /usr/lib/libabsl_*.so)
|
||||
# would not match on Ubuntu anyway, since multiarch puts those under
|
||||
# /usr/lib/<triplet>/.
|
||||
#
|
||||
# Placed down here with the other per-backend gates rather than next to the
|
||||
# shared apt layer: Docker re-keys every layer below an inserted one, so adding
|
||||
# a step above the Vulkan SDK, CUDA, Go and protoc layers would force all of
|
||||
# them to re-execute once for every Go backend image, not just this one.
|
||||
# Nothing between there and here needs any of these packages (the Vulkan and
|
||||
# opus blocks install their own ninja and pkg-config, and the protoc download is
|
||||
# a release binary that needs neither libprotobuf-dev nor protoc from apt), and
|
||||
# nothing here needs anything those layers provide.
|
||||
RUN <<EOT bash
|
||||
if [ "${BACKEND}" = "nemo-speech-cpp" ]; then
|
||||
set -e
|
||||
apt-get update
|
||||
apt-get install -y --no-install-recommends \
|
||||
gcc-12 g++-12 \
|
||||
ninja-build \
|
||||
libprotobuf-dev protobuf-compiler \
|
||||
libre2-dev \
|
||||
autoconf automake libtool pkg-config
|
||||
apt-get clean
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
fi
|
||||
EOT
|
||||
|
||||
RUN git config --global --add safe.directory /LocalAI
|
||||
|
||||
# Prebuild the native engine from a layer that depends on this backend's own
|
||||
|
||||
@@ -15,7 +15,6 @@ service Backend {
|
||||
rpc PredictStream(PredictOptions) returns (stream Reply) {}
|
||||
rpc Embedding(PredictOptions) returns (EmbeddingResult) {}
|
||||
rpc GenerateImage(GenerateImageRequest) returns (Result) {}
|
||||
rpc UpscaleImage(UpscaleImageRequest) returns (Result) {}
|
||||
rpc GenerateVideo(GenerateVideoRequest) returns (Result) {}
|
||||
rpc Generate3D(Generate3DRequest) returns (Result) {}
|
||||
rpc AudioTranscription(TranscriptRequest) returns (TranscriptResult) {}
|
||||
@@ -638,12 +637,6 @@ message GenerateImageRequest {
|
||||
string ModelIdentity = 13;
|
||||
}
|
||||
|
||||
message UpscaleImageRequest {
|
||||
string src = 1; // input image path
|
||||
string dst = 2; // output image path
|
||||
int32 scale = 3; // upscale factor (e.g. 2 or 4)
|
||||
}
|
||||
|
||||
message GenerateVideoRequest {
|
||||
string prompt = 1;
|
||||
string negative_prompt = 2; // Negative prompt for video generation
|
||||
|
||||
@@ -9,7 +9,7 @@
|
||||
# recipe is a make target (not a prepare.sh) so 'make purge && make' is a clean
|
||||
# rebuild and so the bump bot can see the pin.
|
||||
|
||||
AUDIO_CPP_VERSION?=7efbb58def443722ea540d931dd3debee3e4d5e8
|
||||
AUDIO_CPP_VERSION?=f78227c52736a4792a50aa3f82ead7e7385c891b
|
||||
AUDIO_CPP_REPO?=https://github.com/0xShug0/audio.cpp
|
||||
|
||||
CURRENT_MAKEFILE_DIR := $(dir $(abspath $(lastword $(MAKEFILE_LIST))))
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
|
||||
# Pinned to the HEAD of the `prism` branch on https://github.com/PrismML-Eng/llama.cpp.
|
||||
# Auto-bumped nightly by .github/workflows/bump_deps.yaml.
|
||||
BONSAI_VERSION?=9ca265a57f85f2117942490f421f64a226dd9847
|
||||
BONSAI_VERSION?=4dd165625bb6c020285eec8b342af25cf60233dd
|
||||
LLAMA_REPO?=https://github.com/PrismML-Eng/llama.cpp
|
||||
|
||||
CMAKE_ARGS?=
|
||||
|
||||
@@ -69,15 +69,7 @@ target_include_directories(hw_grpc_proto PUBLIC ${CMAKE_CURRENT_BINARY_DIR})
|
||||
|
||||
set(DS4_OBJS "${DS4_DIR}/ds4.o")
|
||||
if(DS4_GPU STREQUAL "cuda")
|
||||
list(APPEND DS4_OBJS
|
||||
"${DS4_DIR}/ds4_cuda.o"
|
||||
"${DS4_DIR}/cuda/mmq/ds4_ggml_stubs.o"
|
||||
"${DS4_DIR}/cuda/mmq/ds4_mmq.o"
|
||||
"${DS4_DIR}/cuda/mmq/ds4_mmq_d2r.o"
|
||||
"${DS4_DIR}/cuda/mmq/quantize.o"
|
||||
"${DS4_DIR}/cuda/mmq/mmid.o"
|
||||
"${DS4_DIR}/cuda/mmq/mmvq.o"
|
||||
"${DS4_DIR}/cuda/mmq/ds4_repack.o")
|
||||
list(APPEND DS4_OBJS "${DS4_DIR}/ds4_cuda.o")
|
||||
elseif(DS4_GPU STREQUAL "metal")
|
||||
list(APPEND DS4_OBJS "${DS4_DIR}/ds4_metal.o")
|
||||
elseif(DS4_GPU STREQUAL "cpu")
|
||||
|
||||
@@ -1,10 +1,10 @@
|
||||
# ds4 backend Makefile.
|
||||
#
|
||||
# Upstream pin lives below as DS4_VERSION?=b0309611041655f4e45671cfd9c9886aff161406
|
||||
# Upstream pin lives below as DS4_VERSION?=54b36ed9ba42da31b24f2d1a5feb075c2475dbb1
|
||||
# (.github/bump_deps.sh) can find and update it - matches the
|
||||
# llama-cpp / ik-llama-cpp / turboquant convention.
|
||||
|
||||
DS4_VERSION?=b0309611041655f4e45671cfd9c9886aff161406
|
||||
DS4_VERSION?=54b36ed9ba42da31b24f2d1a5feb075c2475dbb1
|
||||
DS4_REPO?=https://github.com/antirez/ds4
|
||||
|
||||
CURRENT_MAKEFILE_DIR := $(dir $(abspath $(lastword $(MAKEFILE_LIST))))
|
||||
@@ -23,9 +23,7 @@ CMAKE_ARGS ?= -DCMAKE_BUILD_TYPE=Release
|
||||
# are shared by every GPU mode, so append them unconditionally below.
|
||||
ifeq ($(BUILD_TYPE),cublas)
|
||||
CMAKE_ARGS += -DDS4_GPU=cuda
|
||||
DS4_OBJ_TARGET := ds4.o ds4_cuda.o ds4_distributed.o ds4_tp.o ds4_ssd.o ds4_layer_pack.o \
|
||||
cuda/mmq/ds4_ggml_stubs.o cuda/mmq/ds4_mmq.o cuda/mmq/ds4_mmq_d2r.o \
|
||||
cuda/mmq/quantize.o cuda/mmq/mmid.o cuda/mmq/mmvq.o cuda/mmq/ds4_repack.o
|
||||
DS4_OBJ_TARGET := ds4.o ds4_cuda.o ds4_distributed.o ds4_tp.o ds4_ssd.o ds4_layer_pack.o
|
||||
else ifeq ($(UNAME_S),Darwin)
|
||||
CMAKE_ARGS += -DDS4_GPU=metal
|
||||
DS4_OBJ_TARGET := ds4.o ds4_metal.o ds4_distributed.o ds4_tp.o ds4_ssd.o ds4_layer_pack.o
|
||||
@@ -57,7 +55,7 @@ ds4:
|
||||
# the right per-platform compile flags (Objective-C/Metal on Darwin, nvcc on Linux+CUDA).
|
||||
ds4/ds4.o: ds4
|
||||
ifeq ($(BUILD_TYPE),cublas)
|
||||
+$(MAKE) -C ds4 $(DS4_OBJ_TARGET)
|
||||
+$(MAKE) -C ds4 ds4.o ds4_cuda.o ds4_distributed.o ds4_tp.o ds4_ssd.o ds4_layer_pack.o
|
||||
else ifeq ($(UNAME_S),Darwin)
|
||||
+$(MAKE) -C ds4 ds4.o ds4_metal.o ds4_distributed.o ds4_tp.o ds4_ssd.o ds4_layer_pack.o
|
||||
else
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
|
||||
IK_LLAMA_VERSION?=cf1aa57e1a0fabfd015831718fc99d1aec01ada5
|
||||
IK_LLAMA_VERSION?=3f53a059024039358e9fef75b5dc0c99dbcb40f9
|
||||
LLAMA_REPO?=https://github.com/ikawrakow/ik_llama.cpp
|
||||
|
||||
CMAKE_ARGS?=
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
|
||||
LLAMA_VERSION?=221f0f6356efe2260023208365705ec5d5a7c8f5
|
||||
LLAMA_VERSION?=876a4321163249c43ca4e986818fab5ab081f282
|
||||
LLAMA_REPO?=https://github.com/ggerganov/llama.cpp
|
||||
|
||||
CMAKE_ARGS?=
|
||||
|
||||
@@ -12,11 +12,10 @@ grep -e "flags" /proc/cpuinfo | head -1
|
||||
|
||||
BINARY=llama-cpp-fallback
|
||||
|
||||
# CPU images and most x86 GPU images ship a single llama-cpp-cpu-all built with ggml
|
||||
# CPU images and x86 GPU images ship a single llama-cpp-cpu-all built with ggml
|
||||
# CPU_ALL_VARIANTS: ggml's backend registry dlopens the best libggml-cpu-*.so for this
|
||||
# host, so no shell-side AVX probing. GPU arm64 images still ship llama-cpp-fallback
|
||||
# until their builder toolchains support ggml's complete arm variant matrix, and so do
|
||||
# the SYCL images, whose icpx compiler hangs on the sapphirerapids variant.
|
||||
# until their builder toolchains support ggml's complete arm variant matrix.
|
||||
if [ -e "$CURDIR"/llama-cpp-cpu-all ]; then
|
||||
BINARY=llama-cpp-cpu-all
|
||||
fi
|
||||
|
||||
@@ -12,12 +12,11 @@ grep -e "flags" /proc/cpuinfo | head -1
|
||||
|
||||
BINARY=turboquant-fallback
|
||||
|
||||
# CPU images and most x86 GPU images ship a single turboquant-cpu-all built with ggml
|
||||
# CPU images and x86 GPU images ship a single turboquant-cpu-all built with ggml
|
||||
# CPU_ALL_VARIANTS: ggml's
|
||||
# backend registry dlopens the best libggml-cpu-*.so for this host, so no shell-side
|
||||
# probing. GPU arm64 images still ship turboquant-fallback until their builder toolchains
|
||||
# support ggml's complete arm variant matrix, and so do the SYCL images, whose icpx
|
||||
# compiler hangs on the sapphirerapids variant.
|
||||
# support ggml's complete arm variant matrix.
|
||||
if [ -e "$CURDIR"/turboquant-cpu-all ]; then
|
||||
BINARY=turboquant-cpu-all
|
||||
fi
|
||||
|
||||
@@ -8,7 +8,7 @@ JOBS?=$(shell nproc --ignore=1)
|
||||
|
||||
# CrispASR version (release tag)
|
||||
CRISPASR_REPO?=https://github.com/CrispStrobe/CrispASR
|
||||
CRISPASR_VERSION?=21901d3f7c23554f072964828363e49ddbc2dc68
|
||||
CRISPASR_VERSION?=b5211ac635489049ee8ce86a82d69faa18e8d8da
|
||||
SO_TARGET?=libgocrispasr.so
|
||||
|
||||
CMAKE_ARGS+=-DBUILD_SHARED_LIBS=OFF
|
||||
|
||||
@@ -67,16 +67,7 @@ const defaultTTSSampleRate = 24000
|
||||
// resampling, so the WAV header must match it. Returns ok=false for non-piper
|
||||
// models (key absent) or an unreadable file, letting the caller fall back to
|
||||
// defaultTTSSampleRate.
|
||||
func piperSampleRate(modelPath string) (rate int, ok bool) {
|
||||
// A malformed metadata length can make gguf-parser-go panic before it can
|
||||
// return an error. Keep a bad voice file from crash-looping the backend.
|
||||
defer func() {
|
||||
if recover() != nil {
|
||||
rate = 0
|
||||
ok = false
|
||||
}
|
||||
}()
|
||||
|
||||
func piperSampleRate(modelPath string) (int, bool) {
|
||||
// Only scalar architecture keys are read, so skip the large array metadata
|
||||
// (phoneme map) and mmap the header - same rationale as pkg/vram's reader.
|
||||
f, err := gguf.ParseGGUFFile(modelPath, gguf.UseMMap(), gguf.SkipLargeMetadata())
|
||||
@@ -87,7 +78,7 @@ func piperSampleRate(modelPath string) (rate int, ok bool) {
|
||||
if !ok || kv.ValueType != gguf.GGUFMetadataValueTypeUint32 {
|
||||
return 0, false
|
||||
}
|
||||
rate = int(kv.ValueUint32())
|
||||
rate := int(kv.ValueUint32())
|
||||
if rate <= 0 {
|
||||
return 0, false
|
||||
}
|
||||
|
||||
@@ -3,7 +3,6 @@ package main
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/binary"
|
||||
"math"
|
||||
"os"
|
||||
"path/filepath"
|
||||
|
||||
@@ -103,24 +102,6 @@ var _ = Describe("piper sample rate", func() {
|
||||
_, ok := piperSampleRate(p)
|
||||
Expect(ok).To(BeFalse())
|
||||
})
|
||||
|
||||
It("returns ok=false instead of panicking on a malformed string length", func() {
|
||||
p := filepath.Join(GinkgoT().TempDir(), "malformed.gguf")
|
||||
var b bytes.Buffer
|
||||
b.WriteString("GGUF")
|
||||
Expect(binary.Write(&b, binary.LittleEndian, uint32(3))).To(Succeed())
|
||||
Expect(binary.Write(&b, binary.LittleEndian, uint64(0))).To(Succeed())
|
||||
Expect(binary.Write(&b, binary.LittleEndian, uint64(1))).To(Succeed())
|
||||
key := "general.name"
|
||||
Expect(binary.Write(&b, binary.LittleEndian, uint64(len(key)))).To(Succeed())
|
||||
b.WriteString(key)
|
||||
Expect(binary.Write(&b, binary.LittleEndian, ggufTypeString)).To(Succeed())
|
||||
Expect(binary.Write(&b, binary.LittleEndian, uint64(math.MaxInt64))).To(Succeed())
|
||||
Expect(os.WriteFile(p, b.Bytes(), 0o644)).To(Succeed())
|
||||
|
||||
_, ok := piperSampleRate(p)
|
||||
Expect(ok).To(BeFalse())
|
||||
})
|
||||
})
|
||||
|
||||
// End-to-end through the built .so. Gated on CRISPASR_PIPER_MODEL_PATH (a
|
||||
|
||||
22
backend/go/nemo-speech-cpp/.gitignore
vendored
22
backend/go/nemo-speech-cpp/.gitignore
vendored
@@ -1,22 +0,0 @@
|
||||
# Fetched upstream sources
|
||||
sources/
|
||||
|
||||
# CMake build directories
|
||||
build*/
|
||||
|
||||
# Packaging output
|
||||
package/
|
||||
|
||||
# Compiled backend binary. The second name is what a bare `go build ./...` from
|
||||
# this directory produces (it names the binary after the directory), as opposed
|
||||
# to the -o name the Makefile asks for.
|
||||
nemo-speech-cpp-grpc
|
||||
/nemo-speech-cpp
|
||||
|
||||
# Shared libraries staged in-tree by the Makefile (cp from sources/). The
|
||||
# SOVERSION suffix means the payload is libnemo_speech_*.so.1, hence both globs.
|
||||
*.so
|
||||
*.so.*
|
||||
*.dylib
|
||||
|
||||
compile_commands.json
|
||||
@@ -1,312 +0,0 @@
|
||||
# nemo-speech-cpp backend Makefile.
|
||||
#
|
||||
# Upstream pin lives below as NEMO_SPEECH_VERSION so .github/bump_deps.sh can
|
||||
# find and update it, matching the parakeet-cpp / vibevoice-cpp convention.
|
||||
#
|
||||
# Bumping NEMO_SPEECH_VERSION is a no-op on an existing checkout: sources/ is a
|
||||
# directory target, so make only clones when it is missing and never re-checks
|
||||
# out an already-cloned tree. After a bump run 'make purge && make', the same
|
||||
# rule the parakeet-cpp Makefile documents.
|
||||
#
|
||||
# 'build' is the entry point the backend image calls (backend/Dockerfile.golang
|
||||
# runs 'make -C backend/go/$(BACKEND) build' and then copies package/), so it
|
||||
# has to produce the binary and the package, not just the shared libraries.
|
||||
|
||||
NEMO_SPEECH_VERSION?=2e12e2def8a98ed06666f7ee3ca94e7193e04be4
|
||||
NEMO_SPEECH_REPO?=https://github.com/NVIDIA/NeMo-Speech.cpp
|
||||
|
||||
GOCMD?=go
|
||||
GO_TAGS?=
|
||||
JOBS?=$(shell nproc 2>/dev/null || sysctl -n hw.ncpu 2>/dev/null || echo 4)
|
||||
|
||||
BUILD_TYPE?=
|
||||
NATIVE?=false
|
||||
|
||||
# NEMO_SPEECH_CUBLAS_SHIM defaults ON upstream and builds a drop-in
|
||||
# libcublas.so.13. LocalAI's CUDA images ship the real cuBLAS, so the shim would
|
||||
# shadow it with a slower native GEMM. Always OFF here.
|
||||
CMAKE_ARGS?=-DCMAKE_BUILD_TYPE=Release \
|
||||
-DBUILD_SHARED_LIBS=OFF \
|
||||
-DCMAKE_POSITION_INDEPENDENT_CODE=ON \
|
||||
-DNEMO_SPEECH_CUBLAS_SHIM=OFF \
|
||||
-DNEMO_SPEECH_BUILD_ASR=ON \
|
||||
-DNEMO_SPEECH_BUILD_DIAR=ON \
|
||||
-DNEMO_SPEECH_BUILD_TTS=ON \
|
||||
-DNEMO_SPEECH_BUILD_NMT=ON \
|
||||
-DNEMO_SPEECH_BUILD_CLI=OFF \
|
||||
-DNEMO_SPEECH_BUILD_HTTP=OFF \
|
||||
-DNEMO_SPEECH_BUILD_GRPC=OFF \
|
||||
-DNEMO_SPEECH_WITH_FLASHLIGHT=OFF \
|
||||
-DNEMO_SPEECH_TTS_WITH_ZH=ON \
|
||||
-DNEMO_SPEECH_TTS_WITH_JA=ON
|
||||
|
||||
ifeq ($(NATIVE),false)
|
||||
CMAKE_ARGS+=-DGGML_NATIVE=OFF
|
||||
endif
|
||||
|
||||
# NEMO_SPEECH_TTS_WITH_JA=ON compiles Open JTalk's bundled MeCab, and
|
||||
# mecab/src/dictionary.cpp derives a comparator from std::binary_function, which
|
||||
# C++17 removed. libstdc++ still ships it as deprecated-but-present under
|
||||
# -std=gnu++17, so Linux never notices; libc++ compiles it out and the build dies
|
||||
# with "no template named 'binary_function' in namespace 'std'". Upstream's own
|
||||
# CMakeLists already carries the equivalent workaround for MSVC's STL
|
||||
# (_HAS_AUTO_PTR_ETC plus /FIfunctional) but has no libc++ branch, because
|
||||
# NEMO_SPEECH_TTS_WITH_JA defaults OFF upstream and only LocalAI turns it on.
|
||||
#
|
||||
# libc++ gates the two templates on _LIBCPP_ENABLE_CXX17_REMOVED_UNARY_BINARY_FUNCTION,
|
||||
# and has done since LLVM 16, which is older than any clang Xcode still ships.
|
||||
# The name matters: the older _LIBCPP_ENABLE_CXX17_REMOVED_BINDERS covers
|
||||
# bind1st/bind2nd/ptr_fun/mem_fun and NOT unary_function/binary_function, and the
|
||||
# umbrella _LIBCPP_ENABLE_CXX17_REMOVED_FEATURES no longer exists at all. A wrong
|
||||
# name is silently accepted by the preprocessor and fixes nothing.
|
||||
#
|
||||
# Applied through CMAKE_CXX_FLAGS rather than to the one target because the
|
||||
# tokenizer CMakeLists is upstream's and this tree is a pinned checkout, not a
|
||||
# patched one. Project-wide is also the safer scope: the macro decides whether
|
||||
# libc++'s internal __binary_function alias resolves to std::binary_function or
|
||||
# to __binary_function_keep_layout_base, which is a base class of std::less and
|
||||
# friends, so defining it for a subset of translation units would give those
|
||||
# class templates two spellings in one binary. Both bases are empty and, at
|
||||
# C++17, carry identical members, so the project-wide define changes no layout
|
||||
# and no ABI. On Linux the macro is not a name libstdc++ knows, so the branch is
|
||||
# unreachable there and would be inert even if it were taken.
|
||||
ifeq ($(shell uname -s),Darwin)
|
||||
CXX_COMPAT_FLAGS?=-D_LIBCPP_ENABLE_CXX17_REMOVED_UNARY_BINARY_FUNCTION
|
||||
else
|
||||
CXX_COMPAT_FLAGS?=
|
||||
endif
|
||||
ifneq ($(strip $(CXX_COMPAT_FLAGS)),)
|
||||
CMAKE_ARGS+=-DCMAKE_CXX_FLAGS=$(CXX_COMPAT_FLAGS)
|
||||
endif
|
||||
|
||||
# scripts/build_itn_deps.sh installs the Sparrowhawk/OpenFST runtime here.
|
||||
# NEMO_SPEECH_DEPENDENCY_PREFIX defaults to <src>/.deps upstream, and the ITN
|
||||
# stack goes under its itn/ subdirectory. ITN_MARKER is a real output of that
|
||||
# script (it prints exactly this file on success), so it can drive a make rule.
|
||||
ITN_LIB_DIR=sources/NeMo-Speech.cpp/.deps/itn/lib
|
||||
ITN_MARKER=$(ITN_LIB_DIR)/libsparrowhawk.so
|
||||
|
||||
ITN_CC?=gcc-12
|
||||
ITN_CXX?=g++-12
|
||||
|
||||
# Pin protoc to the apt one. backend/Dockerfile.golang drops protoc 27.1 into
|
||||
# /usr/local/bin, which precedes /usr/bin on PATH, while libprotobuf-dev is the
|
||||
# distro's (3.21 on noble, 3.12 on jammy). Sparrowhawk resolves protoc from PATH
|
||||
# at make time (configure.ac uses AC_CHECK_PROG, so PROTOC substitutes to the
|
||||
# bare word, and src/proto/Makefile.am invokes $(PROTOC)), and it commits no
|
||||
# pregenerated stubs, so this always runs. Code generated by 27.1 includes
|
||||
# google/protobuf/runtime_version.h and a PROTOBUF_VERSION #error guard that the
|
||||
# older headers do not have, so the mismatch breaks the build. configure honours
|
||||
# a pre-set PROTOC ("Let the user override the test"), which is what this is.
|
||||
ITN_PROTOC?=/usr/bin/protoc
|
||||
|
||||
# Text normalization is Linux-only: Sparrowhawk/OpenFST assume a GNU toolchain
|
||||
# and the gcc-12 pin has no macOS analogue. Documented gap, see the spec.
|
||||
#
|
||||
# An already-configured build tree wins over the platform default. Without that,
|
||||
# a tree configured WITH_NORM=OFF would silently try to reconfigure itself to ON
|
||||
# on the next bare `make test`, which means demanding gcc-12 from a developer who
|
||||
# deliberately built without it. An explicit WITH_NORM= on the command line still
|
||||
# overrides both, since command-line variables beat ?= assignments.
|
||||
CMAKE_CACHE=sources/NeMo-Speech.cpp/build/CMakeCache.txt
|
||||
CACHED_WITH_NORM=$(shell sed -n 's/^NEMO_SPEECH_WITH_NORM:BOOL=//p' $(CMAKE_CACHE) 2>/dev/null)
|
||||
ifeq ($(shell uname -s),Darwin)
|
||||
WITH_NORM?=OFF
|
||||
else ifneq ($(CACHED_WITH_NORM),)
|
||||
WITH_NORM?=$(CACHED_WITH_NORM)
|
||||
else
|
||||
WITH_NORM?=ON
|
||||
endif
|
||||
CMAKE_ARGS+=-DNEMO_SPEECH_WITH_NORM=$(WITH_NORM)
|
||||
|
||||
ifeq ($(BUILD_TYPE),cublas)
|
||||
CMAKE_ARGS+=-DGGML_CUDA=ON
|
||||
else ifeq ($(BUILD_TYPE),vulkan)
|
||||
CMAKE_ARGS+=-DGGML_VULKAN=ON
|
||||
else ifeq ($(BUILD_TYPE),metal)
|
||||
CMAKE_ARGS+=-DGGML_METAL=ON
|
||||
endif
|
||||
|
||||
# ggml-patches/ is a CUDA series. Every kernel it adds lives under
|
||||
# src/ggml-cuda/; the only files it touches outside that directory are enum and
|
||||
# name-table entries in include/ggml.h and src/ggml.c plus, in ggml-cpu, a
|
||||
# supports_op returning false and an abort case for the CUDA-only op. Upstream
|
||||
# agrees: its metal-* and vulkan-* CMake presets inherit the cpu-* ones, which
|
||||
# set NEMO_SPEECH_GGML_PATCHED=OFF, and every use of a patch-only symbol in the
|
||||
# ASR sources sits behind NEMO_SPEECH_FUSED_RELPOS_ATTN /
|
||||
# NEMO_SPEECH_FASTCONFORMER_CUDA_FUSIONS (both force-OFF without GGML_CUDA) or
|
||||
# behind NEMO_SPEECH_GGML_PATCHED itself, which guards a Q8_PLANAR flag write
|
||||
# that a non-CUDA buffer already throws before reaching.
|
||||
#
|
||||
# So on macOS the series buys nothing, and it cannot be applied there anyway:
|
||||
# upstream's scripts/apply-ggml-patches.sh uses mapfile, a bash 4 builtin, and
|
||||
# macOS ships bash 3.2 as the only bash on the runner's PATH. Skip the patch
|
||||
# step and tell cmake the linked ggml is stock, which is exactly upstream's own
|
||||
# Metal configuration. Linux keeps applying the series unchanged.
|
||||
ifeq ($(shell uname -s),Darwin)
|
||||
GGML_PATCHED?=OFF
|
||||
else
|
||||
GGML_PATCHED?=ON
|
||||
endif
|
||||
CMAKE_ARGS+=-DNEMO_SPEECH_GGML_PATCHED=$(GGML_PATCHED)
|
||||
|
||||
.PHONY: nemo-speech-cpp-grpc package build clean purge test all stage-libs patch-ggml engine itn
|
||||
|
||||
all: nemo-speech-cpp-grpc package
|
||||
|
||||
sources/NeMo-Speech.cpp:
|
||||
mkdir -p sources
|
||||
cd sources && git clone $(NEMO_SPEECH_REPO) NeMo-Speech.cpp
|
||||
cd sources/NeMo-Speech.cpp && git checkout $(NEMO_SPEECH_VERSION)
|
||||
# NMT links llama.cpp; ja needs open_jtalk; zh needs cppjieba. flashlight and
|
||||
# kenlm are deliberately not initialized, they are out of scope.
|
||||
cd sources/NeMo-Speech.cpp && git submodule update --init --recursive \
|
||||
ggml llama.cpp third_party/open_jtalk third_party/cppjieba third_party/cpp-httplib
|
||||
|
||||
# NEMO_SPEECH_GGML_PATCHED defaults ON and silently assumes the ggml-patches
|
||||
# series is applied. An unpatched checkout builds fine and produces wrong CUDA
|
||||
# encoder output, so a failure here must stop the build rather than warn.
|
||||
#
|
||||
# Upstream's own script is the right tool: it applies the series in filename
|
||||
# order, exits non-zero when a patch does not apply, and decides "already
|
||||
# applied" by comparing the full-series tree hash rather than a timestamp. That
|
||||
# makes it safe to run unconditionally, so there is no sentinel file to go stale
|
||||
# or to wedge the build when deleted.
|
||||
#
|
||||
# Both branches keep the order-only clone prerequisite: it is the only thing
|
||||
# that pulls sources/ in on a WITH_NORM=OFF tree, where the library rule has no
|
||||
# other prerequisite left.
|
||||
ifeq ($(GGML_PATCHED),ON)
|
||||
patch-ggml: | sources/NeMo-Speech.cpp
|
||||
cd sources/NeMo-Speech.cpp && bash scripts/apply-ggml-patches.sh
|
||||
else
|
||||
patch-ggml: | sources/NeMo-Speech.cpp
|
||||
@echo "[ggml-patch] skipped: NEMO_SPEECH_GGML_PATCHED=$(GGML_PATCHED), the series is CUDA-only"
|
||||
endif
|
||||
|
||||
# The Sparrowhawk/OpenFST text-normalization stack, as a target in its own right
|
||||
# keyed on a file the build script actually produces.
|
||||
#
|
||||
# It used to be a side effect of the runtime library rule, which meant make had
|
||||
# no idea whether it existed: once the library was up to date the script could
|
||||
# never run again, so a tree built WITH_NORM=OFF could not be moved to ON, and
|
||||
# anything that needed the ITN prefix was stuck demanding a full clean. As its
|
||||
# own rule it is built on demand, rebuilt independently, and reachable directly
|
||||
# with 'make itn'.
|
||||
#
|
||||
# OpenFST's templates ICE on gcc-13/14 at -O2, hence the gcc-12 pin for this one
|
||||
# step; the runtime itself builds with the image default compiler.
|
||||
$(ITN_MARKER): | sources/NeMo-Speech.cpp
|
||||
@command -v $(ITN_CC) >/dev/null 2>&1 && command -v $(ITN_CXX) >/dev/null 2>&1 || { \
|
||||
echo "ERROR: $(ITN_CC)/$(ITN_CXX) not found, and text normalization needs them:" >&2; \
|
||||
echo " OpenFST's templates ICE on gcc-13 and gcc-14 at -O2." >&2; \
|
||||
echo " Install them, or build this backend with WITH_NORM=OFF." >&2; \
|
||||
exit 1; }
|
||||
# configure's only gate on a preset PROTOC is test -n, so a path that does not
|
||||
# exist is accepted here and surfaces much later as a bare "No such file or
|
||||
# directory" from inside make -C src/proto. Check it up front instead.
|
||||
@command -v $(ITN_PROTOC) >/dev/null 2>&1 || { \
|
||||
echo "ERROR: protoc not found at $(ITN_PROTOC)." >&2; \
|
||||
echo " Install the protobuf-compiler package, whose protoc matches" >&2; \
|
||||
echo " the libprotobuf-dev headers Sparrowhawk compiles against, or" >&2; \
|
||||
echo " point this at a matching one with ITN_PROTOC=/path/to/protoc." >&2; \
|
||||
exit 1; }
|
||||
cd sources/NeMo-Speech.cpp && CC=$(ITN_CC) CXX=$(ITN_CXX) PROTOC=$(ITN_PROTOC) \
|
||||
JOBS=$(JOBS) scripts/build_itn_deps.sh
|
||||
|
||||
itn: $(ITN_MARKER)
|
||||
|
||||
# Only a WITH_NORM=ON build needs the ITN stack, and it must exist before cmake
|
||||
# configures, since the WITH_NORM cmake block find_library()s into the prefix
|
||||
# with REQUIRED.
|
||||
ifeq ($(WITH_NORM),ON)
|
||||
NEMO_RUNTIME_PREREQS=$(ITN_MARKER)
|
||||
endif
|
||||
|
||||
# Upstream sets CMAKE_LIBRARY_OUTPUT_DIRECTORY to ${CMAKE_BINARY_DIR}/bin, so the
|
||||
# shared objects land in build/bin rather than at the top of the build tree.
|
||||
#
|
||||
# patch-ggml is order-only: it is phony and therefore always runs, but an
|
||||
# order-only prerequisite does not mark this target out of date, so an
|
||||
# already-built tree is not relinked on every invocation.
|
||||
sources/NeMo-Speech.cpp/build/bin/libnemo_speech_asr_c.so: $(NEMO_RUNTIME_PREREQS) | patch-ggml
|
||||
cd sources/NeMo-Speech.cpp && cmake -B build -G Ninja $(CMAKE_ARGS)
|
||||
cd sources/NeMo-Speech.cpp && cmake --build build -j$(JOBS)
|
||||
|
||||
# Stage the runtime next to the Go sources so purego.Dlopen finds it during
|
||||
# local development and so package.sh has a single directory to bundle from.
|
||||
#
|
||||
# ASR and NMT build a dedicated _c shared object that links the C++ implementation
|
||||
# in privately. TTS does not: upstream compiles its c_api.cpp straight into
|
||||
# libnemo_speech_tts and only aliases the nemo_speech_tts_c CMake target, so the
|
||||
# TTS C ABI ships without the _c suffix.
|
||||
stage-libs: sources/NeMo-Speech.cpp/build/bin/libnemo_speech_asr_c.so
|
||||
# -a keeps the SOVERSION symlink a symlink instead of duplicating the payload.
|
||||
cp -af sources/NeMo-Speech.cpp/build/bin/libnemo_speech_asr_c.* .
|
||||
cp -af sources/NeMo-Speech.cpp/build/bin/libnemo_speech_tts.* .
|
||||
cp -af sources/NeMo-Speech.cpp/build/bin/libnemo_speech_nmt_c.* .
|
||||
# The _c libraries are thin ABI shims with a DT_NEEDED on the C++
|
||||
# implementation DSO, so dlopen fails without these next to them. TTS needs
|
||||
# no counterpart, its implementation and ABI live in the same object.
|
||||
cp -af sources/NeMo-Speech.cpp/build/bin/libnemo_speech_asr.* .
|
||||
cp -af sources/NeMo-Speech.cpp/build/bin/libnemo_speech_nmt.* .
|
||||
# nemo_speech_text_normalization is STATIC but links sparrowhawk, fstfar and
|
||||
# fst PUBLIC, so those become DT_NEEDED on libnemo_speech_asr.so. They live in
|
||||
# a project-local prefix that nothing else on the system provides, so without
|
||||
# staging them here the packaged backend cannot dlopen at all.
|
||||
#
|
||||
# Keyed on the prefix existing rather than on WITH_NORM, so this stages what
|
||||
# the tree actually built. A WITH_NORM=ON build cannot reach here without the
|
||||
# prefix (the library rule takes ITN_MARKER as a prerequisite), and if a
|
||||
# library that needs Sparrowhawk somehow arrives unstaged, package.sh's
|
||||
# closure guard fails the build rather than shipping it.
|
||||
@if [ -d "$(ITN_LIB_DIR)" ]; then \
|
||||
echo "cp -af $(ITN_LIB_DIR)/*.so* ."; \
|
||||
cp -af $(ITN_LIB_DIR)/*.so* .; \
|
||||
fi
|
||||
|
||||
## Builds the native runtime and stops short of the Go binary. Everything it
|
||||
## touches lives under sources/, a clone pinned by NEMO_SPEECH_VERSION, so
|
||||
## nothing here can observe a change elsewhere in the LocalAI tree.
|
||||
## Dockerfile.golang calls this from a layer that copies in this directory and
|
||||
## nothing else, which keeps the multi-minute ggml/llama.cpp compile in the
|
||||
## registry layer cache across builds whose only change is on the Go side.
|
||||
## Without it that prebuild is skipped and a CUDA build recompiles all of
|
||||
## upstream on every Go-side edit. See .agents/ci-caching.md.
|
||||
engine: stage-libs
|
||||
|
||||
nemo-speech-cpp-grpc: stage-libs
|
||||
# CGO_ENABLED=0 matches whisper / parakeet-cpp / omnivoice-cpp: the runtime is
|
||||
# reached through purego.Dlopen, not cgo, and a static binary is what lets
|
||||
# run.sh route execution through the packaged lib/ld.so.
|
||||
CGO_ENABLED=0 $(GOCMD) build -tags "$(GO_TAGS)" -o nemo-speech-cpp-grpc .
|
||||
|
||||
# The dlopen tests need the staged shared objects on the loader path, the same
|
||||
# way parakeet-cpp sets it up. Depends on stage-libs so that path is not an
|
||||
# empty directory on a clean tree, which would fail the tests confusingly.
|
||||
#
|
||||
# NEMO_SPEECH_REQUIRE_LIBS turns a missing library from a skip into a failure.
|
||||
# The ABI specs are the only thing standing between this backend and silent
|
||||
# memory corruption, so a run that reaches them and quietly skips them is worse
|
||||
# than one that fails: it reports green having checked nothing.
|
||||
test: stage-libs
|
||||
NEMO_SPEECH_REQUIRE_LIBS=1 LD_LIBRARY_PATH=$(CURDIR):$$LD_LIBRARY_PATH $(GOCMD) test ./... -count=1
|
||||
|
||||
package: nemo-speech-cpp-grpc
|
||||
bash package.sh
|
||||
|
||||
# What backend/Dockerfile.golang invokes. It must leave both the binary and a
|
||||
# populated package/ behind, because the final image stage copies package/.
|
||||
build: package
|
||||
|
||||
clean:
|
||||
# Every .so here is staged output (nemo runtime plus, on a WITH_NORM build,
|
||||
# the ITN stack), and the SOVERSION suffix means the payload is *.so.1, so
|
||||
# the globs have to reach past the .so.
|
||||
rm -f nemo-speech-cpp-grpc
|
||||
rm -f *.so *.so.* *.dylib
|
||||
rm -rf package
|
||||
rm -rf sources/NeMo-Speech.cpp/build
|
||||
|
||||
purge: clean
|
||||
rm -rf sources
|
||||
@@ -1,423 +0,0 @@
|
||||
package main
|
||||
|
||||
// purego binds by name at runtime and the config structs cross the ABI by
|
||||
// pointer, so neither a renamed symbol nor a mis-laid-out mirror struct is
|
||||
// visible to the compiler or the linker. Everything here is transcribed from
|
||||
// sources/NeMo-Speech.cpp/include/nemo_speech/{asr,diar,tts,nmt}.h, and
|
||||
// abi_test.go asserts it against the real shared objects.
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"unsafe"
|
||||
|
||||
"github.com/ebitengine/purego"
|
||||
)
|
||||
|
||||
var (
|
||||
asrLib uintptr
|
||||
ttsLib uintptr
|
||||
nmtLib uintptr
|
||||
)
|
||||
|
||||
// ---- ASR ----
|
||||
|
||||
var (
|
||||
ASRCreate func(cfg unsafe.Pointer, out *uintptr) int32
|
||||
ASRDestroy func(recognizer uintptr)
|
||||
ASRRecognizeF32 func(recognizer uintptr, options unsafe.Pointer, samples *float32, nSamples uint64, sampleRate int32, out *uintptr) int32
|
||||
ASRStreamingRecognize func(recognizer uintptr, options unsafe.Pointer, out *uintptr) int32
|
||||
ASRStreamPushF32 func(stream uintptr, samples *float32, nSamples uint64, sampleRate int32) int32
|
||||
ASRStreamForceEndpoint func(stream uintptr) int32
|
||||
ASRStreamFinish func(stream uintptr) int32
|
||||
ASRStreamNext func(stream uintptr, out *uintptr) int32
|
||||
ASRStreamClose func(stream uintptr)
|
||||
ASRRecognitionOptionsDef func() cASRRecognitionOptions
|
||||
|
||||
ASRResultIsFinal func(result uintptr) bool
|
||||
ASRResultAudioProcessed func(result uintptr) float32
|
||||
ASRResultAlternativeCount func(result uintptr) uint64
|
||||
ASRResultTranscript func(result uintptr, alt uint64) string
|
||||
ASRResultConfidence func(result uintptr, alt uint64) float32
|
||||
ASRResultWordCount func(result uintptr, alt uint64) uint64
|
||||
ASRResultWordText func(result uintptr, alt, i uint64) string
|
||||
ASRResultWordStartTime func(result uintptr, alt, i uint64) int32
|
||||
ASRResultWordEndTime func(result uintptr, alt, i uint64) int32
|
||||
ASRResultWordConfidence func(result uintptr, alt, i uint64) float32
|
||||
ASRResultWordSpeakerTag func(result uintptr, alt, i uint64) int32
|
||||
ASRResultLanguageCount func(result uintptr, alt uint64) uint64
|
||||
ASRResultLanguageCode func(result uintptr, alt, i uint64) string
|
||||
ASRResultDestroy func(result uintptr)
|
||||
|
||||
ASRLastError func() string
|
||||
ASRVersion func() string
|
||||
)
|
||||
|
||||
// ---- Diarization (exported from the ASR library) ----
|
||||
|
||||
var (
|
||||
DiarCreate func(cfg unsafe.Pointer, out *uintptr) int32
|
||||
DiarDestroy func(model uintptr)
|
||||
DiarNumSpeakers func(model uintptr) int32
|
||||
DiarSecondsPerFrame func(model uintptr) float64
|
||||
DiarStreamOpen func(model uintptr, out *uintptr) int32
|
||||
DiarStreamPushF32 func(stream uintptr, samples *float32, nSamples uint64, sampleRate int32) int32
|
||||
DiarStreamFinish func(stream uintptr) int32
|
||||
DiarStreamClose func(stream uintptr)
|
||||
// cfg is the optional nemo_speech_diar_segmentation_config (NULL = library
|
||||
// defaults). The two-call count-then-fill pattern is documented on the C
|
||||
// declaration in diar.h.
|
||||
DiarSegments func(stream uintptr, cfg unsafe.Pointer, out unsafe.Pointer, capacity uint64, count *uint64) int32
|
||||
)
|
||||
|
||||
// ---- TTS ----
|
||||
|
||||
var (
|
||||
TTSCreate func(cfg unsafe.Pointer, out *uintptr) int32
|
||||
TTSDestroy func(synthesizer uintptr)
|
||||
TTSSampleRate func(synthesizer uintptr) int32
|
||||
TTSSpeakerCount func(synthesizer uintptr) int32
|
||||
TTSSpeakerName func(synthesizer uintptr, i uint64) string
|
||||
TTSSynthesizeText func(synthesizer uintptr, options unsafe.Pointer, text string, callback uintptr, userData uintptr, statsOut unsafe.Pointer) int32
|
||||
TTSRuntimeConfigDefault func() cTTSRuntimeConfig
|
||||
TTSSynthesisOptionsDefault func() cTTSSynthesisOptions
|
||||
TTSLastError func() string
|
||||
TTSVersion func() string
|
||||
)
|
||||
|
||||
// ---- NMT ----
|
||||
|
||||
var (
|
||||
NMTCreate func(cfg unsafe.Pointer, out *uintptr) int32
|
||||
NMTDestroy func(translator uintptr)
|
||||
NMTTranslate func(translator uintptr, texts *uintptr, nTexts uint64, source, target string, out *uintptr) int32
|
||||
NMTResultCount func(result uintptr) uint64
|
||||
NMTResultText func(result uintptr, i uint64) string
|
||||
NMTResultLanguage func(result uintptr, i uint64) string
|
||||
NMTResultDestroy func(result uintptr)
|
||||
NMTLastError func() string
|
||||
NMTVersion func() string
|
||||
)
|
||||
|
||||
// ---- C struct mirrors ----
|
||||
//
|
||||
// Each mirrors a struct in include/nemo_speech/*.h field for field. The leading
|
||||
// Size field is the C `size_t size` the runtime validates against its own
|
||||
// sizeof, which is what makes a layout mismatch detectable at runtime instead
|
||||
// of silently corrupting memory. Blank fields are System V AMD64 / AAPCS64
|
||||
// padding: C inserts it implicitly, Go does not, so it has to be written out.
|
||||
// See abi_test.go, which pins both every total size and every field offset.
|
||||
|
||||
type cASRBackendConfig struct {
|
||||
Size uintptr
|
||||
GPU int32
|
||||
_ [4]byte // trailing pad to the struct's 8-byte alignment
|
||||
}
|
||||
|
||||
type cASRModelConfig struct {
|
||||
Size uintptr
|
||||
Path uintptr
|
||||
Name uintptr
|
||||
}
|
||||
|
||||
type cASRVADConfig struct {
|
||||
Size uintptr
|
||||
ModelPath uintptr
|
||||
EnableMasking bool
|
||||
_ [3]byte
|
||||
Onset float32
|
||||
Offset float32
|
||||
_ [4]byte
|
||||
}
|
||||
|
||||
type cASRPostprocConfig struct {
|
||||
Size uintptr
|
||||
ProfanityListPath uintptr
|
||||
ITNModelDir uintptr
|
||||
PNCModelPath uintptr
|
||||
}
|
||||
|
||||
type cASRDiarConfig struct {
|
||||
Size uintptr
|
||||
ModelPath uintptr
|
||||
ChunkFrames int32
|
||||
RightContextFrames int32
|
||||
LeftContextFrames int32
|
||||
FIFOFrames int32
|
||||
SpkcacheFrames int32
|
||||
UpdatePeriodFrames int32
|
||||
}
|
||||
|
||||
type cASRRecognizerConfig struct {
|
||||
Size uintptr
|
||||
Backend uintptr
|
||||
Model uintptr
|
||||
Streaming uintptr
|
||||
Decoder uintptr
|
||||
VAD uintptr
|
||||
Endpointing uintptr
|
||||
Postproc uintptr
|
||||
Diar uintptr
|
||||
Batching uintptr
|
||||
}
|
||||
|
||||
type cASRRecognitionOptions struct {
|
||||
Size uintptr
|
||||
RequestID uintptr
|
||||
LanguageCode uintptr
|
||||
InterimResults bool
|
||||
EnableWordTimeOffsets bool
|
||||
EnableAutomaticPunctuation bool
|
||||
VerbatimTranscripts bool
|
||||
ProfanityFilter bool
|
||||
_ [3]byte
|
||||
StopHistoryEouMs int32
|
||||
_ [4]byte
|
||||
SpeechContexts uintptr
|
||||
SpeechContextCount uintptr
|
||||
MaxAlternatives int32
|
||||
EnableSpeakerDiarization bool
|
||||
_ [3]byte
|
||||
MaxSpeakerCount int32
|
||||
_ [4]byte
|
||||
}
|
||||
|
||||
// cDiarModelConfig mirrors nemo_speech_diar_model_config (diar.h). This is the
|
||||
// standalone Sortformer pipeline's own config and is NOT cASRDiarConfig, which
|
||||
// is the diarizer attached to a recognizer: this one carries gpu and preset,
|
||||
// that one does not.
|
||||
//
|
||||
// The six frame counts are sentinel-sensitive. src/asr/c_api.cpp applies each
|
||||
// one only when it is > 0, EXCEPT left_context_frames, which it applies when it
|
||||
// is >= 0. A zero-valued struct would therefore pin the left context to 0
|
||||
// rather than leave the preset's value alone, so loadDiarizer writes -1 into
|
||||
// all six.
|
||||
type cDiarModelConfig struct {
|
||||
Size uintptr
|
||||
ModelPath uintptr
|
||||
GPU int32
|
||||
_ [4]byte // pad to the alignment of the pointer that follows
|
||||
Preset uintptr
|
||||
// Encoder-frame geometry overrides, applied on top of the preset.
|
||||
ChunkFrames int32
|
||||
RightContextFrames int32
|
||||
LeftContextFrames int32
|
||||
FIFOFrames int32
|
||||
SpkcacheFrames int32
|
||||
UpdatePeriodFrames int32
|
||||
}
|
||||
|
||||
// cDiarSegmentationConfig mirrors nemo_speech_diar_segmentation_config
|
||||
// (diar.h): the NeMo ts_vad postprocessing applied when turning per-frame
|
||||
// speaker probabilities into segments.
|
||||
//
|
||||
// onset and offset are float, the four durations are double. That mixture is
|
||||
// the whole reason this mirror needs its offsets pinned: writing all six as
|
||||
// float32 or all six as float64 both produce a struct C would read shifted.
|
||||
type cDiarSegmentationConfig struct {
|
||||
Size uintptr
|
||||
Onset float32
|
||||
Offset float32
|
||||
PadOnsetSec float64
|
||||
PadOffsetSec float64
|
||||
MinGapSec float64
|
||||
MinDurationSec float64
|
||||
}
|
||||
|
||||
// cDiarSegment mirrors nemo_speech_diar_segment (diar.h), the element type
|
||||
// nemo_speech_diar_segments fills.
|
||||
//
|
||||
// It has no leading size field: unlike the config structs it travels from C to
|
||||
// Go, so there is no caller-declared size for the runtime to validate against.
|
||||
// The times are already SECONDS (double), not frame indices, so nothing here
|
||||
// needs the model's seconds-per-frame to be interpreted. Speaker is 1-based,
|
||||
// matching WordInfo.speaker_tag on the ASR surface.
|
||||
type cDiarSegment struct {
|
||||
StartTime float64
|
||||
EndTime float64
|
||||
Speaker int32
|
||||
_ [4]byte // trailing pad to the struct's 8-byte alignment
|
||||
}
|
||||
|
||||
type cTTSModelConfig struct {
|
||||
Size uintptr
|
||||
MagpieModel uintptr
|
||||
CodecModel uintptr
|
||||
TokenizerModelDir uintptr
|
||||
TextNormalizerModelDir uintptr
|
||||
}
|
||||
|
||||
// cTTSRuntimeConfig mirrors nemo_speech_tts_runtime_config. The four backend /
|
||||
// mode fields are C enums, which this toolchain lays out as int32.
|
||||
type cTTSRuntimeConfig struct {
|
||||
Size uintptr
|
||||
Speaker int32
|
||||
Threads int32
|
||||
CodecThreads int32
|
||||
Seed int32
|
||||
Steps int32
|
||||
TopK int32
|
||||
ChunkFrames int32
|
||||
CodecQueueDepth int32
|
||||
CodecHistoryFrames int32
|
||||
CodecFutureFrames int32
|
||||
WindowMs int32
|
||||
Temperature float32
|
||||
OverrideTemperature bool
|
||||
_ [3]byte
|
||||
CFGScale float32
|
||||
OverrideCFGScale bool
|
||||
UseCFG bool
|
||||
UseLocalTransformer bool
|
||||
UseKVCache bool
|
||||
UseStatefulCodec bool
|
||||
CodecCPU bool
|
||||
FlushPartialChunk bool
|
||||
Verbose bool
|
||||
LTBackend int32
|
||||
SamplingBackend int32
|
||||
UMAMode int32
|
||||
LongformMode int32
|
||||
LTFP32 bool
|
||||
_ [7]byte
|
||||
}
|
||||
|
||||
type cTTSSynthesizerConfig struct {
|
||||
Size uintptr
|
||||
Model uintptr
|
||||
Runtime uintptr
|
||||
DefaultLanguageCode uintptr
|
||||
DefaultVoiceName uintptr
|
||||
}
|
||||
|
||||
type cTTSSynthesisOptions struct {
|
||||
Size uintptr
|
||||
RequestID uintptr
|
||||
LanguageCode uintptr
|
||||
Speaker int32
|
||||
Seed int32
|
||||
Steps int32
|
||||
TopK int32
|
||||
Temperature float32
|
||||
OverrideTemperature bool
|
||||
_ [3]byte
|
||||
CFGScale float32
|
||||
OverrideCFGScale bool
|
||||
_ [3]byte
|
||||
VoiceName uintptr
|
||||
OutputSampleRate int32
|
||||
_ [4]byte
|
||||
}
|
||||
|
||||
type cNMTBackendConfig struct {
|
||||
Size uintptr
|
||||
GPU int32
|
||||
_ [4]byte
|
||||
}
|
||||
|
||||
type cNMTModelConfig struct {
|
||||
Size uintptr
|
||||
Path uintptr
|
||||
NCtx int32
|
||||
_ [4]byte
|
||||
}
|
||||
|
||||
type cNMTTranslatorConfig struct {
|
||||
Size uintptr
|
||||
Backend uintptr
|
||||
Model uintptr
|
||||
Generation uintptr
|
||||
Pool uintptr
|
||||
}
|
||||
|
||||
// symbol pairs a Go function pointer with its exported C name. Keeping the
|
||||
// name next to the var means `nm -D libnemo_speech_asr_c.so.1 | grep nemo_speech`
|
||||
// is enough to spot drift after a pin bump.
|
||||
type symbol struct {
|
||||
fn any
|
||||
name string
|
||||
lib *uintptr
|
||||
}
|
||||
|
||||
func symbols() []symbol {
|
||||
return []symbol{
|
||||
{&ASRCreate, "nemo_speech_asr_create", &asrLib},
|
||||
{&ASRDestroy, "nemo_speech_asr_destroy", &asrLib},
|
||||
{&ASRRecognizeF32, "nemo_speech_asr_recognize_f32", &asrLib},
|
||||
{&ASRStreamingRecognize, "nemo_speech_asr_streaming_recognize", &asrLib},
|
||||
{&ASRStreamPushF32, "nemo_speech_asr_stream_push_f32", &asrLib},
|
||||
{&ASRStreamForceEndpoint, "nemo_speech_asr_stream_force_endpoint", &asrLib},
|
||||
{&ASRStreamFinish, "nemo_speech_asr_stream_finish", &asrLib},
|
||||
{&ASRStreamNext, "nemo_speech_asr_stream_next", &asrLib},
|
||||
{&ASRStreamClose, "nemo_speech_asr_stream_close", &asrLib},
|
||||
{&ASRRecognitionOptionsDef, "nemo_speech_asr_recognition_options_default", &asrLib},
|
||||
{&ASRResultIsFinal, "nemo_speech_asr_result_is_final", &asrLib},
|
||||
{&ASRResultAudioProcessed, "nemo_speech_asr_result_audio_processed", &asrLib},
|
||||
{&ASRResultAlternativeCount, "nemo_speech_asr_result_alternative_count", &asrLib},
|
||||
{&ASRResultTranscript, "nemo_speech_asr_result_transcript", &asrLib},
|
||||
{&ASRResultConfidence, "nemo_speech_asr_result_confidence", &asrLib},
|
||||
{&ASRResultWordCount, "nemo_speech_asr_result_word_count", &asrLib},
|
||||
{&ASRResultWordText, "nemo_speech_asr_result_word_text", &asrLib},
|
||||
{&ASRResultWordStartTime, "nemo_speech_asr_result_word_start_time", &asrLib},
|
||||
{&ASRResultWordEndTime, "nemo_speech_asr_result_word_end_time", &asrLib},
|
||||
{&ASRResultWordConfidence, "nemo_speech_asr_result_word_confidence", &asrLib},
|
||||
{&ASRResultWordSpeakerTag, "nemo_speech_asr_result_word_speaker_tag", &asrLib},
|
||||
{&ASRResultLanguageCount, "nemo_speech_asr_result_language_count", &asrLib},
|
||||
{&ASRResultLanguageCode, "nemo_speech_asr_result_language_code", &asrLib},
|
||||
{&ASRResultDestroy, "nemo_speech_asr_result_destroy", &asrLib},
|
||||
{&ASRLastError, "nemo_speech_asr_last_error", &asrLib},
|
||||
{&ASRVersion, "nemo_speech_asr_version", &asrLib},
|
||||
|
||||
{&DiarCreate, "nemo_speech_diar_create", &asrLib},
|
||||
{&DiarDestroy, "nemo_speech_diar_destroy", &asrLib},
|
||||
{&DiarNumSpeakers, "nemo_speech_diar_num_speakers", &asrLib},
|
||||
{&DiarSecondsPerFrame, "nemo_speech_diar_seconds_per_frame", &asrLib},
|
||||
{&DiarStreamOpen, "nemo_speech_diar_stream_open", &asrLib},
|
||||
{&DiarStreamPushF32, "nemo_speech_diar_stream_push_f32", &asrLib},
|
||||
{&DiarStreamFinish, "nemo_speech_diar_stream_finish", &asrLib},
|
||||
{&DiarStreamClose, "nemo_speech_diar_stream_close", &asrLib},
|
||||
{&DiarSegments, "nemo_speech_diar_segments", &asrLib},
|
||||
|
||||
{&TTSCreate, "nemo_speech_tts_create", &ttsLib},
|
||||
{&TTSDestroy, "nemo_speech_tts_destroy", &ttsLib},
|
||||
{&TTSSampleRate, "nemo_speech_tts_sample_rate", &ttsLib},
|
||||
{&TTSSpeakerCount, "nemo_speech_tts_speaker_count", &ttsLib},
|
||||
{&TTSSpeakerName, "nemo_speech_tts_speaker_name", &ttsLib},
|
||||
{&TTSSynthesizeText, "nemo_speech_tts_synthesize_text", &ttsLib},
|
||||
{&TTSRuntimeConfigDefault, "nemo_speech_tts_runtime_config_default", &ttsLib},
|
||||
{&TTSSynthesisOptionsDefault, "nemo_speech_tts_synthesis_options_default", &ttsLib},
|
||||
{&TTSLastError, "nemo_speech_tts_last_error", &ttsLib},
|
||||
{&TTSVersion, "nemo_speech_tts_version", &ttsLib},
|
||||
|
||||
{&NMTCreate, "nemo_speech_nmt_create", &nmtLib},
|
||||
{&NMTDestroy, "nemo_speech_nmt_destroy", &nmtLib},
|
||||
{&NMTTranslate, "nemo_speech_nmt_translate", &nmtLib},
|
||||
{&NMTResultCount, "nemo_speech_nmt_result_count", &nmtLib},
|
||||
{&NMTResultText, "nemo_speech_nmt_result_text", &nmtLib},
|
||||
{&NMTResultLanguage, "nemo_speech_nmt_result_language", &nmtLib},
|
||||
{&NMTResultDestroy, "nemo_speech_nmt_result_destroy", &nmtLib},
|
||||
{&NMTLastError, "nemo_speech_nmt_last_error", &nmtLib},
|
||||
{&NMTVersion, "nemo_speech_nmt_version", &nmtLib},
|
||||
}
|
||||
}
|
||||
|
||||
// registerSymbols binds every entry point. purego panics on a missing symbol,
|
||||
// so this recovers and returns the offending name: after an upstream pin bump a
|
||||
// rename must fail loudly at startup, not at first inference.
|
||||
func registerSymbols() error {
|
||||
for _, s := range symbols() {
|
||||
if err := registerOne(s); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func registerOne(s symbol) (err error) {
|
||||
defer func() {
|
||||
if r := recover(); r != nil {
|
||||
err = fmt.Errorf("nemo-speech-cpp: binding %q: %v", s.name, r)
|
||||
}
|
||||
}()
|
||||
purego.RegisterLibFunc(s.fn, *s.lib, s.name)
|
||||
return nil
|
||||
}
|
||||
@@ -1,317 +0,0 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"os"
|
||||
"unsafe"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
)
|
||||
|
||||
// requireLibs reports whether a missing shared library must fail the specs
|
||||
// instead of skipping them.
|
||||
//
|
||||
// librariesPresent stats bare filenames relative to the working directory,
|
||||
// while openLibraries resolves them through the loader search path, so the two
|
||||
// can legitimately disagree. Under `make test` that is harmless because the
|
||||
// stage-libs prerequisite puts the .so files in the working directory, but any
|
||||
// other invocation would skip every library-backed spec and still report a
|
||||
// green run. The Makefile sets NEMO_SPEECH_REQUIRE_LIBS=1 so no CI path can
|
||||
// pass on a silent skip; leaving it unset keeps the pure-Go specs runnable on a
|
||||
// checkout with no build.
|
||||
func requireLibs() bool {
|
||||
return os.Getenv("NEMO_SPEECH_REQUIRE_LIBS") == "1"
|
||||
}
|
||||
|
||||
// librariesPresent reports whether a local build is available to bind against.
|
||||
func librariesPresent() bool {
|
||||
for _, n := range []string{
|
||||
libraryName("NEMO_SPEECH_ASR_LIBRARY", "libnemo_speech_asr_c"),
|
||||
libraryName("NEMO_SPEECH_TTS_LIBRARY", "libnemo_speech_tts"),
|
||||
libraryName("NEMO_SPEECH_NMT_LIBRARY", "libnemo_speech_nmt_c"),
|
||||
} {
|
||||
if _, err := os.Stat(n); err != nil {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
// layout is one expected number transcribed from the C headers.
|
||||
type layout struct {
|
||||
what string
|
||||
got uintptr
|
||||
want uintptr
|
||||
}
|
||||
|
||||
// The `want` column is what a C compiler reports for the structs in
|
||||
// include/nemo_speech/{asr,diar,tts,nmt}.h under the System V AMD64 / AAPCS64 rules
|
||||
// both supported targets follow. Regenerate after an upstream pin bump with a
|
||||
// throwaway program over the installed headers:
|
||||
//
|
||||
// printf('SIZE %%zu\n', sizeof(nemo_speech_asr_recognition_options));
|
||||
// printf('OFF %%zu\n', offsetof(nemo_speech_asr_recognition_options, max_speaker_count));
|
||||
//
|
||||
// Sizes alone are not enough: two padding mistakes can cancel out and leave the
|
||||
// total unchanged while every field between them reads from the wrong offset,
|
||||
// so each mirror pins its field offsets too.
|
||||
func structSizes() []layout {
|
||||
return []layout{
|
||||
{"cASRBackendConfig", unsafe.Sizeof(cASRBackendConfig{}), 16},
|
||||
{"cASRModelConfig", unsafe.Sizeof(cASRModelConfig{}), 24},
|
||||
{"cASRVADConfig", unsafe.Sizeof(cASRVADConfig{}), 32},
|
||||
{"cASRPostprocConfig", unsafe.Sizeof(cASRPostprocConfig{}), 32},
|
||||
{"cASRDiarConfig", unsafe.Sizeof(cASRDiarConfig{}), 40},
|
||||
{"cASRRecognizerConfig", unsafe.Sizeof(cASRRecognizerConfig{}), 80},
|
||||
{"cASRRecognitionOptions", unsafe.Sizeof(cASRRecognitionOptions{}), 72},
|
||||
{"cDiarModelConfig", unsafe.Sizeof(cDiarModelConfig{}), 56},
|
||||
{"cDiarSegmentationConfig", unsafe.Sizeof(cDiarSegmentationConfig{}), 48},
|
||||
{"cDiarSegment", unsafe.Sizeof(cDiarSegment{}), 24},
|
||||
{"cTTSModelConfig", unsafe.Sizeof(cTTSModelConfig{}), 40},
|
||||
{"cTTSRuntimeConfig", unsafe.Sizeof(cTTSRuntimeConfig{}), 96},
|
||||
{"cTTSSynthesizerConfig", unsafe.Sizeof(cTTSSynthesizerConfig{}), 40},
|
||||
{"cTTSSynthesisOptions", unsafe.Sizeof(cTTSSynthesisOptions{}), 72},
|
||||
{"cNMTBackendConfig", unsafe.Sizeof(cNMTBackendConfig{}), 16},
|
||||
{"cNMTModelConfig", unsafe.Sizeof(cNMTModelConfig{}), 24},
|
||||
{"cNMTTranslatorConfig", unsafe.Sizeof(cNMTTranslatorConfig{}), 40},
|
||||
}
|
||||
}
|
||||
|
||||
func structOffsets() []layout {
|
||||
return []layout{
|
||||
{"cASRBackendConfig.GPU", unsafe.Offsetof(cASRBackendConfig{}.GPU), 8},
|
||||
|
||||
{"cASRModelConfig.Path", unsafe.Offsetof(cASRModelConfig{}.Path), 8},
|
||||
{"cASRModelConfig.Name", unsafe.Offsetof(cASRModelConfig{}.Name), 16},
|
||||
|
||||
{"cASRVADConfig.ModelPath", unsafe.Offsetof(cASRVADConfig{}.ModelPath), 8},
|
||||
{"cASRVADConfig.EnableMasking", unsafe.Offsetof(cASRVADConfig{}.EnableMasking), 16},
|
||||
{"cASRVADConfig.Onset", unsafe.Offsetof(cASRVADConfig{}.Onset), 20},
|
||||
{"cASRVADConfig.Offset", unsafe.Offsetof(cASRVADConfig{}.Offset), 24},
|
||||
|
||||
{"cASRPostprocConfig.ProfanityListPath", unsafe.Offsetof(cASRPostprocConfig{}.ProfanityListPath), 8},
|
||||
{"cASRPostprocConfig.ITNModelDir", unsafe.Offsetof(cASRPostprocConfig{}.ITNModelDir), 16},
|
||||
{"cASRPostprocConfig.PNCModelPath", unsafe.Offsetof(cASRPostprocConfig{}.PNCModelPath), 24},
|
||||
|
||||
{"cASRDiarConfig.ModelPath", unsafe.Offsetof(cASRDiarConfig{}.ModelPath), 8},
|
||||
{"cASRDiarConfig.ChunkFrames", unsafe.Offsetof(cASRDiarConfig{}.ChunkFrames), 16},
|
||||
{"cASRDiarConfig.RightContextFrames", unsafe.Offsetof(cASRDiarConfig{}.RightContextFrames), 20},
|
||||
{"cASRDiarConfig.LeftContextFrames", unsafe.Offsetof(cASRDiarConfig{}.LeftContextFrames), 24},
|
||||
{"cASRDiarConfig.FIFOFrames", unsafe.Offsetof(cASRDiarConfig{}.FIFOFrames), 28},
|
||||
{"cASRDiarConfig.SpkcacheFrames", unsafe.Offsetof(cASRDiarConfig{}.SpkcacheFrames), 32},
|
||||
{"cASRDiarConfig.UpdatePeriodFrames", unsafe.Offsetof(cASRDiarConfig{}.UpdatePeriodFrames), 36},
|
||||
|
||||
{"cASRRecognizerConfig.Backend", unsafe.Offsetof(cASRRecognizerConfig{}.Backend), 8},
|
||||
{"cASRRecognizerConfig.Model", unsafe.Offsetof(cASRRecognizerConfig{}.Model), 16},
|
||||
{"cASRRecognizerConfig.Streaming", unsafe.Offsetof(cASRRecognizerConfig{}.Streaming), 24},
|
||||
{"cASRRecognizerConfig.Decoder", unsafe.Offsetof(cASRRecognizerConfig{}.Decoder), 32},
|
||||
{"cASRRecognizerConfig.VAD", unsafe.Offsetof(cASRRecognizerConfig{}.VAD), 40},
|
||||
{"cASRRecognizerConfig.Endpointing", unsafe.Offsetof(cASRRecognizerConfig{}.Endpointing), 48},
|
||||
{"cASRRecognizerConfig.Postproc", unsafe.Offsetof(cASRRecognizerConfig{}.Postproc), 56},
|
||||
{"cASRRecognizerConfig.Diar", unsafe.Offsetof(cASRRecognizerConfig{}.Diar), 64},
|
||||
{"cASRRecognizerConfig.Batching", unsafe.Offsetof(cASRRecognizerConfig{}.Batching), 72},
|
||||
|
||||
{"cASRRecognitionOptions.RequestID", unsafe.Offsetof(cASRRecognitionOptions{}.RequestID), 8},
|
||||
{"cASRRecognitionOptions.LanguageCode", unsafe.Offsetof(cASRRecognitionOptions{}.LanguageCode), 16},
|
||||
{"cASRRecognitionOptions.InterimResults", unsafe.Offsetof(cASRRecognitionOptions{}.InterimResults), 24},
|
||||
{"cASRRecognitionOptions.EnableWordTimeOffsets", unsafe.Offsetof(cASRRecognitionOptions{}.EnableWordTimeOffsets), 25},
|
||||
{"cASRRecognitionOptions.EnableAutomaticPunctuation", unsafe.Offsetof(cASRRecognitionOptions{}.EnableAutomaticPunctuation), 26},
|
||||
{"cASRRecognitionOptions.VerbatimTranscripts", unsafe.Offsetof(cASRRecognitionOptions{}.VerbatimTranscripts), 27},
|
||||
{"cASRRecognitionOptions.ProfanityFilter", unsafe.Offsetof(cASRRecognitionOptions{}.ProfanityFilter), 28},
|
||||
{"cASRRecognitionOptions.StopHistoryEouMs", unsafe.Offsetof(cASRRecognitionOptions{}.StopHistoryEouMs), 32},
|
||||
{"cASRRecognitionOptions.SpeechContexts", unsafe.Offsetof(cASRRecognitionOptions{}.SpeechContexts), 40},
|
||||
{"cASRRecognitionOptions.SpeechContextCount", unsafe.Offsetof(cASRRecognitionOptions{}.SpeechContextCount), 48},
|
||||
{"cASRRecognitionOptions.MaxAlternatives", unsafe.Offsetof(cASRRecognitionOptions{}.MaxAlternatives), 56},
|
||||
{"cASRRecognitionOptions.EnableSpeakerDiarization", unsafe.Offsetof(cASRRecognitionOptions{}.EnableSpeakerDiarization), 60},
|
||||
{"cASRRecognitionOptions.MaxSpeakerCount", unsafe.Offsetof(cASRRecognitionOptions{}.MaxSpeakerCount), 64},
|
||||
|
||||
{"cDiarModelConfig.ModelPath", unsafe.Offsetof(cDiarModelConfig{}.ModelPath), 8},
|
||||
{"cDiarModelConfig.GPU", unsafe.Offsetof(cDiarModelConfig{}.GPU), 16},
|
||||
{"cDiarModelConfig.Preset", unsafe.Offsetof(cDiarModelConfig{}.Preset), 24},
|
||||
{"cDiarModelConfig.ChunkFrames", unsafe.Offsetof(cDiarModelConfig{}.ChunkFrames), 32},
|
||||
{"cDiarModelConfig.RightContextFrames", unsafe.Offsetof(cDiarModelConfig{}.RightContextFrames), 36},
|
||||
{"cDiarModelConfig.LeftContextFrames", unsafe.Offsetof(cDiarModelConfig{}.LeftContextFrames), 40},
|
||||
{"cDiarModelConfig.FIFOFrames", unsafe.Offsetof(cDiarModelConfig{}.FIFOFrames), 44},
|
||||
{"cDiarModelConfig.SpkcacheFrames", unsafe.Offsetof(cDiarModelConfig{}.SpkcacheFrames), 48},
|
||||
{"cDiarModelConfig.UpdatePeriodFrames", unsafe.Offsetof(cDiarModelConfig{}.UpdatePeriodFrames), 52},
|
||||
|
||||
{"cDiarSegmentationConfig.Onset", unsafe.Offsetof(cDiarSegmentationConfig{}.Onset), 8},
|
||||
{"cDiarSegmentationConfig.Offset", unsafe.Offsetof(cDiarSegmentationConfig{}.Offset), 12},
|
||||
{"cDiarSegmentationConfig.PadOnsetSec", unsafe.Offsetof(cDiarSegmentationConfig{}.PadOnsetSec), 16},
|
||||
{"cDiarSegmentationConfig.PadOffsetSec", unsafe.Offsetof(cDiarSegmentationConfig{}.PadOffsetSec), 24},
|
||||
{"cDiarSegmentationConfig.MinGapSec", unsafe.Offsetof(cDiarSegmentationConfig{}.MinGapSec), 32},
|
||||
{"cDiarSegmentationConfig.MinDurationSec", unsafe.Offsetof(cDiarSegmentationConfig{}.MinDurationSec), 40},
|
||||
|
||||
{"cDiarSegment.StartTime", unsafe.Offsetof(cDiarSegment{}.StartTime), 0},
|
||||
{"cDiarSegment.EndTime", unsafe.Offsetof(cDiarSegment{}.EndTime), 8},
|
||||
{"cDiarSegment.Speaker", unsafe.Offsetof(cDiarSegment{}.Speaker), 16},
|
||||
|
||||
{"cTTSModelConfig.MagpieModel", unsafe.Offsetof(cTTSModelConfig{}.MagpieModel), 8},
|
||||
{"cTTSModelConfig.CodecModel", unsafe.Offsetof(cTTSModelConfig{}.CodecModel), 16},
|
||||
{"cTTSModelConfig.TokenizerModelDir", unsafe.Offsetof(cTTSModelConfig{}.TokenizerModelDir), 24},
|
||||
{"cTTSModelConfig.TextNormalizerModelDir", unsafe.Offsetof(cTTSModelConfig{}.TextNormalizerModelDir), 32},
|
||||
|
||||
{"cTTSRuntimeConfig.Speaker", unsafe.Offsetof(cTTSRuntimeConfig{}.Speaker), 8},
|
||||
{"cTTSRuntimeConfig.Threads", unsafe.Offsetof(cTTSRuntimeConfig{}.Threads), 12},
|
||||
{"cTTSRuntimeConfig.CodecThreads", unsafe.Offsetof(cTTSRuntimeConfig{}.CodecThreads), 16},
|
||||
{"cTTSRuntimeConfig.Seed", unsafe.Offsetof(cTTSRuntimeConfig{}.Seed), 20},
|
||||
{"cTTSRuntimeConfig.Steps", unsafe.Offsetof(cTTSRuntimeConfig{}.Steps), 24},
|
||||
{"cTTSRuntimeConfig.TopK", unsafe.Offsetof(cTTSRuntimeConfig{}.TopK), 28},
|
||||
{"cTTSRuntimeConfig.ChunkFrames", unsafe.Offsetof(cTTSRuntimeConfig{}.ChunkFrames), 32},
|
||||
{"cTTSRuntimeConfig.CodecQueueDepth", unsafe.Offsetof(cTTSRuntimeConfig{}.CodecQueueDepth), 36},
|
||||
{"cTTSRuntimeConfig.CodecHistoryFrames", unsafe.Offsetof(cTTSRuntimeConfig{}.CodecHistoryFrames), 40},
|
||||
{"cTTSRuntimeConfig.CodecFutureFrames", unsafe.Offsetof(cTTSRuntimeConfig{}.CodecFutureFrames), 44},
|
||||
{"cTTSRuntimeConfig.WindowMs", unsafe.Offsetof(cTTSRuntimeConfig{}.WindowMs), 48},
|
||||
{"cTTSRuntimeConfig.Temperature", unsafe.Offsetof(cTTSRuntimeConfig{}.Temperature), 52},
|
||||
{"cTTSRuntimeConfig.OverrideTemperature", unsafe.Offsetof(cTTSRuntimeConfig{}.OverrideTemperature), 56},
|
||||
{"cTTSRuntimeConfig.CFGScale", unsafe.Offsetof(cTTSRuntimeConfig{}.CFGScale), 60},
|
||||
{"cTTSRuntimeConfig.OverrideCFGScale", unsafe.Offsetof(cTTSRuntimeConfig{}.OverrideCFGScale), 64},
|
||||
{"cTTSRuntimeConfig.UseCFG", unsafe.Offsetof(cTTSRuntimeConfig{}.UseCFG), 65},
|
||||
{"cTTSRuntimeConfig.UseLocalTransformer", unsafe.Offsetof(cTTSRuntimeConfig{}.UseLocalTransformer), 66},
|
||||
{"cTTSRuntimeConfig.UseKVCache", unsafe.Offsetof(cTTSRuntimeConfig{}.UseKVCache), 67},
|
||||
{"cTTSRuntimeConfig.UseStatefulCodec", unsafe.Offsetof(cTTSRuntimeConfig{}.UseStatefulCodec), 68},
|
||||
{"cTTSRuntimeConfig.CodecCPU", unsafe.Offsetof(cTTSRuntimeConfig{}.CodecCPU), 69},
|
||||
{"cTTSRuntimeConfig.FlushPartialChunk", unsafe.Offsetof(cTTSRuntimeConfig{}.FlushPartialChunk), 70},
|
||||
{"cTTSRuntimeConfig.Verbose", unsafe.Offsetof(cTTSRuntimeConfig{}.Verbose), 71},
|
||||
{"cTTSRuntimeConfig.LTBackend", unsafe.Offsetof(cTTSRuntimeConfig{}.LTBackend), 72},
|
||||
{"cTTSRuntimeConfig.SamplingBackend", unsafe.Offsetof(cTTSRuntimeConfig{}.SamplingBackend), 76},
|
||||
{"cTTSRuntimeConfig.UMAMode", unsafe.Offsetof(cTTSRuntimeConfig{}.UMAMode), 80},
|
||||
{"cTTSRuntimeConfig.LongformMode", unsafe.Offsetof(cTTSRuntimeConfig{}.LongformMode), 84},
|
||||
{"cTTSRuntimeConfig.LTFP32", unsafe.Offsetof(cTTSRuntimeConfig{}.LTFP32), 88},
|
||||
|
||||
{"cTTSSynthesizerConfig.Model", unsafe.Offsetof(cTTSSynthesizerConfig{}.Model), 8},
|
||||
{"cTTSSynthesizerConfig.Runtime", unsafe.Offsetof(cTTSSynthesizerConfig{}.Runtime), 16},
|
||||
{"cTTSSynthesizerConfig.DefaultLanguageCode", unsafe.Offsetof(cTTSSynthesizerConfig{}.DefaultLanguageCode), 24},
|
||||
{"cTTSSynthesizerConfig.DefaultVoiceName", unsafe.Offsetof(cTTSSynthesizerConfig{}.DefaultVoiceName), 32},
|
||||
|
||||
{"cTTSSynthesisOptions.RequestID", unsafe.Offsetof(cTTSSynthesisOptions{}.RequestID), 8},
|
||||
{"cTTSSynthesisOptions.LanguageCode", unsafe.Offsetof(cTTSSynthesisOptions{}.LanguageCode), 16},
|
||||
{"cTTSSynthesisOptions.Speaker", unsafe.Offsetof(cTTSSynthesisOptions{}.Speaker), 24},
|
||||
{"cTTSSynthesisOptions.Seed", unsafe.Offsetof(cTTSSynthesisOptions{}.Seed), 28},
|
||||
{"cTTSSynthesisOptions.Steps", unsafe.Offsetof(cTTSSynthesisOptions{}.Steps), 32},
|
||||
{"cTTSSynthesisOptions.TopK", unsafe.Offsetof(cTTSSynthesisOptions{}.TopK), 36},
|
||||
{"cTTSSynthesisOptions.Temperature", unsafe.Offsetof(cTTSSynthesisOptions{}.Temperature), 40},
|
||||
{"cTTSSynthesisOptions.OverrideTemperature", unsafe.Offsetof(cTTSSynthesisOptions{}.OverrideTemperature), 44},
|
||||
{"cTTSSynthesisOptions.CFGScale", unsafe.Offsetof(cTTSSynthesisOptions{}.CFGScale), 48},
|
||||
{"cTTSSynthesisOptions.OverrideCFGScale", unsafe.Offsetof(cTTSSynthesisOptions{}.OverrideCFGScale), 52},
|
||||
{"cTTSSynthesisOptions.VoiceName", unsafe.Offsetof(cTTSSynthesisOptions{}.VoiceName), 56},
|
||||
{"cTTSSynthesisOptions.OutputSampleRate", unsafe.Offsetof(cTTSSynthesisOptions{}.OutputSampleRate), 64},
|
||||
|
||||
{"cNMTBackendConfig.GPU", unsafe.Offsetof(cNMTBackendConfig{}.GPU), 8},
|
||||
{"cNMTModelConfig.Path", unsafe.Offsetof(cNMTModelConfig{}.Path), 8},
|
||||
{"cNMTModelConfig.NCtx", unsafe.Offsetof(cNMTModelConfig{}.NCtx), 16},
|
||||
|
||||
{"cNMTTranslatorConfig.Backend", unsafe.Offsetof(cNMTTranslatorConfig{}.Backend), 8},
|
||||
{"cNMTTranslatorConfig.Model", unsafe.Offsetof(cNMTTranslatorConfig{}.Model), 16},
|
||||
{"cNMTTranslatorConfig.Generation", unsafe.Offsetof(cNMTTranslatorConfig{}.Generation), 24},
|
||||
{"cNMTTranslatorConfig.Pool", unsafe.Offsetof(cNMTTranslatorConfig{}.Pool), 32},
|
||||
}
|
||||
}
|
||||
|
||||
var _ = Describe("C struct mirrors", func() {
|
||||
// These need no shared object, so they run on any checkout and catch a
|
||||
// transcription slip the moment it is introduced.
|
||||
It("matches the C sizeof of every mirrored struct", func() {
|
||||
for _, l := range structSizes() {
|
||||
Expect(l.got).To(Equal(l.want), "%s: Go mirror is %d bytes, C says %d", l.what, l.got, l.want)
|
||||
}
|
||||
})
|
||||
|
||||
It("matches the C offset of every mirrored field", func() {
|
||||
for _, l := range structOffsets() {
|
||||
Expect(l.got).To(Equal(l.want), "%s: Go offset %d, C offset %d", l.what, l.got, l.want)
|
||||
}
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("C ABI binding", func() {
|
||||
BeforeEach(func() {
|
||||
if !librariesPresent() {
|
||||
if requireLibs() {
|
||||
cwd, _ := os.Getwd()
|
||||
Fail("NEMO_SPEECH_REQUIRE_LIBS=1 but the shared libraries are not in " + cwd +
|
||||
": these specs are the ABI defence and must not be skipped." +
|
||||
" Run make -C backend/go/nemo-speech-cpp stage-libs")
|
||||
}
|
||||
Skip("shared libraries not built, run make in backend/go/nemo-speech-cpp")
|
||||
}
|
||||
Expect(openLibraries()).To(Succeed())
|
||||
})
|
||||
|
||||
It("resolves every bound symbol", func() {
|
||||
Expect(symbols()).ToNot(BeEmpty())
|
||||
for _, s := range symbols() {
|
||||
Expect(registerOne(s)).To(Succeed())
|
||||
}
|
||||
})
|
||||
|
||||
// The library reports its own sizeof through the size field of each
|
||||
// defaults struct. A Go mirror that disagrees means every field after the
|
||||
// first divergence is read from the wrong offset, which no compiler or
|
||||
// linker check would catch. The three structs below are the only ones with
|
||||
// a defaults entry point, so they are the only ones the runtime can be
|
||||
// asked about directly.
|
||||
It("mirrors the C recognition-options struct layout", func() {
|
||||
def := ASRRecognitionOptionsDef()
|
||||
Expect(def.Size).To(Equal(unsafe.Sizeof(cASRRecognitionOptions{})),
|
||||
"cASRRecognitionOptions does not match the C layout")
|
||||
})
|
||||
|
||||
It("mirrors the C TTS runtime-config struct layout", func() {
|
||||
def := TTSRuntimeConfigDefault()
|
||||
Expect(def.Size).To(Equal(unsafe.Sizeof(cTTSRuntimeConfig{})),
|
||||
"cTTSRuntimeConfig does not match the C layout")
|
||||
})
|
||||
|
||||
It("mirrors the C TTS synthesis-options struct layout", func() {
|
||||
def := TTSSynthesisOptionsDefault()
|
||||
Expect(def.Size).To(Equal(unsafe.Sizeof(cTTSSynthesisOptions{})),
|
||||
"cTTSSynthesisOptions does not match the C layout")
|
||||
})
|
||||
|
||||
// A size match alone cannot see a field read from the wrong offset when two
|
||||
// padding mistakes cancel out, and structOffsets checks the mirrors against
|
||||
// numbers transcribed by the same hand that wrote them. This spec is the
|
||||
// only layer independent of that transcription: it reads values back out of
|
||||
// the running library, so a systematically wrong table cannot hide here.
|
||||
//
|
||||
// Deliberately narrow. An earlier version pinned roughly forty default
|
||||
// values, which would make a legitimate pin bump (threads 4 to 8, or a
|
||||
// flipped flush_partial_chunk) fail with a message that reads like a layout
|
||||
// error. What survives is only the values that are contract, not tuning:
|
||||
//
|
||||
// - max_alternatives is the single non-zero in an otherwise memset-zero
|
||||
// struct, and asr.h documents "<= 1 = 1-best (default)". It pins offset
|
||||
// 56, deep in the tail past the bool run.
|
||||
// - The synthesis-options run of four -1 sentinels, each documented in
|
||||
// tts.h as "< 0 = synthesizer default", pins offsets 24 through 36, and
|
||||
// temperature witnesses that the run stops exactly at offset 40. A
|
||||
// mirror whose tail is shifted by one field spills a -1 into that zero.
|
||||
// - Two -1 sentinels at the ends of the runtime config's long int32 run
|
||||
// pin offset 20 and offset 40 without depending on any tunable.
|
||||
//
|
||||
// Sources: src/asr/c_api.cpp nemo_speech_asr_recognition_options_default,
|
||||
// src/tts/magpietts/runtime.h MagpieRuntimeConfig, src/tts/c_api.cpp
|
||||
// nemo_speech_tts_synthesis_options_default.
|
||||
It("reads the documented default values back through the mirrors", func() {
|
||||
asr := ASRRecognitionOptionsDef()
|
||||
Expect(asr.MaxAlternatives).To(Equal(int32(1)))
|
||||
|
||||
rt := TTSRuntimeConfigDefault()
|
||||
Expect(rt.Seed).To(Equal(int32(-1)))
|
||||
Expect(rt.CodecHistoryFrames).To(Equal(int32(-1)))
|
||||
|
||||
opt := TTSSynthesisOptionsDefault()
|
||||
Expect(opt.Speaker).To(Equal(int32(-1)))
|
||||
Expect(opt.Seed).To(Equal(int32(-1)))
|
||||
Expect(opt.Steps).To(Equal(int32(-1)))
|
||||
Expect(opt.TopK).To(Equal(int32(-1)))
|
||||
Expect(opt.Temperature).To(Equal(float32(0)))
|
||||
})
|
||||
|
||||
It("reports a non-empty version from each library", func() {
|
||||
Expect(ASRVersion()).ToNot(BeEmpty())
|
||||
Expect(TTSVersion()).ToNot(BeEmpty())
|
||||
Expect(NMTVersion()).ToNot(BeEmpty())
|
||||
})
|
||||
})
|
||||
@@ -1,374 +0,0 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"runtime"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
"unsafe"
|
||||
|
||||
pb "github.com/mudler/LocalAI/pkg/grpc/proto"
|
||||
"github.com/mudler/xlog"
|
||||
"google.golang.org/grpc/codes"
|
||||
"google.golang.org/grpc/status"
|
||||
)
|
||||
|
||||
// asrWord is one decoded word with its millisecond offsets and 1-based speaker
|
||||
// tag (0 when diarization was not requested).
|
||||
type asrWord struct {
|
||||
Text string
|
||||
Start int32
|
||||
End int32
|
||||
Speaker int32
|
||||
}
|
||||
|
||||
// pinPtr pins v for the lifetime of p and returns its address in the uintptr
|
||||
// form the config structs carry.
|
||||
//
|
||||
// The config structs mirror C, so their pointer members are uintptr, which the
|
||||
// collector does not trace. Everything reachable only through one of them is
|
||||
// therefore invisible to the GC while C is reading it, exactly as described on
|
||||
// cstr, and needs the same pin. runtime.KeepAlive would cover collection but
|
||||
// says nothing about relocation, and the guarantee wanted here is that the
|
||||
// address C holds stays the address of the object.
|
||||
func pinPtr[T any](p *runtime.Pinner, v *T) uintptr {
|
||||
p.Pin(v)
|
||||
// #nosec G103 -- v is pinned into p on the previous line, so its address is
|
||||
// stable and traced for as long as p lives; every caller defers p.Unpin only
|
||||
// after the create call that reads it. One-way, like cstr: nothing converts
|
||||
// this uintptr back to a pointer.
|
||||
return uintptr(unsafe.Pointer(v))
|
||||
}
|
||||
|
||||
// asrDiarConfig builds the config for the diarizer attached to a recognizer.
|
||||
//
|
||||
// Extracted from loadASR for the same reason diarModelConfig was extracted from
|
||||
// loadDiarizer: the six frame counts are sentinel-sensitive and invisible to
|
||||
// every other check in the tree. src/asr/c_api.cpp:151-165 applies five of them
|
||||
// when they are > 0 but applies left_context_frames when it is >= 0, so a
|
||||
// dropped -1 does not fall back to the model's own streaming geometry, it pins
|
||||
// the left context to zero. The struct is the right shape either way, so the
|
||||
// layout assertions in abi_test.go cannot see it and only a spec on this builder
|
||||
// can.
|
||||
//
|
||||
// diarGeometryDefault is shared with the standalone diarizer rather than
|
||||
// restated: it is the same sentinel, from the same rule, in the same runtime.
|
||||
//
|
||||
// modelPath is a C pointer from cstr, not a Go string, and the caller owns its
|
||||
// release.
|
||||
func asrDiarConfig(modelPath uintptr) cASRDiarConfig {
|
||||
return cASRDiarConfig{
|
||||
Size: unsafe.Sizeof(cASRDiarConfig{}),
|
||||
ModelPath: modelPath,
|
||||
ChunkFrames: diarGeometryDefault,
|
||||
RightContextFrames: diarGeometryDefault,
|
||||
LeftContextFrames: diarGeometryDefault,
|
||||
FIFOFrames: diarGeometryDefault,
|
||||
SpkcacheFrames: diarGeometryDefault,
|
||||
UpdatePeriodFrames: diarGeometryDefault,
|
||||
}
|
||||
}
|
||||
|
||||
// loadASR creates the recognizer, attaching VAD, PnC, ITN and diarization when
|
||||
// the corresponding options were set.
|
||||
//
|
||||
// Every field below is assigned by name against include/nemo_speech/asr.h. The
|
||||
// sub-configs are optional pointers: a nil one means "library defaults", which
|
||||
// is why each is populated only when its option was given rather than always
|
||||
// being attached with empty strings.
|
||||
//
|
||||
// Each struct's Size is load-bearing, not decoration. The runtime decides a
|
||||
// field is present with HAS_FIELD (src/asr/c_api.cpp), which tests the caller's
|
||||
// size against offsetof(field) + sizeof(field), so a config sent with Size 0
|
||||
// has every field ignored and the model silently loads with defaults.
|
||||
//
|
||||
// This must not take engineMu: Load is its only caller and already holds it.
|
||||
func (n *NemoSpeech) loadASR(modelFile string) error {
|
||||
// nemo_speech_asr_create deep-copies every const char* into a std::string
|
||||
// (src/asr/c_api.cpp to_config, via str_or_empty) and retains no pointer
|
||||
// afterwards, so pinning for the duration of the create call is both
|
||||
// necessary and sufficient.
|
||||
var pinner runtime.Pinner
|
||||
defer pinner.Unpin()
|
||||
|
||||
pathP, freePath := cstr(modelFile)
|
||||
defer freePath()
|
||||
|
||||
model := cASRModelConfig{Size: unsafe.Sizeof(cASRModelConfig{}), Path: pathP}
|
||||
backend := cASRBackendConfig{Size: unsafe.Sizeof(cASRBackendConfig{}), GPU: n.opts.gpu}
|
||||
|
||||
cfg := cASRRecognizerConfig{
|
||||
Size: unsafe.Sizeof(cASRRecognizerConfig{}),
|
||||
Backend: pinPtr(&pinner, &backend),
|
||||
Model: pinPtr(&pinner, &model),
|
||||
}
|
||||
|
||||
var vad cASRVADConfig
|
||||
if n.opts.vadModel != "" {
|
||||
p, free := cstr(n.opts.vadModel)
|
||||
defer free()
|
||||
vad = cASRVADConfig{Size: unsafe.Sizeof(cASRVADConfig{}), ModelPath: p}
|
||||
cfg.VAD = pinPtr(&pinner, &vad)
|
||||
}
|
||||
|
||||
var postproc cASRPostprocConfig
|
||||
if n.opts.itnDir != "" || n.opts.pncModel != "" {
|
||||
itnP, freeITN := cstr(n.opts.itnDir)
|
||||
defer freeITN()
|
||||
pncP, freePNC := cstr(n.opts.pncModel)
|
||||
defer freePNC()
|
||||
postproc = cASRPostprocConfig{
|
||||
Size: unsafe.Sizeof(cASRPostprocConfig{}),
|
||||
ITNModelDir: itnP,
|
||||
PNCModelPath: pncP,
|
||||
}
|
||||
cfg.Postproc = pinPtr(&pinner, &postproc)
|
||||
}
|
||||
|
||||
var diar cASRDiarConfig
|
||||
if n.opts.diarModel != "" {
|
||||
p, free := cstr(n.opts.diarModel)
|
||||
defer free()
|
||||
diar = asrDiarConfig(p)
|
||||
cfg.Diar = pinPtr(&pinner, &diar)
|
||||
}
|
||||
|
||||
xlog.Info("nemo-speech-cpp: creating recognizer",
|
||||
"gpu", n.opts.gpu,
|
||||
"vad", n.opts.vadModel != "",
|
||||
"pnc", n.opts.pncModel != "",
|
||||
"itn", n.opts.itnDir != "",
|
||||
"diarization", n.opts.diarModel != "")
|
||||
|
||||
// #nosec G103 -- cfg is a local POD struct passed as a pointer for the
|
||||
// duration of this call only; every uintptr member it carries is either a
|
||||
// cstr allocation or a pinPtr address, all pinned above and released by the
|
||||
// defers, and nemo_speech_asr_create deep-copies and retains nothing.
|
||||
if st := ASRCreate(unsafe.Pointer(&cfg), &n.recognizer); st != 0 {
|
||||
return statusErrorf(st, "nemo-speech-cpp: asr create: %s", ASRLastError())
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// recognizeF32 runs one offline decode and returns the result handle, which the
|
||||
// caller must destroy.
|
||||
//
|
||||
// The empty-input guard is here rather than at the call site because &pcm[0]
|
||||
// panics on a zero-length slice: Go never reaches the C side's own "empty
|
||||
// audio" rejection. A silent clip or a truncated upload decodes to zero
|
||||
// samples, which is ordinary input, not an exotic one.
|
||||
//
|
||||
// The caller must hold engineMu.
|
||||
func recognizeF32(recognizer uintptr, opts *cASRRecognitionOptions, pcm []float32, sampleRate int32) (uintptr, error) {
|
||||
if len(pcm) == 0 {
|
||||
return 0, status.Error(codes.InvalidArgument, "nemo-speech-cpp: empty audio")
|
||||
}
|
||||
|
||||
var result uintptr
|
||||
// #nosec G103 -- opts is the caller's live struct, borrowed for this call
|
||||
// only; its LanguageCode is a cstr allocation the caller keeps pinned across
|
||||
// it. &pcm[0] is guarded by the empty check above and the length handed over
|
||||
// is exactly len(pcm), so the runtime cannot read past the slice.
|
||||
if st := ASRRecognizeF32(recognizer, unsafe.Pointer(opts),
|
||||
&pcm[0], uint64(len(pcm)), sampleRate, &result); st != 0 {
|
||||
return 0, statusErrorf(st, "nemo-speech-cpp: recognize: %s", ASRLastError())
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
// msToNanos converts a runtime word offset to the wire unit. The runtime
|
||||
// reports milliseconds (src/asr/types.h); TranscriptSegment.start/end and
|
||||
// TranscriptWord.start/end are int64 nanoseconds, which core/backend reads
|
||||
// straight into a time.Duration.
|
||||
func msToNanos(ms int32) int64 {
|
||||
return int64(ms) * int64(time.Millisecond)
|
||||
}
|
||||
|
||||
// extractWords pulls the top alternative's words out of a result handle.
|
||||
func extractWords(result uintptr) []asrWord {
|
||||
if ASRResultAlternativeCount(result) == 0 {
|
||||
return nil
|
||||
}
|
||||
count := ASRResultWordCount(result, 0)
|
||||
words := make([]asrWord, 0, count)
|
||||
for i := uint64(0); i < count; i++ {
|
||||
words = append(words, asrWord{
|
||||
Text: ASRResultWordText(result, 0, i),
|
||||
Start: ASRResultWordStartTime(result, 0, i),
|
||||
End: ASRResultWordEndTime(result, 0, i),
|
||||
Speaker: ASRResultWordSpeakerTag(result, 0, i),
|
||||
})
|
||||
}
|
||||
return words
|
||||
}
|
||||
|
||||
// wordsRequested reports whether the caller asked for word-level timestamps.
|
||||
// The OpenAI transcription API gates word timings behind
|
||||
// timestamp_granularities[] containing "word" and defaults to segment level
|
||||
// otherwise; every backend here follows that contract (see
|
||||
// backend/go/parakeet-cpp).
|
||||
func wordsRequested(granularities []string) bool {
|
||||
for _, g := range granularities {
|
||||
if strings.EqualFold(strings.TrimSpace(g), "word") {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// wordsToSegments groups words into one segment per consecutive speaker run.
|
||||
// Without diarization every word carries speaker 0, so this collapses to a
|
||||
// single segment.
|
||||
//
|
||||
// The boundary is a CHANGE of speaker, not the first appearance of one: a
|
||||
// conversation that returns to an earlier speaker has to start a new turn
|
||||
// rather than reopen the old one.
|
||||
//
|
||||
// withWords additionally attaches the per-word timings that
|
||||
// core/backend/transcript.go turns into the response's word list. It is off by
|
||||
// default because the OpenAI contract asks for word timestamps explicitly, and
|
||||
// a long transcript pays for every word twice otherwise.
|
||||
func wordsToSegments(words []asrWord, withWords bool) []*pb.TranscriptSegment {
|
||||
if len(words) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
var segs []*pb.TranscriptSegment
|
||||
start := 0
|
||||
flush := func(end int) {
|
||||
run := words[start:end]
|
||||
texts := make([]string, 0, len(run))
|
||||
for _, w := range run {
|
||||
texts = append(texts, w.Text)
|
||||
}
|
||||
seg := &pb.TranscriptSegment{
|
||||
// #nosec G115 -- TranscriptSegment.Id is int32 on the wire, and segs
|
||||
// holds one entry per speaker run over the words of a single decode
|
||||
// result, which exhausts memory long before it reaches 2^31.
|
||||
Id: int32(len(segs)),
|
||||
Text: strings.Join(texts, " "),
|
||||
Start: msToNanos(run[0].Start),
|
||||
End: msToNanos(run[len(run)-1].End),
|
||||
}
|
||||
// The speaker tag is 1-based with 0 meaning untagged, so an undiarized
|
||||
// run must stay unlabelled rather than be attributed to a speaker "0".
|
||||
if run[0].Speaker > 0 {
|
||||
seg.Speaker = strconv.Itoa(int(run[0].Speaker))
|
||||
}
|
||||
if withWords {
|
||||
seg.Words = wordsToProto(run)
|
||||
}
|
||||
segs = append(segs, seg)
|
||||
}
|
||||
|
||||
for i := 1; i < len(words); i++ {
|
||||
if words[i].Speaker != words[start].Speaker {
|
||||
flush(i)
|
||||
start = i
|
||||
}
|
||||
}
|
||||
flush(len(words))
|
||||
return segs
|
||||
}
|
||||
|
||||
// AudioTranscription decodes the audio at req.Dst and returns one offline
|
||||
// transcription.
|
||||
//
|
||||
// The whole body runs inside withEngine, so the family check and the C calls
|
||||
// that trust the handle happen under a single acquisition of engineMu. Decoding
|
||||
// the audio is in there too: pkg/grpc/server.go already serialises RPCs on this
|
||||
// backend through base.SingleThread, so the lock costs no concurrency, and the
|
||||
// alternative (check, unlock, decode, relock) is the exact gap Free can land in.
|
||||
func (n *NemoSpeech) AudioTranscription(ctx context.Context, req *pb.TranscriptRequest) (pb.TranscriptResult, error) {
|
||||
var out *pb.TranscriptResult
|
||||
if err := n.withEngine(familyASR, func() error {
|
||||
r, err := n.transcribe(req)
|
||||
out = r
|
||||
return err
|
||||
}); err != nil {
|
||||
return pb.TranscriptResult{}, err
|
||||
}
|
||||
// transcribe returns a non-nil result whenever it returns a nil error, so
|
||||
// this cannot fire today. It is a guard rather than a comment because the
|
||||
// alternative to stating the invariant is a nil dereference in an RPC
|
||||
// handler if a later edit ever adds a success path that forgets to set it.
|
||||
if out == nil {
|
||||
return pb.TranscriptResult{}, status.Error(codes.Internal,
|
||||
"nemo-speech-cpp: transcription produced no result")
|
||||
}
|
||||
|
||||
// Assembled field by field rather than dereferenced: the RPC signature
|
||||
// returns the proto message by value, but the message embeds a mutex, so
|
||||
// copying the struct is a copylocks violation. Every backend in this tree
|
||||
// gets around it the same way, by only ever returning a composite literal.
|
||||
return pb.TranscriptResult{
|
||||
Text: out.Text,
|
||||
Segments: out.Segments,
|
||||
Language: out.Language,
|
||||
Duration: out.Duration,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// transcribe is AudioTranscription's body. The caller must hold engineMu.
|
||||
func (n *NemoSpeech) transcribe(req *pb.TranscriptRequest) (*pb.TranscriptResult, error) {
|
||||
if req.GetDst() == "" {
|
||||
return nil, status.Error(codes.InvalidArgument,
|
||||
"nemo-speech-cpp: TranscriptRequest.dst (audio path) is required")
|
||||
}
|
||||
|
||||
pcm, sampleRate, err := decodeAudioMono16k(req.GetDst())
|
||||
if err != nil {
|
||||
return nil, status.Errorf(codes.InvalidArgument,
|
||||
"nemo-speech-cpp: read audio: %v", err)
|
||||
}
|
||||
// Rejected here, before anything crosses the ABI, and not only inside
|
||||
// recognizeF32: a silent or truncated upload decodes to zero samples, and
|
||||
// there is no point building options and pinning strings for a request
|
||||
// that cannot produce a transcript. recognizeF32 keeps its own guard as a
|
||||
// precondition on the function.
|
||||
if len(pcm) == 0 {
|
||||
return nil, status.Error(codes.InvalidArgument, "nemo-speech-cpp: empty audio")
|
||||
}
|
||||
|
||||
// A per-request language wins over the model-level default; both may be
|
||||
// empty, which the runtime reads as auto/model default.
|
||||
language := req.GetLanguage()
|
||||
if language == "" {
|
||||
language = n.opts.languageCode
|
||||
}
|
||||
langP, freeLang := cstr(language)
|
||||
defer freeLang()
|
||||
|
||||
opts := ASRRecognitionOptionsDef()
|
||||
opts.LanguageCode = langP
|
||||
// Segments are built out of word offsets, so they are always asked for.
|
||||
opts.EnableWordTimeOffsets = true
|
||||
// Keyed on the recognizer owning a diar model, not on req.Diarize: asr.h
|
||||
// documents that a request asking for diarization from a recognizer created
|
||||
// without one fails with INVALID_ARGUMENT, and setting diar_model is already
|
||||
// the operator's opt-in.
|
||||
opts.EnableSpeakerDiarization = n.opts.diarModel != ""
|
||||
|
||||
result, err := recognizeF32(n.recognizer, &opts, pcm, sampleRate)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer ASRResultDestroy(result)
|
||||
|
||||
out := &pb.TranscriptResult{
|
||||
Text: ASRResultTranscript(result, 0),
|
||||
Segments: wordsToSegments(extractWords(result),
|
||||
wordsRequested(req.GetTimestampGranularities())),
|
||||
}
|
||||
// Multilingual models report what they decided the audio was; monolingual
|
||||
// ones report nothing, and an empty language is better than echoing back
|
||||
// whatever the caller guessed.
|
||||
if ASRResultLanguageCount(result, 0) > 0 {
|
||||
out.Language = ASRResultLanguageCode(result, 0, 0)
|
||||
}
|
||||
if sampleRate > 0 {
|
||||
out.Duration = float32(len(pcm)) / float32(sampleRate)
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
@@ -1,526 +0,0 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"unsafe"
|
||||
|
||||
pb "github.com/mudler/LocalAI/pkg/grpc/proto"
|
||||
"github.com/mudler/xlog"
|
||||
"google.golang.org/grpc/codes"
|
||||
"google.golang.org/grpc/status"
|
||||
)
|
||||
|
||||
// streamChunkSamples is one push into a streaming session. At 16 kHz mono 1600
|
||||
// samples is 100 ms, short enough that the decoder is polled often enough to
|
||||
// see an endpoint promptly and short enough that a cancelled request stops
|
||||
// within one push.
|
||||
const streamChunkSamples = 1600
|
||||
|
||||
// The rates nemo_speech_asr_stream_push_f32 will resample from (asr.h). Outside
|
||||
// this range the runtime has nothing to do with the audio, and 0 is NOT
|
||||
// "unknown": it means "these samples are already at the model rate".
|
||||
const (
|
||||
minStreamSampleRate = 8000
|
||||
maxStreamSampleRate = 96000
|
||||
// TranscriptLiveConfig.sample_rate documents 0 as 16 kHz, which is a
|
||||
// different meaning from the C API's 0, so it is resolved before the push.
|
||||
defaultLiveSampleRate = 16000
|
||||
)
|
||||
|
||||
// streamResult is one result lifted out of C memory. Everything is copied
|
||||
// before nemo_speech_asr_result_destroy runs, so a streamResult outlives the
|
||||
// handle it came from.
|
||||
type streamResult struct {
|
||||
Text string
|
||||
Final bool
|
||||
Words []asrWord
|
||||
}
|
||||
|
||||
// asrSession is the streaming half of the ASR C API, narrowed to the four
|
||||
// entry points the two streaming RPCs use.
|
||||
//
|
||||
// It is an interface because there is no NeMo GGUF small enough to keep in the
|
||||
// tree, so the loops on top of it (chunking, the need-more-audio drain, the
|
||||
// live config/reset protocol) would otherwise have no test at all. The seam is
|
||||
// at the ABI, not at the model: a fake session scripts what the C API returns,
|
||||
// it does not pretend to transcribe anything.
|
||||
type asrSession interface {
|
||||
// push buffers audio. It does not decode; next drives that.
|
||||
push(pcm []float32, sampleRate int32) error
|
||||
// finish flushes the decoder tail. The end-of-stream final then comes back
|
||||
// from next.
|
||||
finish() error
|
||||
// next pulls one result. ok=false means the decoder needs more audio,
|
||||
// which is a pause in the stream and not an error or an end.
|
||||
next() (result streamResult, ok bool, err error)
|
||||
close()
|
||||
}
|
||||
|
||||
// sessionOpener creates a session for one language. n.openSession is the
|
||||
// C-backed implementation.
|
||||
type sessionOpener func(language string) (asrSession, error)
|
||||
|
||||
// cSession is the real asrSession, over one nemo_speech_asr_stream.
|
||||
type cSession struct {
|
||||
handle uintptr
|
||||
}
|
||||
|
||||
func (s *cSession) push(pcm []float32, sampleRate int32) error {
|
||||
// &pcm[0] panics on an empty slice, and an empty frame is ordinary input
|
||||
// from a live caller: it is a keepalive, not audio.
|
||||
if len(pcm) == 0 {
|
||||
return nil
|
||||
}
|
||||
if st := ASRStreamPushF32(s.handle, &pcm[0], uint64(len(pcm)), sampleRate); st != 0 {
|
||||
return statusErrorf(st, "nemo-speech-cpp: stream push: %s", ASRLastError())
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *cSession) finish() error {
|
||||
if st := ASRStreamFinish(s.handle); st != 0 {
|
||||
return statusErrorf(st, "nemo-speech-cpp: stream finish: %s", ASRLastError())
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *cSession) next() (streamResult, bool, error) {
|
||||
var handle uintptr
|
||||
if st := ASRStreamNext(s.handle, &handle); st != 0 {
|
||||
return streamResult{}, false, statusErrorf(st, "nemo-speech-cpp: stream next: %s", ASRLastError())
|
||||
}
|
||||
// OK with a NULL handle is the documented "need more audio". Reading it as
|
||||
// an error aborts every stream at the first gap; reading it as "keep
|
||||
// pulling" spins forever.
|
||||
if handle == 0 {
|
||||
return streamResult{}, false, nil
|
||||
}
|
||||
// Destroyed here rather than by the caller: everything below is copied out
|
||||
// of C memory into Go values, so nothing survives that would need it, and
|
||||
// a caller that returned early would otherwise leak the result.
|
||||
defer ASRResultDestroy(handle)
|
||||
|
||||
return streamResult{
|
||||
Text: ASRResultTranscript(handle, 0),
|
||||
Final: ASRResultIsFinal(handle),
|
||||
Words: extractWords(handle),
|
||||
}, true, nil
|
||||
}
|
||||
|
||||
func (s *cSession) close() { ASRStreamClose(s.handle) }
|
||||
|
||||
// openSession starts a streaming recognition on the loaded recognizer.
|
||||
//
|
||||
// The caller must hold engineMu.
|
||||
//
|
||||
// nemo_speech_asr_streaming_recognize copies the options (src/asr/c_api.cpp
|
||||
// to_options) and keeps no pointer into them, so the language buffer only has
|
||||
// to stay pinned across this call, exactly as in loadASR.
|
||||
func (n *NemoSpeech) openSession(language string) (asrSession, error) {
|
||||
// A per-request language wins over the model-level default; both may be
|
||||
// empty, which the runtime reads as auto/model default.
|
||||
if language == "" {
|
||||
language = n.opts.languageCode
|
||||
}
|
||||
langP, freeLang := cstr(language)
|
||||
defer freeLang()
|
||||
|
||||
opts := ASRRecognitionOptionsDef()
|
||||
opts.LanguageCode = langP
|
||||
// Segments and the live word list are built out of word offsets, so they
|
||||
// are always asked for.
|
||||
opts.EnableWordTimeOffsets = true
|
||||
// Keyed on the recognizer owning a diar model rather than on the request:
|
||||
// asr.h documents that asking a recognizer created without one for
|
||||
// diarization fails with INVALID_ARGUMENT.
|
||||
opts.EnableSpeakerDiarization = n.opts.diarModel != ""
|
||||
// interim_results is left off deliberately. The runtime emits interims from
|
||||
// next() regardless of it, and they are filtered here rather than
|
||||
// forwarded: see streamPCM's emit for why the wire contract cannot carry
|
||||
// them.
|
||||
|
||||
var handle uintptr
|
||||
// #nosec G103 -- opts is a local POD struct borrowed for this call only, and
|
||||
// its one uintptr member (LanguageCode) is the cstr allocation pinned by the
|
||||
// deferred freeLang above. to_options copies the struct, so nothing here
|
||||
// outlives the call.
|
||||
if st := ASRStreamingRecognize(n.recognizer, unsafe.Pointer(&opts), &handle); st != 0 {
|
||||
return nil, statusErrorf(st, "nemo-speech-cpp: streaming recognize: %s", ASRLastError())
|
||||
}
|
||||
xlog.Debug("nemo-speech-cpp: streaming session open", "language", language)
|
||||
return &cSession{handle: handle}, nil
|
||||
}
|
||||
|
||||
// chunkPCM slices pcm into fixed-size chunks, leaving the final chunk short
|
||||
// rather than padding it: silence padding would push audio the caller never
|
||||
// sent through the encoder and shift the tail word timings.
|
||||
func chunkPCM(pcm []float32, size int) [][]float32 {
|
||||
if len(pcm) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make([][]float32, 0, (len(pcm)+size-1)/size)
|
||||
for off := 0; off < len(pcm); off += size {
|
||||
out = append(out, pcm[off:min(off+size, len(pcm))])
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// drain pulls every result the session currently has, handing each to emit.
|
||||
// It returns when the session reports it needs more audio, which is the loop's
|
||||
// only terminating condition.
|
||||
func drain(sess asrSession, emit func(streamResult) error) error {
|
||||
for {
|
||||
r, ok, err := sess.next()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
if err := emit(r); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// streamPCM drives one whole clip through an open session, emitting each
|
||||
// finalized utterance as a delta and closing with the assembled result.
|
||||
//
|
||||
// Only finals become deltas, and the reason is the wire contract:
|
||||
// TranscriptStreamResponse.delta is newly-FINALIZED text that consumers
|
||||
// CONCATENATE (core/http/endpoints/openai/transcription.go, and the realtime
|
||||
// semantic-VAD path). An interim is the decoder's running hypothesis for the
|
||||
// utterance in flight, so forwarding "he", "hell", "hello", "Hello." would
|
||||
// assemble to "hehellhelloHello." rather than to the transcript. That the
|
||||
// runtime also postprocesses finals only (build_result_ in
|
||||
// src/asr/recognizer.cpp runs ITN and strip_formatting on the final, so it
|
||||
// rewrites rather than extends the interim) means there is no diffing trick
|
||||
// that would rescue them either.
|
||||
//
|
||||
// The cost is that the first delta of an utterance arrives at its endpoint
|
||||
// rather than mid-word.
|
||||
func streamPCM(ctx context.Context, sess asrSession, pcm []float32, sampleRate int32, wantWords bool, results chan<- *pb.TranscriptStreamResponse) error {
|
||||
if len(pcm) == 0 {
|
||||
return status.Error(codes.InvalidArgument, "nemo-speech-cpp: empty audio")
|
||||
}
|
||||
|
||||
var (
|
||||
full strings.Builder
|
||||
segments []*pb.TranscriptSegment
|
||||
// sawEndpoint records a final that arrived before the tail flush, i.e.
|
||||
// a real endpoint rather than the end of the file.
|
||||
sawEndpoint bool
|
||||
flushing bool
|
||||
tailText string
|
||||
)
|
||||
|
||||
emit := func(r streamResult) error {
|
||||
if !r.Final {
|
||||
return nil
|
||||
}
|
||||
if flushing {
|
||||
tailText += r.Text
|
||||
} else {
|
||||
sawEndpoint = true
|
||||
}
|
||||
if r.Text == "" {
|
||||
return nil
|
||||
}
|
||||
|
||||
// The separator is part of the delta, not added when assembling the
|
||||
// final text, so concatenating the deltas reproduces FinalResult.Text
|
||||
// exactly. Utterance transcripts carry no leading or trailing space of
|
||||
// their own (the runner clears its buffer at each endpoint).
|
||||
delta := r.Text
|
||||
if full.Len() > 0 {
|
||||
delta = " " + delta
|
||||
}
|
||||
full.WriteString(delta)
|
||||
|
||||
// One segment run per utterance, renumbered into the running sequence.
|
||||
// wordsToSegments splits a run further on a speaker change, so a
|
||||
// diarized utterance contributes one segment per turn.
|
||||
segs := wordsToSegments(r.Words, wantWords)
|
||||
if len(segs) == 0 {
|
||||
// Word offsets were requested but a decoder head may still return
|
||||
// none; a segment carrying just the text beats dropping it.
|
||||
segs = []*pb.TranscriptSegment{{Text: r.Text}}
|
||||
}
|
||||
for _, s := range segs {
|
||||
// #nosec G115 -- TranscriptSegment.Id is int32 on the wire, and
|
||||
// segments holds one entry per speaker run per finalized utterance of
|
||||
// a single request, which exhausts memory long before it reaches 2^31.
|
||||
s.Id = int32(len(segments))
|
||||
segments = append(segments, s)
|
||||
}
|
||||
|
||||
results <- &pb.TranscriptStreamResponse{Delta: delta}
|
||||
return nil
|
||||
}
|
||||
|
||||
for _, chunk := range chunkPCM(pcm, streamChunkSamples) {
|
||||
// The RPC body holds engineMu for the whole stream, so Free waits on
|
||||
// it. Without this check a client that disconnected mid-file would pin
|
||||
// the model against unload until the whole clip had been pushed.
|
||||
if err := ctx.Err(); err != nil {
|
||||
return status.Error(codes.Canceled, "nemo-speech-cpp: transcription cancelled")
|
||||
}
|
||||
if err := sess.push(chunk, sampleRate); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := drain(sess, emit); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
flushing = true
|
||||
if err := sess.finish(); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := drain(sess, emit); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
final := &pb.TranscriptResult{
|
||||
Text: full.String(),
|
||||
Segments: segments,
|
||||
// The tail flush returns whatever the decoder was still holding.
|
||||
// Nothing held back after at least one endpoint means the last
|
||||
// endpoint consumed the audio, which is what "the clip ended on an
|
||||
// utterance boundary" means here. Text coming back means it ended
|
||||
// mid-utterance.
|
||||
Eou: sawEndpoint && tailText == "",
|
||||
}
|
||||
if sampleRate > 0 {
|
||||
final.Duration = float32(len(pcm)) / float32(sampleRate)
|
||||
}
|
||||
results <- &pb.TranscriptStreamResponse{FinalResult: final}
|
||||
return nil
|
||||
}
|
||||
|
||||
// runLive drives one bidirectional live session. The protocol is the one
|
||||
// documented on the RPC in backend.proto: a Config first, a ready ack once the
|
||||
// session is open, deltas as utterances finalize, and a terminal result when
|
||||
// the caller closes its send side.
|
||||
//
|
||||
// There is no context here on purpose. The gRPC host closes `in` when the
|
||||
// stream context is cancelled (pkg/grpc/server.go's recv pump), so ranging
|
||||
// over it is what stops this loop, and that is also what releases engineMu for
|
||||
// a waiting Free.
|
||||
func runLive(open sessionOpener, in <-chan *pb.TranscriptLiveRequest, out chan<- *pb.TranscriptLiveResponse) error {
|
||||
first, ok := <-in
|
||||
if !ok {
|
||||
// The caller closed without sending anything. Nothing was opened, so
|
||||
// there is nothing to report.
|
||||
return nil
|
||||
}
|
||||
cfg := first.GetConfig()
|
||||
if cfg == nil {
|
||||
return status.Error(codes.InvalidArgument,
|
||||
"nemo-speech-cpp: the first live message must carry a config")
|
||||
}
|
||||
rate, err := liveSampleRate(cfg)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
sess, err := open(cfg.GetLanguage())
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
// A mid-stream config replaces sess, so this closes whichever session is
|
||||
// current when the RPC unwinds.
|
||||
defer func() { sess.close() }()
|
||||
|
||||
// Callers block on the first Recv waiting for this and degrade to
|
||||
// non-live transcription when it does not arrive, so it goes out before
|
||||
// any audio is read.
|
||||
out <- &pb.TranscriptLiveResponse{Ready: true}
|
||||
|
||||
var (
|
||||
full strings.Builder
|
||||
flushing bool
|
||||
)
|
||||
emit := func(r streamResult) error {
|
||||
// Finals only, for the same reason as streamPCM: an interim is a
|
||||
// hypothesis the final rewrites, and delta is newly-finalized text.
|
||||
if !r.Final || (r.Text == "" && len(r.Words) == 0) {
|
||||
return nil
|
||||
}
|
||||
|
||||
// The separator goes INTO the delta, exactly as in streamPCM, because
|
||||
// the live consumer is the one that actually concatenates: the realtime
|
||||
// semantic-VAD path joins the accumulated deltas with the empty string
|
||||
// and only clears them at a turn reset, never at an endpoint. Adding
|
||||
// the space when assembling the terminal text instead would make the
|
||||
// running caption read "one.two." while the committed transcript read
|
||||
// "one. two.".
|
||||
delta := r.Text
|
||||
if delta != "" && full.Len() > 0 {
|
||||
delta = " " + delta
|
||||
}
|
||||
full.WriteString(delta)
|
||||
|
||||
out <- &pb.TranscriptLiveResponse{
|
||||
Delta: delta,
|
||||
// A final that arrives while audio is still coming IS the model's
|
||||
// endpoint: the decoder resets its utterance there and the next one
|
||||
// starts fresh, which is the turn boundary the realtime detector
|
||||
// waits on. The final that comes back from the tail flush is the
|
||||
// end of the STREAM, not a user yielding a turn, so it carries no
|
||||
// eou even though the send side has already closed.
|
||||
Eou: !flushing,
|
||||
Words: wordsToProto(r.Words),
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
for req := range in {
|
||||
switch payload := req.GetPayload().(type) {
|
||||
case *pb.TranscriptLiveRequest_Config:
|
||||
// A rate cannot change inside a stream (asr.h) and the decoder
|
||||
// keeps utterance state, so a reconfigure has to be a fresh
|
||||
// session rather than a reconfigured one.
|
||||
newRate, err := liveSampleRate(payload.Config)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
// Opened before the old one is closed so a failure here leaves a
|
||||
// live session for the deferred close, not a dangling handle.
|
||||
next, err := open(payload.Config.GetLanguage())
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
sess.close()
|
||||
sess, rate = next, newRate
|
||||
full.Reset()
|
||||
case *pb.TranscriptLiveRequest_Audio:
|
||||
pcm := payload.Audio.GetPcm()
|
||||
if len(pcm) == 0 {
|
||||
continue
|
||||
}
|
||||
if err := sess.push(pcm, rate); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := drain(sess, emit); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Send side closed: flush the tail and emit the terminal result. Like the
|
||||
// other backends' live path this carries Text only; per-utterance segments
|
||||
// and the duration are the file path's concern.
|
||||
flushing = true
|
||||
if err := sess.finish(); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := drain(sess, emit); err != nil {
|
||||
return err
|
||||
}
|
||||
// Not trimmed: the terminal text is the verbatim concatenation of the
|
||||
// deltas, which is the invariant the concatenating consumers rely on. The
|
||||
// first delta never carries the separator, so there is no leading space to
|
||||
// trim off in the first place.
|
||||
out <- &pb.TranscriptLiveResponse{
|
||||
FinalResult: &pb.TranscriptResult{Text: full.String()},
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// liveSampleRate resolves TranscriptLiveConfig.sample_rate to the rate the C
|
||||
// API is given. The proto's 0 means 16 kHz; the C API's 0 means "already at the
|
||||
// model rate", so the two cannot be forwarded to each other.
|
||||
func liveSampleRate(cfg *pb.TranscriptLiveConfig) (int32, error) {
|
||||
rate := cfg.GetSampleRate()
|
||||
if rate == 0 {
|
||||
return defaultLiveSampleRate, nil
|
||||
}
|
||||
if rate < minStreamSampleRate || rate > maxStreamSampleRate {
|
||||
return 0, status.Errorf(codes.InvalidArgument,
|
||||
"nemo-speech-cpp: unsupported live sample_rate %d (accepted: 0 or %d-%d Hz)",
|
||||
rate, minStreamSampleRate, maxStreamSampleRate)
|
||||
}
|
||||
return rate, nil
|
||||
}
|
||||
|
||||
// wordsToProto converts decoded words to the wire form. TranscriptWord.start
|
||||
// and .end are int64 nanoseconds; the runtime reports milliseconds.
|
||||
func wordsToProto(words []asrWord) []*pb.TranscriptWord {
|
||||
if len(words) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make([]*pb.TranscriptWord, len(words))
|
||||
for i, w := range words {
|
||||
out[i] = &pb.TranscriptWord{
|
||||
Text: w.Text,
|
||||
Start: msToNanos(w.Start),
|
||||
End: msToNanos(w.End),
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// AudioTranscriptionStream decodes the audio at req.Dst through the streaming
|
||||
// recognizer, emitting each finalized utterance as it lands.
|
||||
//
|
||||
// The body runs inside withEngine for the reason documented on withEngine, and
|
||||
// that holds engineMu for the whole stream: Free waits rather than destroying
|
||||
// the recognizer under a half-finished stream. streamPCM honours ctx so the
|
||||
// wait is bounded by the client's disconnect rather than by its silence.
|
||||
func (n *NemoSpeech) AudioTranscriptionStream(ctx context.Context, req *pb.TranscriptRequest, results chan *pb.TranscriptStreamResponse) error {
|
||||
// The host ranges over this channel and only returns once it closes, so
|
||||
// every path out of here, rejection included, has to close it.
|
||||
defer close(results)
|
||||
|
||||
return n.withEngine(familyASR, func() error {
|
||||
return n.transcribeStream(ctx, req, results)
|
||||
})
|
||||
}
|
||||
|
||||
// transcribeStream is AudioTranscriptionStream's body. The caller must hold
|
||||
// engineMu.
|
||||
func (n *NemoSpeech) transcribeStream(ctx context.Context, req *pb.TranscriptRequest, results chan<- *pb.TranscriptStreamResponse) error {
|
||||
if req.GetDst() == "" {
|
||||
return status.Error(codes.InvalidArgument,
|
||||
"nemo-speech-cpp: TranscriptRequest.dst (audio path) is required")
|
||||
}
|
||||
// Checked before the decode so a client that has already gone away does
|
||||
// not pay for an ffmpeg run, and so a cancellation is never reported as a
|
||||
// broken file.
|
||||
if err := ctx.Err(); err != nil {
|
||||
return status.Error(codes.Canceled, "nemo-speech-cpp: transcription cancelled")
|
||||
}
|
||||
|
||||
pcm, sampleRate, err := decodeAudioMono16k(req.GetDst())
|
||||
if err != nil {
|
||||
return status.Errorf(codes.InvalidArgument, "nemo-speech-cpp: read audio: %v", err)
|
||||
}
|
||||
// Before the session is opened, for the same reason as the offline path:
|
||||
// there is no transcript to be had from zero samples, and opening a stream
|
||||
// only to close it again asks the runtime to allocate decoder state for
|
||||
// nothing.
|
||||
if len(pcm) == 0 {
|
||||
return status.Error(codes.InvalidArgument, "nemo-speech-cpp: empty audio")
|
||||
}
|
||||
|
||||
sess, err := n.openSession(req.GetLanguage())
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer sess.close()
|
||||
|
||||
return streamPCM(ctx, sess, pcm, sampleRate,
|
||||
wordsRequested(req.GetTimestampGranularities()), results)
|
||||
}
|
||||
|
||||
// AudioTranscriptionLive serves the bidirectional live RPC over one streaming
|
||||
// session. See runLive for the protocol and withEngine for the locking.
|
||||
func (n *NemoSpeech) AudioTranscriptionLive(in <-chan *pb.TranscriptLiveRequest, out chan<- *pb.TranscriptLiveResponse) error {
|
||||
defer close(out)
|
||||
|
||||
return n.withEngine(familyASR, func() error {
|
||||
return runLive(n.openSession, in, out)
|
||||
})
|
||||
}
|
||||
@@ -1,699 +0,0 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"path/filepath"
|
||||
"time"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
"google.golang.org/grpc/codes"
|
||||
"google.golang.org/grpc/status"
|
||||
|
||||
pb "github.com/mudler/LocalAI/pkg/grpc/proto"
|
||||
)
|
||||
|
||||
// fakeSession is a scripted asrSession. It stands in for the streaming C API,
|
||||
// not for a model: no NeMo GGUF is small enough to keep in the tree, and the
|
||||
// need-more-audio drain is the easiest thing in this file to get subtly wrong
|
||||
// (a mishandled NULL either spins forever or drops every result).
|
||||
//
|
||||
// script is one batch of results per drain. next() hands back the current
|
||||
// batch one result at a time and then reports "need more audio" exactly once,
|
||||
// which advances to the next batch. That is precisely the C contract:
|
||||
// nemo_speech_asr_stream_next returns OK with a NULL handle when the decoder
|
||||
// has consumed the buffered audio, and the loop must resume after the next
|
||||
// push rather than treat it as the end of the stream.
|
||||
type fakeSession struct {
|
||||
script [][]streamResult
|
||||
batch int
|
||||
pos int
|
||||
|
||||
pushed [][]float32
|
||||
rates []int32
|
||||
finished int
|
||||
closed int
|
||||
|
||||
pushErr error
|
||||
finishErr error
|
||||
nextErr error
|
||||
}
|
||||
|
||||
func (f *fakeSession) push(pcm []float32, sampleRate int32) error {
|
||||
if f.pushErr != nil {
|
||||
return f.pushErr
|
||||
}
|
||||
f.pushed = append(f.pushed, pcm)
|
||||
f.rates = append(f.rates, sampleRate)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (f *fakeSession) finish() error {
|
||||
if f.finishErr != nil {
|
||||
return f.finishErr
|
||||
}
|
||||
f.finished++
|
||||
return nil
|
||||
}
|
||||
|
||||
func (f *fakeSession) next() (streamResult, bool, error) {
|
||||
if f.nextErr != nil {
|
||||
return streamResult{}, false, f.nextErr
|
||||
}
|
||||
if f.batch >= len(f.script) {
|
||||
return streamResult{}, false, nil
|
||||
}
|
||||
if f.pos >= len(f.script[f.batch]) {
|
||||
f.batch++
|
||||
f.pos = 0
|
||||
return streamResult{}, false, nil
|
||||
}
|
||||
r := f.script[f.batch][f.pos]
|
||||
f.pos++
|
||||
return r, true, nil
|
||||
}
|
||||
|
||||
func (f *fakeSession) close() { f.closed++ }
|
||||
|
||||
// samples returns the flat concatenation of everything pushed, so a spec can
|
||||
// assert the whole clip reached the engine without caring how it was sliced.
|
||||
func (f *fakeSession) samples() []float32 {
|
||||
var out []float32
|
||||
for _, c := range f.pushed {
|
||||
out = append(out, c...)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// collect drains a response channel into a slice. The channels are unbuffered
|
||||
// in the specs on purpose: a producer that stops honouring cancellation would
|
||||
// otherwise fill a buffer and look healthy.
|
||||
func collect[T any](ch chan T) chan []T {
|
||||
done := make(chan []T, 1)
|
||||
go func() {
|
||||
var got []T
|
||||
for v := range ch {
|
||||
got = append(got, v)
|
||||
}
|
||||
done <- got
|
||||
}()
|
||||
return done
|
||||
}
|
||||
|
||||
var _ = Describe("chunkPCM", func() {
|
||||
It("splits into equal chunks when evenly divisible", func() {
|
||||
chunks := chunkPCM(make([]float32, 400), 100)
|
||||
Expect(chunks).To(HaveLen(4))
|
||||
for _, c := range chunks {
|
||||
Expect(c).To(HaveLen(100))
|
||||
}
|
||||
})
|
||||
|
||||
// Padding the tail with silence would push phantom audio through the
|
||||
// encoder and shift the tail word timings, so the final chunk stays short.
|
||||
It("makes the final chunk short rather than padding it", func() {
|
||||
chunks := chunkPCM(make([]float32, 250), 100)
|
||||
Expect(chunks).To(HaveLen(3))
|
||||
Expect(chunks[2]).To(HaveLen(50))
|
||||
})
|
||||
|
||||
It("returns one chunk when the input is shorter than the chunk size", func() {
|
||||
chunks := chunkPCM(make([]float32, 10), 100)
|
||||
Expect(chunks).To(HaveLen(1))
|
||||
Expect(chunks[0]).To(HaveLen(10))
|
||||
})
|
||||
|
||||
It("returns nothing for empty input", func() {
|
||||
Expect(chunkPCM(nil, 100)).To(BeEmpty())
|
||||
Expect(chunkPCM([]float32{}, 100)).To(BeEmpty())
|
||||
})
|
||||
|
||||
// Every spec above works on all-zero audio, so none of them can tell a
|
||||
// correct slicing from one that reorders or repeats windows. Audio fed out
|
||||
// of order still decodes, it just decodes to nonsense.
|
||||
It("preserves sample order across the chunk boundaries", func() {
|
||||
pcm := []float32{1, 2, 3, 4, 5}
|
||||
chunks := chunkPCM(pcm, 2)
|
||||
Expect(chunks).To(HaveLen(3))
|
||||
Expect(chunks[0]).To(Equal([]float32{1, 2}))
|
||||
Expect(chunks[1]).To(Equal([]float32{3, 4}))
|
||||
Expect(chunks[2]).To(Equal([]float32{5}))
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("drain", func() {
|
||||
It("emits every result in a batch and stops on need-more-audio", func() {
|
||||
sess := &fakeSession{script: [][]streamResult{
|
||||
{{Text: "a"}, {Text: "b", Final: true}},
|
||||
{{Text: "c"}},
|
||||
}}
|
||||
var got []string
|
||||
Expect(drain(sess, func(r streamResult) error {
|
||||
got = append(got, r.Text)
|
||||
return nil
|
||||
})).To(Succeed())
|
||||
Expect(got).To(Equal([]string{"a", "b"}))
|
||||
})
|
||||
|
||||
// The NULL handle is a pause, not an end: the next drain, after more audio
|
||||
// has been pushed, must pick the stream back up.
|
||||
It("resumes on the next drain after a need-more-audio pause", func() {
|
||||
sess := &fakeSession{script: [][]streamResult{{{Text: "a"}}, {{Text: "b"}}}}
|
||||
var got []string
|
||||
emit := func(r streamResult) error { got = append(got, r.Text); return nil }
|
||||
Expect(drain(sess, emit)).To(Succeed())
|
||||
Expect(drain(sess, emit)).To(Succeed())
|
||||
Expect(got).To(Equal([]string{"a", "b"}))
|
||||
})
|
||||
|
||||
It("returns nothing and no error for a stream with no results ready", func() {
|
||||
var got []string
|
||||
Expect(drain(&fakeSession{}, func(r streamResult) error {
|
||||
got = append(got, r.Text)
|
||||
return nil
|
||||
})).To(Succeed())
|
||||
Expect(got).To(BeEmpty())
|
||||
})
|
||||
|
||||
It("propagates a failure from the runtime", func() {
|
||||
sess := &fakeSession{nextErr: errors.New("boom")}
|
||||
Expect(drain(sess, func(streamResult) error { return nil })).To(MatchError(ContainSubstring("boom")))
|
||||
})
|
||||
|
||||
It("stops pulling once emit fails", func() {
|
||||
sess := &fakeSession{script: [][]streamResult{{{Text: "a"}, {Text: "b"}}}}
|
||||
Expect(drain(sess, func(streamResult) error {
|
||||
return errors.New("send failed")
|
||||
})).To(MatchError(ContainSubstring("send failed")))
|
||||
Expect(sess.pos).To(Equal(1))
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("streamPCM", func() {
|
||||
streamWords := func(ctx context.Context, sess asrSession, pcm []float32, rate int32, wantWords bool) ([]*pb.TranscriptStreamResponse, error) {
|
||||
GinkgoHelper()
|
||||
results := make(chan *pb.TranscriptStreamResponse)
|
||||
done := collect(results)
|
||||
err := streamPCM(ctx, sess, pcm, rate, wantWords, results)
|
||||
close(results)
|
||||
return <-done, err
|
||||
}
|
||||
stream := func(ctx context.Context, sess asrSession, pcm []float32, rate int32) ([]*pb.TranscriptStreamResponse, error) {
|
||||
GinkgoHelper()
|
||||
return streamWords(ctx, sess, pcm, rate, false)
|
||||
}
|
||||
|
||||
It("pushes the whole clip in chunks at the clip's own sample rate", func() {
|
||||
sess := &fakeSession{}
|
||||
pcm := make([]float32, streamChunkSamples*2+7)
|
||||
_, err := stream(context.Background(), sess, pcm, 16000)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(sess.pushed).To(HaveLen(3))
|
||||
Expect(sess.samples()).To(HaveLen(len(pcm)))
|
||||
for _, r := range sess.rates {
|
||||
Expect(r).To(Equal(int32(16000)))
|
||||
}
|
||||
})
|
||||
|
||||
It("finishes the stream once, after the last chunk", func() {
|
||||
sess := &fakeSession{}
|
||||
_, err := stream(context.Background(), sess, make([]float32, 10), 16000)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(sess.finished).To(Equal(1))
|
||||
})
|
||||
|
||||
// Interims are the decoder's running hypothesis for the utterance in
|
||||
// flight. The wire contract is that delta is newly FINALIZED text and that
|
||||
// concatenating the deltas reproduces the transcript, so forwarding an
|
||||
// interim would duplicate every word it later re-sends inside the final.
|
||||
It("emits a delta per final and nothing for interims", func() {
|
||||
sess := &fakeSession{script: [][]streamResult{{
|
||||
{Text: "hel"},
|
||||
{Text: "hello"},
|
||||
{Text: "Hello.", Final: true},
|
||||
}}}
|
||||
got, err := stream(context.Background(), sess, make([]float32, 10), 16000)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
var deltas []string
|
||||
for _, r := range got {
|
||||
if r.GetDelta() != "" {
|
||||
deltas = append(deltas, r.GetDelta())
|
||||
}
|
||||
}
|
||||
Expect(deltas).To(Equal([]string{"Hello."}))
|
||||
})
|
||||
|
||||
It("reproduces the final transcript by concatenating the deltas", func() {
|
||||
sess := &fakeSession{script: [][]streamResult{{
|
||||
{Text: "One.", Final: true},
|
||||
{Text: "Two.", Final: true},
|
||||
}}}
|
||||
got, err := stream(context.Background(), sess, make([]float32, 10), 16000)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
var joined string
|
||||
var final *pb.TranscriptResult
|
||||
for _, r := range got {
|
||||
joined += r.GetDelta()
|
||||
if r.GetFinalResult() != nil {
|
||||
final = r.GetFinalResult()
|
||||
}
|
||||
}
|
||||
Expect(final).ToNot(BeNil())
|
||||
Expect(final.GetText()).To(Equal("One. Two."))
|
||||
Expect(joined).To(Equal(final.GetText()))
|
||||
})
|
||||
|
||||
It("sends the terminal final result last and only once", func() {
|
||||
sess := &fakeSession{script: [][]streamResult{{{Text: "hi", Final: true}}}}
|
||||
got, err := stream(context.Background(), sess, make([]float32, 10), 16000)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(got).ToNot(BeEmpty())
|
||||
|
||||
var finals int
|
||||
for _, r := range got {
|
||||
if r.GetFinalResult() != nil {
|
||||
finals++
|
||||
}
|
||||
}
|
||||
Expect(finals).To(Equal(1))
|
||||
Expect(got[len(got)-1].GetFinalResult()).ToNot(BeNil())
|
||||
})
|
||||
|
||||
It("reports the clip duration in seconds", func() {
|
||||
sess := &fakeSession{}
|
||||
got, err := stream(context.Background(), sess, make([]float32, 8000), 16000)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(got[len(got)-1].GetFinalResult().GetDuration()).To(BeNumerically("~", 0.5, 1e-6))
|
||||
})
|
||||
|
||||
It("builds per-utterance segments with nanosecond timestamps", func() {
|
||||
sess := &fakeSession{script: [][]streamResult{{
|
||||
{Text: "one", Final: true, Words: []asrWord{{Text: "one", Start: 0, End: 500}}},
|
||||
{Text: "two", Final: true, Words: []asrWord{{Text: "two", Start: 900, End: 1400}}},
|
||||
}}}
|
||||
got, err := stream(context.Background(), sess, make([]float32, 10), 16000)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
segs := got[len(got)-1].GetFinalResult().GetSegments()
|
||||
Expect(segs).To(HaveLen(2))
|
||||
Expect(segs[0].GetId()).To(Equal(int32(0)))
|
||||
Expect(segs[1].GetId()).To(Equal(int32(1)))
|
||||
Expect(time.Duration(segs[1].GetStart())).To(Equal(900 * time.Millisecond))
|
||||
Expect(time.Duration(segs[1].GetEnd())).To(Equal(1400 * time.Millisecond))
|
||||
})
|
||||
|
||||
// core/backend/transcript.go builds the response's word list out of
|
||||
// TranscriptSegment.Words, so leaving it unset makes
|
||||
// timestamp_granularities: ["word"] come back empty.
|
||||
It("attaches the word timings only when they were asked for", func() {
|
||||
script := func() [][]streamResult {
|
||||
return [][]streamResult{{{Text: "one", Final: true,
|
||||
Words: []asrWord{{Text: "one", Start: 100, End: 500}}}}}
|
||||
}
|
||||
|
||||
got, err := streamWords(context.Background(), &fakeSession{script: script()}, make([]float32, 10), 16000, true)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
segs := got[len(got)-1].GetFinalResult().GetSegments()
|
||||
Expect(segs[0].GetWords()).To(HaveLen(1))
|
||||
Expect(segs[0].GetWords()[0].GetText()).To(Equal("one"))
|
||||
Expect(time.Duration(segs[0].GetWords()[0].GetStart())).To(Equal(100 * time.Millisecond))
|
||||
|
||||
got, err = streamWords(context.Background(), &fakeSession{script: script()}, make([]float32, 10), 16000, false)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
segs = got[len(got)-1].GetFinalResult().GetSegments()
|
||||
Expect(segs[0].GetText()).To(Equal("one"))
|
||||
Expect(segs[0].GetWords()).To(BeEmpty())
|
||||
})
|
||||
|
||||
// The flush that nemo_speech_asr_stream_finish triggers returns whatever the
|
||||
// decoder was still holding. Nothing held back means the last endpoint
|
||||
// consumed the audio, which is exactly "the clip ended on an utterance
|
||||
// boundary"; text coming back means it ended mid-utterance.
|
||||
It("marks eou when the tail flush had nothing left to emit", func() {
|
||||
sess := &fakeSession{script: [][]streamResult{
|
||||
{{Text: "done.", Final: true}},
|
||||
{{Text: "", Final: true}},
|
||||
}}
|
||||
got, err := stream(context.Background(), sess, make([]float32, 10), 16000)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(got[len(got)-1].GetFinalResult().GetEou()).To(BeTrue())
|
||||
})
|
||||
|
||||
It("does not mark eou when the tail flush produced text", func() {
|
||||
sess := &fakeSession{script: [][]streamResult{
|
||||
{{Text: "done.", Final: true}},
|
||||
{{Text: "and more", Final: true}},
|
||||
}}
|
||||
got, err := stream(context.Background(), sess, make([]float32, 10), 16000)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(got[len(got)-1].GetFinalResult().GetEou()).To(BeFalse())
|
||||
})
|
||||
|
||||
// The RPC body runs inside withEngine, so it holds the engine mutex for the
|
||||
// whole stream and Free waits on it. A loop that ignored cancellation would
|
||||
// pin the model against unload for as long as a disconnected client's audio
|
||||
// takes to push.
|
||||
It("stops promptly when the request context is cancelled", func() {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
sess := &fakeSession{}
|
||||
_, err := stream(ctx, sess, make([]float32, streamChunkSamples*4), 16000)
|
||||
Expect(status.Code(err)).To(Equal(codes.Canceled))
|
||||
Expect(sess.pushed).To(BeEmpty())
|
||||
})
|
||||
|
||||
It("reports a push failure", func() {
|
||||
sess := &fakeSession{pushErr: errors.New("push blew up")}
|
||||
_, err := stream(context.Background(), sess, make([]float32, 10), 16000)
|
||||
Expect(err).To(MatchError(ContainSubstring("push blew up")))
|
||||
})
|
||||
|
||||
It("reports a finish failure", func() {
|
||||
sess := &fakeSession{finishErr: errors.New("finish blew up")}
|
||||
_, err := stream(context.Background(), sess, make([]float32, 10), 16000)
|
||||
Expect(err).To(MatchError(ContainSubstring("finish blew up")))
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("runLive", func() {
|
||||
// live drives runLive against a fake opener and returns everything the RPC
|
||||
// wrote plus the sessions it opened.
|
||||
live := func(reqs []*pb.TranscriptLiveRequest, script ...[][]streamResult) ([]*pb.TranscriptLiveResponse, []*fakeSession, error) {
|
||||
GinkgoHelper()
|
||||
var opened []*fakeSession
|
||||
open := func(language string) (asrSession, error) {
|
||||
s := &fakeSession{}
|
||||
if len(opened) < len(script) {
|
||||
s.script = script[len(opened)]
|
||||
}
|
||||
opened = append(opened, s)
|
||||
return s, nil
|
||||
}
|
||||
|
||||
in := make(chan *pb.TranscriptLiveRequest)
|
||||
out := make(chan *pb.TranscriptLiveResponse)
|
||||
done := collect(out)
|
||||
go func() {
|
||||
defer close(in)
|
||||
for _, r := range reqs {
|
||||
in <- r
|
||||
}
|
||||
}()
|
||||
err := runLive(open, in, out)
|
||||
close(out)
|
||||
return <-done, opened, err
|
||||
}
|
||||
|
||||
cfg := func(rate int32) *pb.TranscriptLiveRequest {
|
||||
return &pb.TranscriptLiveRequest{Payload: &pb.TranscriptLiveRequest_Config{
|
||||
Config: &pb.TranscriptLiveConfig{SampleRate: rate},
|
||||
}}
|
||||
}
|
||||
audio := func(pcm ...float32) *pb.TranscriptLiveRequest {
|
||||
return &pb.TranscriptLiveRequest{Payload: &pb.TranscriptLiveRequest_Audio{
|
||||
Audio: &pb.TranscriptLiveAudio{Pcm: pcm},
|
||||
}}
|
||||
}
|
||||
|
||||
It("requires the first message to carry a config", func() {
|
||||
_, opened, err := live([]*pb.TranscriptLiveRequest{audio(1, 2, 3)})
|
||||
Expect(status.Code(err)).To(Equal(codes.InvalidArgument))
|
||||
Expect(opened).To(BeEmpty())
|
||||
})
|
||||
|
||||
It("returns without error when the caller closes without sending anything", func() {
|
||||
got, opened, err := live(nil)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(got).To(BeEmpty())
|
||||
Expect(opened).To(BeEmpty())
|
||||
})
|
||||
|
||||
// Callers block on the first Recv waiting for this ack, and degrade to
|
||||
// non-live transcription when it does not arrive.
|
||||
It("acknowledges a successful open before any transcript", func() {
|
||||
got, _, err := live([]*pb.TranscriptLiveRequest{cfg(0)})
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(got).ToNot(BeEmpty())
|
||||
Expect(got[0].GetReady()).To(BeTrue())
|
||||
})
|
||||
|
||||
// The proto documents 0 as "16 kHz". The C API reads 0 as "these samples
|
||||
// are already at the model rate" and skips resampling, so forwarding the
|
||||
// zero through would silently mean something else.
|
||||
It("resolves the default sample rate to 16 kHz before pushing", func() {
|
||||
_, opened, err := live([]*pb.TranscriptLiveRequest{cfg(0), audio(1, 2, 3)})
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(opened).To(HaveLen(1))
|
||||
Expect(opened[0].rates).To(Equal([]int32{16000}))
|
||||
})
|
||||
|
||||
It("pushes at the configured sample rate", func() {
|
||||
_, opened, err := live([]*pb.TranscriptLiveRequest{cfg(8000), audio(1, 2, 3)})
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(opened[0].rates).To(Equal([]int32{8000}))
|
||||
Expect(opened[0].samples()).To(Equal([]float32{1, 2, 3}))
|
||||
})
|
||||
|
||||
It("rejects a sample rate the runtime cannot resample", func() {
|
||||
_, opened, err := live([]*pb.TranscriptLiveRequest{cfg(4000)})
|
||||
Expect(status.Code(err)).To(Equal(codes.InvalidArgument))
|
||||
Expect(opened).To(BeEmpty())
|
||||
})
|
||||
|
||||
It("ignores an empty audio frame instead of pushing it", func() {
|
||||
_, opened, err := live([]*pb.TranscriptLiveRequest{cfg(0), audio()})
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(opened[0].pushed).To(BeEmpty())
|
||||
})
|
||||
|
||||
It("streams a delta with its words and marks the utterance boundary", func() {
|
||||
got, _, err := live(
|
||||
[]*pb.TranscriptLiveRequest{cfg(0), audio(1)},
|
||||
[][]streamResult{{
|
||||
{Text: "partial"},
|
||||
{Text: "Hello there.", Final: true, Words: []asrWord{
|
||||
{Text: "Hello", Start: 100, End: 400},
|
||||
{Text: "there", Start: 400, End: 900},
|
||||
}},
|
||||
}},
|
||||
)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
var deltas []*pb.TranscriptLiveResponse
|
||||
for _, r := range got {
|
||||
if r.GetDelta() != "" {
|
||||
deltas = append(deltas, r)
|
||||
}
|
||||
}
|
||||
Expect(deltas).To(HaveLen(1))
|
||||
Expect(deltas[0].GetDelta()).To(Equal("Hello there."))
|
||||
Expect(deltas[0].GetEou()).To(BeTrue())
|
||||
Expect(deltas[0].GetWords()).To(HaveLen(2))
|
||||
Expect(time.Duration(deltas[0].GetWords()[1].GetStart())).To(Equal(400 * time.Millisecond))
|
||||
Expect(time.Duration(deltas[0].GetWords()[1].GetEnd())).To(Equal(900 * time.Millisecond))
|
||||
})
|
||||
|
||||
It("finishes and closes the session when the caller closes the send side", func() {
|
||||
got, opened, err := live(
|
||||
[]*pb.TranscriptLiveRequest{cfg(0), audio(1)},
|
||||
[][]streamResult{{{Text: "one.", Final: true}}, {{Text: "two.", Final: true}}},
|
||||
)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(opened[0].finished).To(Equal(1))
|
||||
Expect(opened[0].closed).To(Equal(1))
|
||||
Expect(got[len(got)-1].GetFinalResult()).ToNot(BeNil())
|
||||
Expect(got[len(got)-1].GetFinalResult().GetText()).To(Equal("one. two."))
|
||||
})
|
||||
|
||||
// The live path is the one with a consumer that really concatenates: the
|
||||
// realtime semantic-VAD path joins the accumulated deltas with the empty
|
||||
// string and clears them only at a turn reset, never at an utterance
|
||||
// boundary. A separator added when assembling the terminal text instead of
|
||||
// inside the delta makes the running caption read "one.two." while the
|
||||
// committed transcript reads "one. two.".
|
||||
It("reproduces the final transcript by concatenating the deltas", func() {
|
||||
got, _, err := live(
|
||||
[]*pb.TranscriptLiveRequest{cfg(0), audio(1)},
|
||||
[][]streamResult{{{Text: "one.", Final: true}}, {{Text: "two.", Final: true}}},
|
||||
)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
var joined string
|
||||
var final *pb.TranscriptResult
|
||||
for _, r := range got {
|
||||
joined += r.GetDelta()
|
||||
if r.GetFinalResult() != nil {
|
||||
final = r.GetFinalResult()
|
||||
}
|
||||
}
|
||||
Expect(final).ToNot(BeNil())
|
||||
Expect(final.GetText()).To(Equal("one. two."))
|
||||
Expect(joined).To(Equal(final.GetText()))
|
||||
})
|
||||
|
||||
// Eou is the model's endpoint, which is a user yielding the turn. The final
|
||||
// that comes back from the tail flush is the end of the stream: the send
|
||||
// side has already closed, so reporting a turn boundary there tells the
|
||||
// turn detector something that did not happen.
|
||||
It("marks the endpoint finals but not the tail flush", func() {
|
||||
got, _, err := live(
|
||||
[]*pb.TranscriptLiveRequest{cfg(0), audio(1)},
|
||||
[][]streamResult{{{Text: "one.", Final: true}}, {{Text: "two.", Final: true}}},
|
||||
)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
var eous []bool
|
||||
for _, r := range got {
|
||||
if r.GetDelta() != "" {
|
||||
eous = append(eous, r.GetEou())
|
||||
}
|
||||
}
|
||||
Expect(eous).To(Equal([]bool{true, false}))
|
||||
})
|
||||
|
||||
// A rate cannot change inside a stream and the decoder keeps no state
|
||||
// across a reset, so a second config has to be a fresh session, not a
|
||||
// reconfigured one.
|
||||
It("opens a fresh session on a mid-stream config and drops the old transcript", func() {
|
||||
got, opened, err := live(
|
||||
[]*pb.TranscriptLiveRequest{cfg(0), audio(1), cfg(0), audio(2)},
|
||||
[][]streamResult{{{Text: "dropped.", Final: true}}},
|
||||
[][]streamResult{{{Text: "kept.", Final: true}}},
|
||||
)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(opened).To(HaveLen(2))
|
||||
Expect(opened[0].closed).To(Equal(1))
|
||||
Expect(got[len(got)-1].GetFinalResult().GetText()).To(Equal("kept."))
|
||||
})
|
||||
|
||||
It("reports a push failure and still closes the session", func() {
|
||||
var opened []*fakeSession
|
||||
open := func(string) (asrSession, error) {
|
||||
s := &fakeSession{pushErr: errors.New("push blew up")}
|
||||
opened = append(opened, s)
|
||||
return s, nil
|
||||
}
|
||||
in := make(chan *pb.TranscriptLiveRequest, 2)
|
||||
in <- cfg(0)
|
||||
in <- audio(1, 2)
|
||||
close(in)
|
||||
out := make(chan *pb.TranscriptLiveResponse, 8)
|
||||
|
||||
err := runLive(open, in, out)
|
||||
Expect(err).To(MatchError(ContainSubstring("push blew up")))
|
||||
Expect(opened[0].closed).To(Equal(1))
|
||||
})
|
||||
|
||||
It("propagates a failure to open the session", func() {
|
||||
open := func(string) (asrSession, error) { return nil, errors.New("no streaming here") }
|
||||
in := make(chan *pb.TranscriptLiveRequest, 1)
|
||||
in <- cfg(0)
|
||||
close(in)
|
||||
out := make(chan *pb.TranscriptLiveResponse, 8)
|
||||
Expect(runLive(open, in, out)).To(MatchError(ContainSubstring("no streaming here")))
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("AudioTranscriptionStream", func() {
|
||||
run := func(ctx context.Context, n *NemoSpeech, req *pb.TranscriptRequest) ([]*pb.TranscriptStreamResponse, error) {
|
||||
GinkgoHelper()
|
||||
results := make(chan *pb.TranscriptStreamResponse)
|
||||
done := collect(results)
|
||||
err := n.AudioTranscriptionStream(ctx, req, results)
|
||||
return <-done, err
|
||||
}
|
||||
|
||||
// The RPC owns the channel: the gRPC host ranges over it and only returns
|
||||
// once it closes, so a rejection path that forgets to close hangs the call
|
||||
// instead of failing it.
|
||||
It("closes the results channel on every rejection path", func() {
|
||||
for _, n := range []*NemoSpeech{{fam: familyTTS}, {fam: familyASR}, {}} {
|
||||
_, err := run(context.Background(), n, &pb.TranscriptRequest{})
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(n.engineMu.TryLock()).To(BeTrue())
|
||||
n.engineMu.Unlock()
|
||||
}
|
||||
})
|
||||
|
||||
It("refuses a model loaded as another family", func() {
|
||||
n := &NemoSpeech{fam: familyTTS}
|
||||
_, err := run(context.Background(), n, &pb.TranscriptRequest{Dst: "x.wav"})
|
||||
Expect(status.Code(err)).To(Equal(codes.Unimplemented))
|
||||
Expect(err.Error()).To(ContainSubstring("tts"))
|
||||
})
|
||||
|
||||
It("requires a destination path", func() {
|
||||
n := &NemoSpeech{fam: familyASR}
|
||||
_, err := run(context.Background(), n, &pb.TranscriptRequest{})
|
||||
Expect(status.Code(err)).To(Equal(codes.InvalidArgument))
|
||||
})
|
||||
|
||||
// Cancellation is checked before the decode so a client that has already
|
||||
// gone away does not pay for an ffmpeg run, and so the check cannot be
|
||||
// mistaken for the decode failing.
|
||||
It("returns cancelled without touching the audio", func() {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
n := &NemoSpeech{fam: familyASR}
|
||||
_, err := run(ctx, n, &pb.TranscriptRequest{
|
||||
Dst: filepath.Join(GinkgoT().TempDir(), "absent.wav"),
|
||||
})
|
||||
Expect(status.Code(err)).To(Equal(codes.Canceled))
|
||||
})
|
||||
|
||||
It("reports an audio file it cannot read", func() {
|
||||
n := &NemoSpeech{fam: familyASR}
|
||||
_, err := run(context.Background(), n, &pb.TranscriptRequest{
|
||||
Dst: filepath.Join(GinkgoT().TempDir(), "absent.wav"),
|
||||
})
|
||||
Expect(status.Code(err)).To(Equal(codes.InvalidArgument))
|
||||
})
|
||||
|
||||
// Same ordering constraint as the offline path: a clip that decodes to no
|
||||
// samples has to be refused before a session is opened, which is also
|
||||
// before any bound entry point is called. Nothing is loaded here, so a
|
||||
// guard placed after the open would panic instead of failing.
|
||||
It("refuses a decodable clip that carries no samples, before opening a session", func() {
|
||||
path := filepath.Join(GinkgoT().TempDir(), "silence.wav")
|
||||
writeMono16kWAV(path, 0)
|
||||
|
||||
n := &NemoSpeech{fam: familyASR}
|
||||
var err error
|
||||
Expect(func() {
|
||||
_, err = run(context.Background(), n, &pb.TranscriptRequest{Dst: path})
|
||||
}).ToNot(Panic())
|
||||
Expect(status.Code(err)).To(Equal(codes.InvalidArgument))
|
||||
Expect(err.Error()).To(ContainSubstring("empty audio"))
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("AudioTranscriptionLive", func() {
|
||||
It("refuses a model loaded as another family and closes the output", func() {
|
||||
n := &NemoSpeech{fam: familyNMT}
|
||||
in := make(chan *pb.TranscriptLiveRequest)
|
||||
close(in)
|
||||
out := make(chan *pb.TranscriptLiveResponse)
|
||||
done := collect(out)
|
||||
|
||||
err := n.AudioTranscriptionLive(in, out)
|
||||
Expect(status.Code(err)).To(Equal(codes.Unimplemented))
|
||||
Expect(<-done).To(BeEmpty())
|
||||
Expect(n.engineMu.TryLock()).To(BeTrue())
|
||||
n.engineMu.Unlock()
|
||||
})
|
||||
|
||||
It("refuses an unloaded model", func() {
|
||||
n := &NemoSpeech{}
|
||||
in := make(chan *pb.TranscriptLiveRequest)
|
||||
close(in)
|
||||
out := make(chan *pb.TranscriptLiveResponse)
|
||||
done := collect(out)
|
||||
|
||||
err := n.AudioTranscriptionLive(in, out)
|
||||
Expect(status.Code(err)).To(Equal(codes.Unimplemented))
|
||||
Expect(<-done).To(BeEmpty())
|
||||
})
|
||||
})
|
||||
@@ -1,367 +0,0 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"math"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"time"
|
||||
"unsafe"
|
||||
|
||||
"github.com/go-audio/audio"
|
||||
"github.com/go-audio/wav"
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
"google.golang.org/grpc/codes"
|
||||
"google.golang.org/grpc/status"
|
||||
|
||||
pb "github.com/mudler/LocalAI/pkg/grpc/proto"
|
||||
)
|
||||
|
||||
// writeMono16kWAV writes `frames` samples of 16 kHz mono 16-bit silence.
|
||||
// That is already AudioToWav's target format, so the decode path copies the
|
||||
// file through instead of shelling out to ffmpeg, which the test host may not
|
||||
// have.
|
||||
func writeMono16kWAV(path string, frames int) {
|
||||
GinkgoHelper()
|
||||
f, err := os.Create(path)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
enc := wav.NewEncoder(f, 16000, 16, 1, 1)
|
||||
Expect(enc.Write(&audio.IntBuffer{
|
||||
Format: &audio.Format{NumChannels: 1, SampleRate: 16000},
|
||||
SourceBitDepth: 16,
|
||||
Data: make([]int, frames),
|
||||
})).To(Succeed())
|
||||
Expect(enc.Close()).To(Succeed())
|
||||
Expect(f.Close()).To(Succeed())
|
||||
}
|
||||
|
||||
var _ = Describe("wordsToSegments", func() {
|
||||
It("groups words into one segment per speaker run", func() {
|
||||
words := []asrWord{
|
||||
{Text: "hello", Start: 0, End: 400, Speaker: 1},
|
||||
{Text: "there", Start: 400, End: 800, Speaker: 1},
|
||||
{Text: "hi", Start: 900, End: 1200, Speaker: 2},
|
||||
}
|
||||
segs := wordsToSegments(words, false)
|
||||
Expect(segs).To(HaveLen(2))
|
||||
Expect(segs[0].Text).To(Equal("hello there"))
|
||||
Expect(segs[1].Text).To(Equal("hi"))
|
||||
})
|
||||
|
||||
// A run is bounded by a CHANGE of speaker, not by the speaker id being new.
|
||||
// Grouping that keyed on the id itself (a map, or a comparison against the
|
||||
// first word) would merge the two A turns into one segment spanning B, and
|
||||
// the three-word spec above cannot see that because it never returns to an
|
||||
// earlier speaker.
|
||||
It("starts a new segment when an earlier speaker takes another turn", func() {
|
||||
words := []asrWord{
|
||||
{Text: "one", Start: 0, End: 100, Speaker: 1},
|
||||
{Text: "two", Start: 100, End: 200, Speaker: 2},
|
||||
{Text: "three", Start: 200, End: 300, Speaker: 1},
|
||||
}
|
||||
segs := wordsToSegments(words, false)
|
||||
Expect(segs).To(HaveLen(3))
|
||||
Expect(segs[0].Text).To(Equal("one"))
|
||||
Expect(segs[1].Text).To(Equal("two"))
|
||||
Expect(segs[2].Text).To(Equal("three"))
|
||||
})
|
||||
|
||||
// TranscriptSegment.start/end are int64 nanoseconds, not seconds:
|
||||
// core/backend/transcript.go reads them straight into a time.Duration. The
|
||||
// runtime reports word offsets in milliseconds (src/asr/types.h:46).
|
||||
It("converts millisecond word times to nanoseconds", func() {
|
||||
words := []asrWord{{Text: "a", Start: 1500, End: 2250, Speaker: 0}}
|
||||
segs := wordsToSegments(words, false)
|
||||
Expect(segs).To(HaveLen(1))
|
||||
Expect(time.Duration(segs[0].Start)).To(Equal(1500 * time.Millisecond))
|
||||
Expect(time.Duration(segs[0].End)).To(Equal(2250 * time.Millisecond))
|
||||
})
|
||||
|
||||
It("spans a segment from its first word's start to its last word's end", func() {
|
||||
words := []asrWord{
|
||||
{Text: "a", Start: 100, End: 200, Speaker: 0},
|
||||
{Text: "b", Start: 500, End: 900, Speaker: 0},
|
||||
}
|
||||
segs := wordsToSegments(words, false)
|
||||
Expect(segs).To(HaveLen(1))
|
||||
Expect(time.Duration(segs[0].Start)).To(Equal(100 * time.Millisecond))
|
||||
Expect(time.Duration(segs[0].End)).To(Equal(900 * time.Millisecond))
|
||||
})
|
||||
|
||||
It("produces a single segment when no speaker tags are present", func() {
|
||||
words := []asrWord{
|
||||
{Text: "a", Start: 0, End: 100, Speaker: 0},
|
||||
{Text: "b", Start: 100, End: 200, Speaker: 0},
|
||||
}
|
||||
segs := wordsToSegments(words, false)
|
||||
Expect(segs).To(HaveLen(1))
|
||||
Expect(segs[0].Text).To(Equal("a b"))
|
||||
})
|
||||
|
||||
It("returns no segments for no words", func() {
|
||||
Expect(wordsToSegments(nil, false)).To(BeEmpty())
|
||||
Expect(wordsToSegments([]asrWord{}, false)).To(BeEmpty())
|
||||
})
|
||||
|
||||
It("numbers the segments from zero in order", func() {
|
||||
words := []asrWord{
|
||||
{Text: "a", Speaker: 1},
|
||||
{Text: "b", Speaker: 2},
|
||||
{Text: "c", Speaker: 3},
|
||||
}
|
||||
segs := wordsToSegments(words, false)
|
||||
Expect(segs).To(HaveLen(3))
|
||||
for i, s := range segs {
|
||||
Expect(s.Id).To(Equal(int32(i)))
|
||||
}
|
||||
})
|
||||
|
||||
// TranscriptSegment.Words is what core/backend/transcript.go turns into the
|
||||
// response's word list, so an unset one makes timestamp_granularities:
|
||||
// ["word"] come back empty however good the timings were.
|
||||
It("attaches the per-word timings only when they were asked for", func() {
|
||||
words := []asrWord{
|
||||
{Text: "a", Start: 0, End: 100},
|
||||
{Text: "b", Start: 100, End: 250},
|
||||
}
|
||||
with := wordsToSegments(words, true)
|
||||
Expect(with[0].Words).To(HaveLen(2))
|
||||
Expect(with[0].Words[1].Text).To(Equal("b"))
|
||||
Expect(time.Duration(with[0].Words[1].Start)).To(Equal(100 * time.Millisecond))
|
||||
Expect(time.Duration(with[0].Words[1].End)).To(Equal(250 * time.Millisecond))
|
||||
|
||||
Expect(wordsToSegments(words, false)[0].Words).To(BeEmpty())
|
||||
})
|
||||
|
||||
// A speaker change splits the run, and each segment must carry only its own
|
||||
// words rather than the whole utterance's.
|
||||
It("gives each speaker run only its own words", func() {
|
||||
segs := wordsToSegments([]asrWord{
|
||||
{Text: "a", Speaker: 1},
|
||||
{Text: "b", Speaker: 2},
|
||||
}, true)
|
||||
Expect(segs).To(HaveLen(2))
|
||||
Expect(segs[0].Words).To(HaveLen(1))
|
||||
Expect(segs[0].Words[0].Text).To(Equal("a"))
|
||||
Expect(segs[1].Words[0].Text).To(Equal("b"))
|
||||
})
|
||||
|
||||
// The C ABI documents the speaker tag as 1-based with 0 meaning "untagged",
|
||||
// so a run of untagged words must not come back attributed to a speaker
|
||||
// literally named "0".
|
||||
It("labels a diarized run and leaves an untagged one unlabelled", func() {
|
||||
Expect(wordsToSegments([]asrWord{{Text: "a", Speaker: 2}}, false)[0].Speaker).To(Equal("2"))
|
||||
Expect(wordsToSegments([]asrWord{{Text: "a", Speaker: 0}}, false)[0].Speaker).To(BeEmpty())
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("wordsRequested", func() {
|
||||
It("recognises the OpenAI word granularity in any casing or padding", func() {
|
||||
Expect(wordsRequested([]string{"word"})).To(BeTrue())
|
||||
Expect(wordsRequested([]string{"segment", " Word "})).To(BeTrue())
|
||||
})
|
||||
|
||||
It("defaults to segment level", func() {
|
||||
Expect(wordsRequested(nil)).To(BeFalse())
|
||||
Expect(wordsRequested([]string{"segment"})).To(BeFalse())
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("recognizeF32", func() {
|
||||
// &pcm[0] panics on a zero-length slice, and a silent or empty upload is
|
||||
// ordinary input rather than an exotic one. The C side rejects empty audio
|
||||
// too, but Go never gets that far.
|
||||
It("refuses empty audio instead of indexing an empty slice", func() {
|
||||
// A zero options struct is enough: the guard has to fire before the
|
||||
// options are ever handed across the ABI, and building real ones would
|
||||
// need the library bound, which this spec deliberately does not.
|
||||
opts := cASRRecognitionOptions{}
|
||||
for _, pcm := range [][]float32{nil, {}} {
|
||||
var (
|
||||
handle uintptr
|
||||
err error
|
||||
)
|
||||
Expect(func() { handle, err = recognizeF32(0, &opts, pcm, 16000) }).ToNot(Panic())
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(status.Code(err)).To(Equal(codes.InvalidArgument))
|
||||
Expect(err.Error()).To(ContainSubstring("empty audio"))
|
||||
Expect(handle).To(BeZero())
|
||||
}
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("AudioTranscription", func() {
|
||||
// The gate has to fire before anything expensive: a model loaded as TTS
|
||||
// cannot transcribe whatever the request says, and reading the audio first
|
||||
// would report a file problem for a configuration one.
|
||||
It("refuses a model loaded as another family, before it reads the audio", func() {
|
||||
n := &NemoSpeech{fam: familyTTS}
|
||||
_, err := n.AudioTranscription(context.Background(), &pb.TranscriptRequest{
|
||||
Dst: filepath.Join(GinkgoT().TempDir(), "does-not-exist.wav"),
|
||||
})
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(status.Code(err)).To(Equal(codes.Unimplemented))
|
||||
Expect(err.Error()).To(ContainSubstring("tts"))
|
||||
})
|
||||
|
||||
It("refuses an unloaded model", func() {
|
||||
n := &NemoSpeech{}
|
||||
_, err := n.AudioTranscription(context.Background(), &pb.TranscriptRequest{Dst: "ignored.wav"})
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(status.Code(err)).To(Equal(codes.Unimplemented))
|
||||
})
|
||||
|
||||
// A lock leaked on a rejection path deadlocks the next request rather than
|
||||
// failing it, which is far harder to diagnose than the failure itself.
|
||||
It("releases the engine lock on every rejection path", func() {
|
||||
n := &NemoSpeech{fam: familyTTS}
|
||||
_, err := n.AudioTranscription(context.Background(), &pb.TranscriptRequest{Dst: "x.wav"})
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(n.engineMu.TryLock()).To(BeTrue())
|
||||
n.engineMu.Unlock()
|
||||
})
|
||||
|
||||
It("reports an audio file it cannot read", func() {
|
||||
n := &NemoSpeech{fam: familyASR}
|
||||
_, err := n.AudioTranscription(context.Background(), &pb.TranscriptRequest{
|
||||
Dst: filepath.Join(GinkgoT().TempDir(), "absent.wav"),
|
||||
})
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(status.Code(err)).To(Equal(codes.InvalidArgument))
|
||||
Expect(n.engineMu.TryLock()).To(BeTrue())
|
||||
n.engineMu.Unlock()
|
||||
})
|
||||
|
||||
It("requires a destination path", func() {
|
||||
n := &NemoSpeech{fam: familyASR}
|
||||
_, err := n.AudioTranscription(context.Background(), &pb.TranscriptRequest{})
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(status.Code(err)).To(Equal(codes.InvalidArgument))
|
||||
})
|
||||
|
||||
// The whole rejection path end to end, on the input that actually reaches
|
||||
// it: a silent or truncated upload decodes to zero samples, and the guard
|
||||
// has to fire between the decode and the ABI. Nothing is loaded here (no
|
||||
// recognizer, and the specs that bind the library may not have run), so
|
||||
// this also pins the ORDER: a guard placed after the options are built
|
||||
// calls a nil-bound entry point and panics rather than failing.
|
||||
It("refuses a decodable clip that carries no samples", func() {
|
||||
path := filepath.Join(GinkgoT().TempDir(), "silence.wav")
|
||||
writeMono16kWAV(path, 0)
|
||||
|
||||
n := &NemoSpeech{fam: familyASR, recognizer: 0}
|
||||
var err error
|
||||
Expect(func() {
|
||||
_, err = n.AudioTranscription(context.Background(), &pb.TranscriptRequest{Dst: path})
|
||||
}).ToNot(Panic())
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(status.Code(err)).To(Equal(codes.InvalidArgument))
|
||||
Expect(err.Error()).To(ContainSubstring("empty audio"))
|
||||
Expect(n.engineMu.TryLock()).To(BeTrue())
|
||||
n.engineMu.Unlock()
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("sampleRateOf", func() {
|
||||
// 0 is not "unknown" to this runtime: nemo_speech_asr_recognize_f32 and
|
||||
// nemo_speech_asr_stream_push_f32 both read a 0 rate as "these samples are
|
||||
// already at the model rate" and skip resampling. Falling back to it for an
|
||||
// undecodable header would silently pitch-shift the audio instead of
|
||||
// failing, so an unknown rate has to be an error.
|
||||
It("rejects a buffer whose format the decoder did not fill in", func() {
|
||||
_, err := sampleRateOf(&audio.IntBuffer{})
|
||||
Expect(err).To(HaveOccurred())
|
||||
})
|
||||
|
||||
It("rejects a non-positive sample rate", func() {
|
||||
_, err := sampleRateOf(&audio.IntBuffer{Format: &audio.Format{SampleRate: 0, NumChannels: 1}})
|
||||
Expect(err).To(HaveOccurred())
|
||||
})
|
||||
|
||||
// The WAV header carries the sample rate as an unsigned 32-bit field, which
|
||||
// go-audio widens to int. Anything above the int32 range therefore passes a
|
||||
// "> 0" test and then narrows to a NEGATIVE rate, which the runtime would take
|
||||
// as a resampling ratio rather than reject. The failure is silent, so the
|
||||
// bound is asserted rather than left to the caller.
|
||||
//
|
||||
// Written as a conversion plus one rather than as the constant MaxInt32+1:
|
||||
// the untyped form does not fit an int on a 32-bit build and would not
|
||||
// compile there, while this wraps to a negative rate the same guard rejects.
|
||||
It("rejects a rate that would not survive the narrowing to int32", func() {
|
||||
_, err := sampleRateOf(&audio.IntBuffer{
|
||||
Format: &audio.Format{SampleRate: int(math.MaxInt32) + 1, NumChannels: 1},
|
||||
})
|
||||
Expect(err).To(HaveOccurred())
|
||||
})
|
||||
|
||||
It("returns the decoded rate", func() {
|
||||
rate, err := sampleRateOf(&audio.IntBuffer{Format: &audio.Format{SampleRate: 22050, NumChannels: 1}})
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(rate).To(Equal(int32(22050)))
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("decodeAudioMono16k", func() {
|
||||
It("decodes a 16 kHz mono WAV to float32 samples at its own rate", func() {
|
||||
path := filepath.Join(GinkgoT().TempDir(), "silence.wav")
|
||||
writeMono16kWAV(path, 800)
|
||||
|
||||
pcm, rate, err := decodeAudioMono16k(path)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(rate).To(Equal(int32(16000)))
|
||||
Expect(pcm).To(HaveLen(800))
|
||||
})
|
||||
|
||||
// A zero-frame WAV is what a truncated upload decodes to, and it is the
|
||||
// input recognizeF32's guard exists for.
|
||||
It("decodes a WAV with no frames to an empty slice", func() {
|
||||
path := filepath.Join(GinkgoT().TempDir(), "empty.wav")
|
||||
writeMono16kWAV(path, 0)
|
||||
|
||||
pcm, _, err := decodeAudioMono16k(path)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(pcm).To(BeEmpty())
|
||||
})
|
||||
|
||||
It("reports a file that does not exist", func() {
|
||||
_, _, err := decodeAudioMono16k(filepath.Join(GinkgoT().TempDir(), "nope.wav"))
|
||||
Expect(err).To(HaveOccurred())
|
||||
})
|
||||
})
|
||||
|
||||
// The six frame counts on the recognizer-attached diarizer are
|
||||
// sentinel-sensitive and invisible to every other check in the tree.
|
||||
// src/asr/c_api.cpp:151-165 applies five of them when they are > 0 but applies
|
||||
// left_context_frames when it is >= 0, so a dropped -1 does not fall back to
|
||||
// the model's own streaming geometry, it pins the left context to zero. The
|
||||
// struct is the right shape either way, so abi_test.go's layout assertions
|
||||
// cannot see it.
|
||||
var _ = Describe("asrDiarConfig", func() {
|
||||
It("keeps the model path it was given", func() {
|
||||
Expect(asrDiarConfig(42).ModelPath).To(Equal(uintptr(42)))
|
||||
})
|
||||
|
||||
// A config sent with the wrong size has every field past it ignored by
|
||||
// HAS_FIELD, and the diarizer attaches with defaults instead of failing.
|
||||
It("declares the size the runtime validates against", func() {
|
||||
Expect(asrDiarConfig(42).Size).To(Equal(unsafe.Sizeof(cASRDiarConfig{})))
|
||||
})
|
||||
|
||||
It("leaves every frame count at the sentinel that means default", func() {
|
||||
cfg := asrDiarConfig(42)
|
||||
Expect(cfg.ChunkFrames).To(Equal(diarGeometryDefault))
|
||||
Expect(cfg.RightContextFrames).To(Equal(diarGeometryDefault))
|
||||
Expect(cfg.LeftContextFrames).To(Equal(diarGeometryDefault))
|
||||
Expect(cfg.FIFOFrames).To(Equal(diarGeometryDefault))
|
||||
Expect(cfg.SpkcacheFrames).To(Equal(diarGeometryDefault))
|
||||
Expect(cfg.UpdatePeriodFrames).To(Equal(diarGeometryDefault))
|
||||
})
|
||||
|
||||
// Stated separately from the field-by-field assertions above: the whole
|
||||
// group is only "unset" to the runtime while the sentinel stays negative,
|
||||
// and zero is a value it would apply to the left context.
|
||||
It("uses a negative sentinel, not zero", func() {
|
||||
Expect(diarGeometryDefault).To(BeNumerically("<", 0))
|
||||
})
|
||||
})
|
||||
@@ -1,86 +0,0 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"math"
|
||||
"os"
|
||||
"path/filepath"
|
||||
|
||||
"github.com/go-audio/audio"
|
||||
"github.com/go-audio/wav"
|
||||
"github.com/mudler/LocalAI/pkg/utils"
|
||||
)
|
||||
|
||||
// decodeAudioMono16k converts an arbitrary audio file to 16 kHz mono PCM and
|
||||
// returns the float32 samples together with the rate they are actually at.
|
||||
//
|
||||
// pkg/utils exposes the ffmpeg normalisation (AudioToWav) but no decode, so
|
||||
// every Go ASR backend pairs it with go-audio itself. This mirrors
|
||||
// backend/go/parakeet-cpp rather than adding a shared helper: the backends
|
||||
// differ in what they need back (parakeet wants a duration, this one wants the
|
||||
// sample rate to hand to the runtime), so a shared signature would be a
|
||||
// lowest-common-denominator of both.
|
||||
func decodeAudioMono16k(path string) ([]float32, int32, error) {
|
||||
dir, err := os.MkdirTemp("", "nemo-speech")
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
defer func() { _ = os.RemoveAll(dir) }()
|
||||
|
||||
// A WAV already at 16 kHz mono 16-bit is hardlinked or copied through
|
||||
// without spawning ffmpeg, so the common case costs nothing.
|
||||
converted := filepath.Join(dir, "converted.wav")
|
||||
if err := utils.AudioToWav(path, converted); err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
|
||||
// #nosec G304 -- converted is filepath.Join of a directory this function just
|
||||
// created with os.MkdirTemp and a constant basename. The request-controlled
|
||||
// path is the INPUT to AudioToWav and never reaches this open.
|
||||
fh, err := os.Open(converted)
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
defer func() { _ = fh.Close() }()
|
||||
|
||||
buf, err := wav.NewDecoder(fh).FullPCMBuffer()
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
|
||||
// The rate is read back from the decoded file rather than assumed to be
|
||||
// 16000. AudioToWav always lands there today, but the runtime resamples
|
||||
// anything from 8 to 96 kHz off this number, so a wrong one would not fail,
|
||||
// it would silently pitch-shift the audio and quietly degrade the transcript.
|
||||
rate, err := sampleRateOf(buf)
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
return buf.AsFloat32Buffer().Data, rate, nil
|
||||
}
|
||||
|
||||
// sampleRateOf reads the decoded rate back off the buffer.
|
||||
//
|
||||
// It is an error rather than a zero fallback because 0 is not "unknown" to this
|
||||
// runtime: nemo_speech_asr_recognize_f32 and nemo_speech_asr_stream_push_f32
|
||||
// both read a 0 rate as "these samples are already at the model rate" and skip
|
||||
// resampling (include/nemo_speech/asr.h). Handing 0 over for a header the
|
||||
// decoder could not read would not fail, it would silently pitch-shift the
|
||||
// audio and quietly degrade the transcript, which is the same failure the
|
||||
// caller comment warns about for a wrong rate.
|
||||
//
|
||||
// The upper bound is what makes the narrowing to int32 safe rather than merely
|
||||
// unlikely. go-audio reads the WAV header's sample rate as an unsigned 32-bit
|
||||
// field into an int, so on a 64-bit build a header claiming more than 2^31-1
|
||||
// survives the "> 0" test and then narrows to a NEGATIVE rate, which the runtime
|
||||
// would take as a resampling ratio. Nothing this backend decodes can reach that
|
||||
// today (AudioToWav either passes through a WAV it has confirmed is exactly
|
||||
// 16 kHz or runs ffmpeg with -ar 16000), but that is a property of a helper in
|
||||
// another package, and this function exists precisely because the rate is read
|
||||
// back rather than assumed.
|
||||
func sampleRateOf(buf *audio.IntBuffer) (int32, error) {
|
||||
if buf.Format == nil || buf.Format.SampleRate <= 0 || buf.Format.SampleRate > math.MaxInt32 {
|
||||
return 0, errors.New("nemo-speech-cpp: decoded audio has no usable sample rate")
|
||||
}
|
||||
return int32(buf.Format.SampleRate), nil
|
||||
}
|
||||
@@ -1,502 +0,0 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"strconv"
|
||||
"unsafe"
|
||||
|
||||
pb "github.com/mudler/LocalAI/pkg/grpc/proto"
|
||||
"github.com/mudler/xlog"
|
||||
"google.golang.org/grpc/codes"
|
||||
"google.golang.org/grpc/status"
|
||||
)
|
||||
|
||||
// diarSegmentsMaxAttempts bounds the count-then-fill retry.
|
||||
//
|
||||
// On a finished stream the count is stable and one attempt is always enough.
|
||||
// The bound exists because the RPC holds engineMu for its whole body, so a
|
||||
// runtime whose count kept growing would not merely spin, it would block the
|
||||
// unload behind it.
|
||||
const diarSegmentsMaxAttempts = 4
|
||||
|
||||
// maxDiarSegments caps the buffer collectSegments will allocate from a count
|
||||
// the C side reported.
|
||||
//
|
||||
// make() panics rather than erroring on a length it cannot satisfy, and a
|
||||
// panic in an RPC handler takes the backend process down, so an uninitialised
|
||||
// or corrupted size_t coming back across the ABI would kill the model rather
|
||||
// than fail the request. The ceiling turns that into a diagnosable error.
|
||||
//
|
||||
// It is set far above anything real: a segment spans at least one 80 ms frame,
|
||||
// so 2^22 segments is upwards of 93 hours of audio, and the buffer itself
|
||||
// would already be 100 MB at 24 bytes each.
|
||||
const maxDiarSegments = 1 << 22
|
||||
|
||||
// diarSegmenter is the result half of the diarization C API: the two-call
|
||||
// protocol nemo_speech_diar_segments documents.
|
||||
//
|
||||
// The two calls are the same C function with a different `out`, but they are
|
||||
// separate methods here because their contracts differ. countSegments passes
|
||||
// out=NULL, which the runtime answers by writing *count and returning OK
|
||||
// without touching a buffer. fillSegments passes a real buffer and gets
|
||||
// INVALID_ARGUMENT if it is too short, having written *count first, which is
|
||||
// what makes a growth retry possible at all.
|
||||
type diarSegmenter interface {
|
||||
// countSegments is the size query. It never fails for lack of a buffer.
|
||||
countSegments() (uint64, error)
|
||||
// fillSegments fills buf and returns the count the runtime reported. That
|
||||
// count is meaningful even alongside an error: on a short buffer the
|
||||
// runtime writes it before rejecting the call.
|
||||
fillSegments(buf []cDiarSegment) (uint64, error)
|
||||
}
|
||||
|
||||
// diarStream is one diarization job over the C API, narrowed to what the RPC
|
||||
// uses.
|
||||
//
|
||||
// It is an interface for the same reason asrSession is: no Sortformer GGUF is
|
||||
// small enough to keep in the tree, so without a seam at the ABI the loop on
|
||||
// top of it (the empty guard, chunking, finish-before-query, the growth retry)
|
||||
// would have no test at all. A fake here scripts what C returns; it does not
|
||||
// pretend to diarize anything.
|
||||
type diarStream interface {
|
||||
diarSegmenter
|
||||
push(pcm []float32, sampleRate int32) error
|
||||
finish() error
|
||||
close()
|
||||
}
|
||||
|
||||
// diarStreamOpener creates a job. n.openDiarStream is the C-backed one.
|
||||
//
|
||||
// The segmentation config is handed over at open time rather than per query
|
||||
// because it belongs to the whole job: every segments call on one stream must
|
||||
// use the same postprocessing or the segment ids would not be comparable
|
||||
// between calls.
|
||||
type diarStreamOpener func(cfg *cDiarSegmentationConfig) (diarStream, error)
|
||||
|
||||
// cDiarStream is the real diarStream, over one nemo_speech_diar_stream.
|
||||
type cDiarStream struct {
|
||||
handle uintptr
|
||||
cfg *cDiarSegmentationConfig
|
||||
}
|
||||
|
||||
// cfgPtr hands the segmentation config to C, or NULL when the request asked
|
||||
// for no postprocessing. NULL is not the same as a zeroed struct in spirit
|
||||
// even though src/asr/c_api.cpp treats them alike today: diar.h documents NULL
|
||||
// as "library defaults", so it is the one form that cannot be invalidated by a
|
||||
// future field whose sentinel is not zero.
|
||||
func (s *cDiarStream) cfgPtr() unsafe.Pointer {
|
||||
if s.cfg == nil {
|
||||
return nil
|
||||
}
|
||||
// #nosec G103 -- a plain *T to unsafe.Pointer conversion of a non-nil,
|
||||
// GC-traced field. cDiarSegmentationConfig is pure scalars (no uintptr
|
||||
// members to pin) and the stream owns it for its whole life, so the only
|
||||
// requirement is that it outlive the DiarSegments call, which it does.
|
||||
return unsafe.Pointer(s.cfg)
|
||||
}
|
||||
|
||||
func (s *cDiarStream) push(pcm []float32, sampleRate int32) error {
|
||||
// &pcm[0] panics on an empty slice before the C side ever sees the call.
|
||||
if len(pcm) == 0 {
|
||||
return nil
|
||||
}
|
||||
if st := DiarStreamPushF32(s.handle, &pcm[0], uint64(len(pcm)), sampleRate); st != 0 {
|
||||
return statusErrorf(st, "nemo-speech-cpp: diarization push: %s", ASRLastError())
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *cDiarStream) finish() error {
|
||||
if st := DiarStreamFinish(s.handle); st != 0 {
|
||||
return statusErrorf(st, "nemo-speech-cpp: diarization finish: %s", ASRLastError())
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *cDiarStream) close() { DiarStreamClose(s.handle) }
|
||||
|
||||
func (s *cDiarStream) countSegments() (uint64, error) {
|
||||
var count uint64
|
||||
// out=NULL and capacity=0: the size query. The runtime reads capacity only
|
||||
// once it has a buffer to check it against.
|
||||
if st := DiarSegments(s.handle, s.cfgPtr(), nil, 0, &count); st != 0 {
|
||||
return 0, statusErrorf(st,
|
||||
"nemo-speech-cpp: diarization segment count: %s", ASRLastError())
|
||||
}
|
||||
return count, nil
|
||||
}
|
||||
|
||||
func (s *cDiarStream) fillSegments(buf []cDiarSegment) (uint64, error) {
|
||||
if len(buf) == 0 {
|
||||
// A NULL out would silently turn this into a second size query, and the
|
||||
// caller would read it as "filled nothing" rather than "asked nothing".
|
||||
return 0, status.Error(codes.Internal,
|
||||
"nemo-speech-cpp: diarization segment fill needs a buffer")
|
||||
}
|
||||
var count uint64
|
||||
// #nosec G103 -- &buf[0] is guarded by the empty check above, and the
|
||||
// capacity handed over is exactly len(buf), so the runtime cannot write past
|
||||
// the caller's allocation. collectSegments sizes buf under maxDiarSegments
|
||||
// and rejects a reported count larger than it rather than slicing to it.
|
||||
st := DiarSegments(s.handle, s.cfgPtr(), unsafe.Pointer(&buf[0]), uint64(len(buf)), &count)
|
||||
if st != 0 {
|
||||
// count is returned alongside the error on purpose: a too-small buffer
|
||||
// is rejected only after the runtime has written the size it wanted.
|
||||
return count, statusErrorf(st,
|
||||
"nemo-speech-cpp: diarization segments: %s", ASRLastError())
|
||||
}
|
||||
return count, nil
|
||||
}
|
||||
|
||||
// diarGeometryDefault is the sentinel that means "keep the preset's value" for
|
||||
// every one of nemo_speech_diar_model_config's six frame counts.
|
||||
//
|
||||
// It has to be negative, not zero, and that is not a style choice.
|
||||
// src/asr/c_api.cpp:497-512 applies five of the six overrides when they are
|
||||
// > 0 but applies left_context_frames when it is >= 0, so a zero-valued config
|
||||
// reads as "unset" for five fields and as an explicit left context of zero for
|
||||
// the sixth. That silently changes the model's streaming geometry, and no
|
||||
// layout assertion can see it because the struct is the right shape either way.
|
||||
const diarGeometryDefault int32 = -1
|
||||
|
||||
// diarModelConfig builds the create-time config for the standalone diarizer.
|
||||
//
|
||||
// Extracted from loadDiarizer purely so the sentinels above can be asserted:
|
||||
// they are invisible to every other check in the tree, including the layout
|
||||
// assertions, so a spec pinning them is the only thing standing between a
|
||||
// dropped -1 and a quietly mis-configured model.
|
||||
//
|
||||
// modelPath is a C pointer from cstr, not a Go string, and the caller owns its
|
||||
// release. preset is deliberately left NULL, which diar.h reads as "streaming".
|
||||
// The "offline" preset is a different accuracy/latency tradeoff for long files
|
||||
// and is worth exposing, but not on an unverified guess: no Sortformer GGUF
|
||||
// exists here to measure the difference on.
|
||||
func diarModelConfig(modelPath uintptr, gpu int32) cDiarModelConfig {
|
||||
return cDiarModelConfig{
|
||||
Size: unsafe.Sizeof(cDiarModelConfig{}),
|
||||
ModelPath: modelPath,
|
||||
GPU: gpu,
|
||||
ChunkFrames: diarGeometryDefault,
|
||||
RightContextFrames: diarGeometryDefault,
|
||||
LeftContextFrames: diarGeometryDefault,
|
||||
FIFOFrames: diarGeometryDefault,
|
||||
SpkcacheFrames: diarGeometryDefault,
|
||||
UpdatePeriodFrames: diarGeometryDefault,
|
||||
}
|
||||
}
|
||||
|
||||
// loadDiarizer creates the standalone Sortformer diarizer.
|
||||
//
|
||||
// This must not take engineMu: Load is its only caller and already holds it.
|
||||
func (n *NemoSpeech) loadDiarizer(modelFile string) error {
|
||||
pathP, freePath := cstr(modelFile)
|
||||
defer freePath()
|
||||
|
||||
cfg := diarModelConfig(pathP, n.opts.gpu)
|
||||
|
||||
xlog.Info("nemo-speech-cpp: creating diarizer", "gpu", n.opts.gpu)
|
||||
|
||||
// #nosec G103 -- cfg is a local POD struct borrowed for this call only. Its
|
||||
// only uintptr member is ModelPath, the cstr allocation pinned by the
|
||||
// deferred freePath above (Preset is deliberately NULL), and
|
||||
// nemo_speech_diar_create deep-copies the path and retains nothing.
|
||||
if st := DiarCreate(unsafe.Pointer(&cfg), &n.diarizer); st != 0 {
|
||||
return statusErrorf(st, "nemo-speech-cpp: diarizer create: %s", ASRLastError())
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// openDiarStream starts a diarization job on the loaded model.
|
||||
//
|
||||
// The caller must hold engineMu.
|
||||
func (n *NemoSpeech) openDiarStream(cfg *cDiarSegmentationConfig) (diarStream, error) {
|
||||
var handle uintptr
|
||||
if st := DiarStreamOpen(n.diarizer, &handle); st != 0 {
|
||||
return nil, statusErrorf(st,
|
||||
"nemo-speech-cpp: diarization stream open: %s", ASRLastError())
|
||||
}
|
||||
return &cDiarStream{handle: handle, cfg: cfg}, nil
|
||||
}
|
||||
|
||||
// sizeofDiarSegmentationConfig is the size the runtime validates the config
|
||||
// against. It is a function so the specs can assert the value the config
|
||||
// actually carries rather than restate the number.
|
||||
func sizeofDiarSegmentationConfig() uintptr {
|
||||
return unsafe.Sizeof(cDiarSegmentationConfig{})
|
||||
}
|
||||
|
||||
// segmentationConfig maps the request's postprocessing knobs onto
|
||||
// nemo_speech_diar_segmentation_config, or returns nil when none were set.
|
||||
//
|
||||
// Only two of DiarizeRequest's tuning fields have a real equivalent here, and
|
||||
// both are exact rather than approximate: NeMo's ts_vad postprocessing is the
|
||||
// same algorithm the proto's wording describes.
|
||||
//
|
||||
// - min_duration_on ("discard segments shorter than this") is min_duration_sec
|
||||
// ("drop segments shorter than this"), which c_api.cpp assigns to
|
||||
// DiarSegmentationCfg.min_duration_on.
|
||||
// - min_duration_off ("merge gaps shorter than this") is min_gap_sec ("fill
|
||||
// silence gaps shorter than this"), assigned to min_duration_off.
|
||||
//
|
||||
// The names cross over between the proto and the C header, which is exactly the
|
||||
// kind of transposition a layout assertion cannot see, so each mapping is
|
||||
// pinned by its own spec.
|
||||
//
|
||||
// Nothing is written for a non-positive value: the runtime tests every field
|
||||
// with > 0 and keeps its default otherwise, so a zero here means "unset" on
|
||||
// both sides.
|
||||
func segmentationConfig(req *pb.DiarizeRequest) *cDiarSegmentationConfig {
|
||||
cfg := cDiarSegmentationConfig{Size: sizeofDiarSegmentationConfig()}
|
||||
|
||||
var set bool
|
||||
if v := req.GetMinDurationOn(); v > 0 {
|
||||
cfg.MinDurationSec = float64(v)
|
||||
set = true
|
||||
}
|
||||
if v := req.GetMinDurationOff(); v > 0 {
|
||||
cfg.MinGapSec = float64(v)
|
||||
set = true
|
||||
}
|
||||
if !set {
|
||||
return nil
|
||||
}
|
||||
return &cfg
|
||||
}
|
||||
|
||||
// unsupportedRequestFields names the DiarizeRequest fields this backend cannot
|
||||
// honour, so they are logged rather than silently dropped.
|
||||
//
|
||||
// Each is a deliberate omission, not a gap waiting to be filled:
|
||||
//
|
||||
// - num_speakers, min_speakers, max_speakers: Sortformer is end-to-end and
|
||||
// its speaker capacity is fixed by the checkpoint (v2: 4).
|
||||
// nemo_speech_diar_num_speakers reports that capacity, it does not set it,
|
||||
// and there is no config field for a target count.
|
||||
// - clustering_threshold: there is no clustering stage. The nearest knob is
|
||||
// the onset/offset probability hysteresis, which is a different quantity on
|
||||
// a different scale, so mapping one onto the other would invent an
|
||||
// equivalence the header does not have.
|
||||
// - include_text: this pipeline carries no ASR at all (diar.h: "no ASR
|
||||
// involved"). Word-level speaker tags on a transcript are the ASR surface's
|
||||
// job, through diar_model plus enable_speaker_diarization.
|
||||
// - threads: neither nemo_speech_diar_model_config nor the segmentation
|
||||
// config has a thread count.
|
||||
func unsupportedRequestFields(req *pb.DiarizeRequest) []string {
|
||||
var out []string
|
||||
if req.GetNumSpeakers() != 0 {
|
||||
out = append(out, "num_speakers")
|
||||
}
|
||||
if req.GetMinSpeakers() != 0 {
|
||||
out = append(out, "min_speakers")
|
||||
}
|
||||
if req.GetMaxSpeakers() != 0 {
|
||||
out = append(out, "max_speakers")
|
||||
}
|
||||
if req.GetClusteringThreshold() != 0 {
|
||||
out = append(out, "clustering_threshold")
|
||||
}
|
||||
if req.GetIncludeText() {
|
||||
out = append(out, "include_text")
|
||||
}
|
||||
if req.GetThreads() != 0 {
|
||||
out = append(out, "threads")
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// collectSegments runs the count-then-fill protocol and returns the segments.
|
||||
//
|
||||
// The growth retry is not defensive padding. nemo_speech_diar_segments writes
|
||||
// *count and only then rejects a buffer that is too small, so the size a
|
||||
// rejected call reports is the size to retry with; without the retry a stream
|
||||
// that gained a segment between the two calls would fail the whole request.
|
||||
// Truncating to the first count instead would be worse still, dropping turns
|
||||
// with nothing to show for it.
|
||||
func collectSegments(s diarSegmenter) ([]cDiarSegment, error) {
|
||||
want, err := s.countSegments()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
for range diarSegmentsMaxAttempts {
|
||||
if want == 0 {
|
||||
// No segments means no fill: the fill call needs a non-empty buffer
|
||||
// to be distinguishable from a second size query.
|
||||
return nil, nil
|
||||
}
|
||||
if want > maxDiarSegments {
|
||||
return nil, status.Errorf(codes.Internal,
|
||||
"nemo-speech-cpp: diarization reported %d segments, above the %d ceiling", want, maxDiarSegments)
|
||||
}
|
||||
|
||||
buf := make([]cDiarSegment, want)
|
||||
got, fillErr := s.fillSegments(buf)
|
||||
if fillErr == nil {
|
||||
if got > want {
|
||||
// The runtime cannot report this on success (it rejects a short
|
||||
// buffer instead), so it means the ABI is not what this code
|
||||
// thinks it is. Slicing to it would read past the allocation.
|
||||
return nil, status.Errorf(codes.Internal,
|
||||
"nemo-speech-cpp: diarization returned %d segments for a %d-segment buffer", got, want)
|
||||
}
|
||||
return buf[:got], nil
|
||||
}
|
||||
// A count that did not grow means the call failed for some other
|
||||
// reason, and retrying the same size would just fail the same way.
|
||||
if got <= want {
|
||||
return nil, fillErr
|
||||
}
|
||||
want = got
|
||||
}
|
||||
|
||||
return nil, status.Error(codes.Internal,
|
||||
"nemo-speech-cpp: diarization segment count kept growing, giving up")
|
||||
}
|
||||
|
||||
// toDiarizeSegments converts the runtime's segments to the wire form.
|
||||
//
|
||||
// No unit conversion happens here, and that is the point: nemo_speech_diar_segment
|
||||
// carries start_time and end_time in SECONDS already (diar.h), and
|
||||
// DiarizeSegment.start/end are seconds too. The frame indices the model works
|
||||
// in never reach this layer, so nemo_speech_diar_seconds_per_frame is not
|
||||
// involved. The narrowing to float32 is the proto's choice of type; at 80 ms
|
||||
// resolution it is lossless for any clip short enough to hold in memory.
|
||||
//
|
||||
// The speaker label is the runtime's 1-based tag rendered as a decimal string,
|
||||
// which is what wordsToSegments emits for the ASR path. The same speaker has to
|
||||
// read the same way whether the caller diarized a file or transcribed it.
|
||||
func toDiarizeSegments(in []cDiarSegment) []*pb.DiarizeSegment {
|
||||
if len(in) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make([]*pb.DiarizeSegment, 0, len(in))
|
||||
for i, s := range in {
|
||||
out = append(out, &pb.DiarizeSegment{
|
||||
Id: int32(i),
|
||||
Start: float32(s.StartTime),
|
||||
End: float32(s.EndTime),
|
||||
Speaker: strconv.Itoa(int(s.Speaker)),
|
||||
})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// distinctSpeakers counts the speaker labels present in the segments.
|
||||
//
|
||||
// This is what DiarizeResponse.num_speakers is documented to hold, and it is
|
||||
// NOT nemo_speech_diar_num_speakers: that reports the checkpoint's capacity
|
||||
// (four for Sortformer v2), so a two-person interview would come back claiming
|
||||
// four speakers.
|
||||
func distinctSpeakers(segs []*pb.DiarizeSegment) int32 {
|
||||
seen := make(map[string]struct{}, len(segs))
|
||||
for _, s := range segs {
|
||||
seen[s.GetSpeaker()] = struct{}{}
|
||||
}
|
||||
// #nosec G115 -- seen holds at most one entry per segment, and collectSegments
|
||||
// refuses any count above maxDiarSegments (2^22), so this is orders of
|
||||
// magnitude below the int32 the proto field is.
|
||||
return int32(len(seen))
|
||||
}
|
||||
|
||||
// diarizePCM drives one whole clip through a diarization job.
|
||||
//
|
||||
// The caller must hold engineMu.
|
||||
func diarizePCM(open diarStreamOpener, pcm []float32, sampleRate int32, cfg *cDiarSegmentationConfig) (*pb.DiarizeResponse, error) {
|
||||
// Before the stream is opened, not inside the push: a silent or truncated
|
||||
// upload decodes to zero samples, &pcm[0] panics on that, and there is no
|
||||
// diarization to be had from it anyway.
|
||||
if len(pcm) == 0 {
|
||||
return nil, status.Error(codes.InvalidArgument, "nemo-speech-cpp: empty audio")
|
||||
}
|
||||
|
||||
stream, err := open(cfg)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer stream.close()
|
||||
|
||||
// Chunked rather than pushed whole so the runtime advances as it goes
|
||||
// instead of buffering the entire clip before the first chunk boundary.
|
||||
for _, chunk := range chunkPCM(pcm, streamChunkSamples) {
|
||||
if err := stream.push(chunk, sampleRate); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
// Before the query, always: finish is what labels the audio tail, so
|
||||
// segmenting first drops the last turn of every clip.
|
||||
if err := stream.finish(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
raw, err := collectSegments(stream)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
segs := toDiarizeSegments(raw)
|
||||
out := &pb.DiarizeResponse{
|
||||
Segments: segs,
|
||||
NumSpeakers: distinctSpeakers(segs),
|
||||
}
|
||||
// 0 is the proto's "unknown" and the C API's "already at the model rate",
|
||||
// so a rate that means the latter must not be divided by.
|
||||
if sampleRate > 0 {
|
||||
out.Duration = float32(len(pcm)) / float32(sampleRate)
|
||||
}
|
||||
// Language and the per-segment text stay empty: there is no ASR in this
|
||||
// pipeline to fill them, and the proto documents both as optional.
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// Diarize labels who spoke when in the audio at req.Dst.
|
||||
//
|
||||
// The whole body runs inside withEngine, so the family check and the C calls
|
||||
// that trust the handle happen under a single acquisition of engineMu. The
|
||||
// audio decode is in there too, for the reason documented on
|
||||
// AudioTranscription: the backend already serialises RPCs, so the wider hold
|
||||
// costs nothing, and the narrower one is the gap Free can land in.
|
||||
func (n *NemoSpeech) Diarize(req *pb.DiarizeRequest) (pb.DiarizeResponse, error) {
|
||||
var out *pb.DiarizeResponse
|
||||
if err := n.withEngine(familyDiarization, func() error {
|
||||
r, err := n.diarize(req)
|
||||
out = r
|
||||
return err
|
||||
}); err != nil {
|
||||
return pb.DiarizeResponse{}, err
|
||||
}
|
||||
if out == nil {
|
||||
return pb.DiarizeResponse{}, status.Error(codes.Internal,
|
||||
"nemo-speech-cpp: diarization produced no result")
|
||||
}
|
||||
|
||||
// Assembled field by field rather than dereferenced: the RPC returns the
|
||||
// message by value and the message embeds a mutex, so copying the struct is
|
||||
// a copylocks violation.
|
||||
return pb.DiarizeResponse{
|
||||
Segments: out.Segments,
|
||||
NumSpeakers: out.NumSpeakers,
|
||||
Duration: out.Duration,
|
||||
Language: out.Language,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// diarize is Diarize's body. The caller must hold engineMu.
|
||||
func (n *NemoSpeech) diarize(req *pb.DiarizeRequest) (*pb.DiarizeResponse, error) {
|
||||
if req.GetDst() == "" {
|
||||
return nil, status.Error(codes.InvalidArgument,
|
||||
"nemo-speech-cpp: DiarizeRequest.dst (audio path) is required")
|
||||
}
|
||||
// Logged rather than rejected: a client that asks for a speaker count still
|
||||
// wants the diarization it can have, and a request that names a field this
|
||||
// backend drops should say so somewhere the operator can find it.
|
||||
if dropped := unsupportedRequestFields(req); len(dropped) > 0 {
|
||||
xlog.Warn("nemo-speech-cpp: ignoring diarization request fields this model has no equivalent for",
|
||||
"fields", dropped)
|
||||
}
|
||||
|
||||
pcm, sampleRate, err := decodeAudioMono16k(req.GetDst())
|
||||
if err != nil {
|
||||
return nil, status.Errorf(codes.InvalidArgument, "nemo-speech-cpp: read audio: %v", err)
|
||||
}
|
||||
|
||||
return diarizePCM(n.openDiarStream, pcm, sampleRate, segmentationConfig(req))
|
||||
}
|
||||
@@ -1,540 +0,0 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"path/filepath"
|
||||
"unsafe"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
"google.golang.org/grpc/codes"
|
||||
"google.golang.org/grpc/status"
|
||||
|
||||
pb "github.com/mudler/LocalAI/pkg/grpc/proto"
|
||||
)
|
||||
|
||||
// fakeDiarStream scripts what the C API returns for one diarization job.
|
||||
//
|
||||
// There is no Sortformer GGUF in the tree, so this is the only way the loop on
|
||||
// top of the ABI (the empty guard, chunking, the count-then-fill protocol, the
|
||||
// buffer growth retry) gets tested at all. It fakes the C contract, not the
|
||||
// model: segs is whatever nemo_speech_diar_segments would have produced.
|
||||
type fakeDiarStream struct {
|
||||
segs []cDiarSegment
|
||||
|
||||
// countErr and fillErrs script failures. fillErrs is consumed one entry per
|
||||
// fillSegments call so a growth retry can be scripted.
|
||||
countErr error
|
||||
fillErrs []error
|
||||
// queryCount, when non-zero, is what the size query reports instead of
|
||||
// len(segs), so a runtime that under-reported can be scripted.
|
||||
queryCount uint64
|
||||
// growTo, when non-zero, is the count reported by the FIRST fillSegments
|
||||
// call, standing in for a runtime whose segment list outgrew the size query.
|
||||
growTo uint64
|
||||
|
||||
pushed [][]float32
|
||||
rates []int32
|
||||
finished int
|
||||
closed int
|
||||
counts int
|
||||
fills int
|
||||
// opened records the segmentation config the opener was handed.
|
||||
cfg *cDiarSegmentationConfig
|
||||
}
|
||||
|
||||
func (f *fakeDiarStream) push(pcm []float32, rate int32) error {
|
||||
f.pushed = append(f.pushed, pcm)
|
||||
f.rates = append(f.rates, rate)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (f *fakeDiarStream) finish() error {
|
||||
f.finished++
|
||||
return nil
|
||||
}
|
||||
|
||||
func (f *fakeDiarStream) close() { f.closed++ }
|
||||
|
||||
func (f *fakeDiarStream) countSegments() (uint64, error) {
|
||||
f.counts++
|
||||
if f.countErr != nil {
|
||||
return 0, f.countErr
|
||||
}
|
||||
if f.queryCount > 0 {
|
||||
return f.queryCount, nil
|
||||
}
|
||||
return uint64(len(f.segs)), nil
|
||||
}
|
||||
|
||||
func (f *fakeDiarStream) fillSegments(buf []cDiarSegment) (uint64, error) {
|
||||
f.fills++
|
||||
var err error
|
||||
if len(f.fillErrs) > 0 {
|
||||
err, f.fillErrs = f.fillErrs[0], f.fillErrs[1:]
|
||||
}
|
||||
if f.fills == 1 && f.growTo > 0 {
|
||||
// The runtime writes *count before it rejects a short buffer, so a
|
||||
// growth failure still reports the count the caller needs.
|
||||
return f.growTo, err
|
||||
}
|
||||
n := copy(buf, f.segs)
|
||||
return uint64(n), err
|
||||
}
|
||||
|
||||
// allPushed flattens what the fake received, so a spec can assert the audio
|
||||
// arrived intact regardless of how it was chunked.
|
||||
func (f *fakeDiarStream) allPushed() []float32 {
|
||||
var out []float32
|
||||
for _, c := range f.pushed {
|
||||
out = append(out, c...)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func (f *fakeDiarStream) opener() diarStreamOpener {
|
||||
return func(cfg *cDiarSegmentationConfig) (diarStream, error) {
|
||||
f.cfg = cfg
|
||||
return f, nil
|
||||
}
|
||||
}
|
||||
|
||||
var _ = Describe("Diarize", func() {
|
||||
It("refuses when the loaded model is not a diarization model", func() {
|
||||
n := &NemoSpeech{fam: familyASR}
|
||||
_, err := n.Diarize(&pb.DiarizeRequest{Dst: "/tmp/whatever.wav"})
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(status.Code(err)).To(Equal(codes.Unimplemented))
|
||||
})
|
||||
|
||||
It("refuses on a model that was never loaded", func() {
|
||||
n := &NemoSpeech{}
|
||||
_, err := n.Diarize(&pb.DiarizeRequest{Dst: "/tmp/whatever.wav"})
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(status.Code(err)).To(Equal(codes.Unimplemented))
|
||||
})
|
||||
|
||||
// The family gate has to run before anything reads the request, or a
|
||||
// misrouted request would be reported as a bad path rather than as a model
|
||||
// that cannot diarize.
|
||||
It("reports a missing audio path on a diarization model", func() {
|
||||
n := &NemoSpeech{fam: familyDiarization}
|
||||
_, err := n.Diarize(&pb.DiarizeRequest{})
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(status.Code(err)).To(Equal(codes.InvalidArgument))
|
||||
Expect(err.Error()).To(ContainSubstring("dst"))
|
||||
})
|
||||
|
||||
It("reports audio it cannot read", func() {
|
||||
n := &NemoSpeech{fam: familyDiarization}
|
||||
missing := filepath.Join(GinkgoT().TempDir(), "absent.wav")
|
||||
_, err := n.Diarize(&pb.DiarizeRequest{Dst: missing})
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(status.Code(err)).To(Equal(codes.InvalidArgument))
|
||||
Expect(err.Error()).To(ContainSubstring("read audio"))
|
||||
})
|
||||
|
||||
// A rejection must leave the mutex free, or the next request deadlocks
|
||||
// rather than fails.
|
||||
It("releases the engine lock on every rejection path", func() {
|
||||
n := &NemoSpeech{fam: familyDiarization}
|
||||
_, err := n.Diarize(&pb.DiarizeRequest{})
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(n.engineMu.TryLock()).To(BeTrue())
|
||||
n.engineMu.Unlock()
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("diarizePCM", func() {
|
||||
// Task 7 found that a purego-bound entry point reached with zero samples
|
||||
// panics on &pcm[0], so a silent clip must be rejected before the stream is
|
||||
// ever opened, not inside the push.
|
||||
It("rejects empty audio without opening a stream", func() {
|
||||
opened := false
|
||||
open := func(*cDiarSegmentationConfig) (diarStream, error) {
|
||||
opened = true
|
||||
return &fakeDiarStream{}, nil
|
||||
}
|
||||
_, err := diarizePCM(open, nil, 16000, nil)
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(status.Code(err)).To(Equal(codes.InvalidArgument))
|
||||
Expect(err.Error()).To(ContainSubstring("empty audio"))
|
||||
Expect(opened).To(BeFalse())
|
||||
})
|
||||
|
||||
It("pushes the whole clip, finishes, and closes the stream", func() {
|
||||
pcm := make([]float32, streamChunkSamples*2+7)
|
||||
for i := range pcm {
|
||||
pcm[i] = float32(i)
|
||||
}
|
||||
f := &fakeDiarStream{}
|
||||
_, err := diarizePCM(f.opener(), pcm, 16000, nil)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
Expect(f.allPushed()).To(Equal(pcm))
|
||||
Expect(f.pushed).To(HaveLen(3), "the clip must be chunked, not pushed whole")
|
||||
Expect(f.rates).To(HaveEach(int32(16000)))
|
||||
Expect(f.finished).To(Equal(1))
|
||||
Expect(f.closed).To(Equal(1))
|
||||
})
|
||||
|
||||
// Segments must come from a finished stream: the tail of the audio is only
|
||||
// labelled by finish, so asking first silently drops the last turn.
|
||||
It("finishes the stream before it asks for segments", func() {
|
||||
f := &fakeDiarStream{segs: []cDiarSegment{{StartTime: 0, EndTime: 1, Speaker: 1}}}
|
||||
f.fillErrs = nil
|
||||
var finishedAtCount int
|
||||
wrapped := func(cfg *cDiarSegmentationConfig) (diarStream, error) {
|
||||
f.cfg = cfg
|
||||
return &countObserver{fakeDiarStream: f, seen: &finishedAtCount}, nil
|
||||
}
|
||||
_, err := diarizePCM(wrapped, []float32{1, 2, 3}, 16000, nil)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(finishedAtCount).To(Equal(1), "the size query ran before finish")
|
||||
})
|
||||
|
||||
It("converts the runtime's seconds straight through and numbers the segments", func() {
|
||||
f := &fakeDiarStream{segs: []cDiarSegment{
|
||||
{StartTime: 0, EndTime: 0.8, Speaker: 1},
|
||||
{StartTime: 0.8, EndTime: 2.0, Speaker: 2},
|
||||
}}
|
||||
res, err := diarizePCM(f.opener(), []float32{1, 2, 3}, 16000, nil)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(res.Segments).To(HaveLen(2))
|
||||
|
||||
Expect(res.Segments[0].GetId()).To(Equal(int32(0)))
|
||||
Expect(res.Segments[0].GetStart()).To(BeNumerically("~", 0.0, 1e-6))
|
||||
Expect(res.Segments[0].GetEnd()).To(BeNumerically("~", 0.8, 1e-6))
|
||||
Expect(res.Segments[0].GetSpeaker()).To(Equal("1"))
|
||||
|
||||
Expect(res.Segments[1].GetId()).To(Equal(int32(1)))
|
||||
Expect(res.Segments[1].GetStart()).To(BeNumerically("~", 0.8, 1e-6))
|
||||
Expect(res.Segments[1].GetEnd()).To(BeNumerically("~", 2.0, 1e-6))
|
||||
Expect(res.Segments[1].GetSpeaker()).To(Equal("2"))
|
||||
})
|
||||
|
||||
It("reports the clip duration in seconds", func() {
|
||||
f := &fakeDiarStream{}
|
||||
res, err := diarizePCM(f.opener(), make([]float32, 32000), 16000, nil)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(res.GetDuration()).To(BeNumerically("~", 2.0, 1e-6))
|
||||
})
|
||||
|
||||
// 0 is the proto's documented "unknown", and it is also the C API's "these
|
||||
// samples are already at the model rate", so a rate that cannot be trusted
|
||||
// must not be turned into a duration.
|
||||
It("reports no duration when the sample rate is unknown", func() {
|
||||
f := &fakeDiarStream{}
|
||||
res, err := diarizePCM(f.opener(), make([]float32, 32000), 0, nil)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(res.GetDuration()).To(BeZero())
|
||||
})
|
||||
|
||||
It("hands the segmentation config to the opener", func() {
|
||||
f := &fakeDiarStream{}
|
||||
cfg := &cDiarSegmentationConfig{MinDurationSec: 0.5}
|
||||
_, err := diarizePCM(f.opener(), []float32{1}, 16000, cfg)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(f.cfg).To(BeIdenticalTo(cfg))
|
||||
})
|
||||
|
||||
It("closes the stream when the segment query fails", func() {
|
||||
f := &fakeDiarStream{countErr: errors.New("boom")}
|
||||
_, err := diarizePCM(f.opener(), []float32{1}, 16000, nil)
|
||||
Expect(err).To(MatchError(ContainSubstring("boom")))
|
||||
Expect(f.closed).To(Equal(1))
|
||||
})
|
||||
|
||||
// The pipeline carries no ASR, so text and language stay empty whatever the
|
||||
// caller asked for.
|
||||
It("leaves the transcript fields empty", func() {
|
||||
f := &fakeDiarStream{segs: []cDiarSegment{{StartTime: 0, EndTime: 1, Speaker: 1}}}
|
||||
res, err := diarizePCM(f.opener(), []float32{1}, 16000, nil)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(res.GetLanguage()).To(BeEmpty())
|
||||
Expect(res.Segments[0].GetText()).To(BeEmpty())
|
||||
})
|
||||
})
|
||||
|
||||
// countObserver records how many size queries had run by the time finish was
|
||||
// called, so the ordering can be asserted without reaching into diarizePCM.
|
||||
type countObserver struct {
|
||||
*fakeDiarStream
|
||||
seen *int
|
||||
}
|
||||
|
||||
func (c *countObserver) finish() error {
|
||||
*c.seen = c.counts + 1 // finish must run before the first query
|
||||
return c.fakeDiarStream.finish()
|
||||
}
|
||||
|
||||
var _ = Describe("distinctSpeakers", func() {
|
||||
It("counts labels, not segments", func() {
|
||||
segs := []*pb.DiarizeSegment{
|
||||
{Speaker: "1"}, {Speaker: "2"}, {Speaker: "1"}, {Speaker: "2"}, {Speaker: "1"},
|
||||
}
|
||||
Expect(distinctSpeakers(segs)).To(Equal(int32(2)))
|
||||
})
|
||||
|
||||
It("counts a single-speaker recording as one", func() {
|
||||
segs := []*pb.DiarizeSegment{{Speaker: "1"}, {Speaker: "1"}, {Speaker: "1"}}
|
||||
Expect(distinctSpeakers(segs)).To(Equal(int32(1)))
|
||||
})
|
||||
|
||||
// Four segments over three labels, not three over three: with the segment
|
||||
// count and the label count equal, a `return len(segs)` would satisfy this
|
||||
// spec and it would assert nothing.
|
||||
It("counts every distinct label once", func() {
|
||||
segs := []*pb.DiarizeSegment{{Speaker: "1"}, {Speaker: "2"}, {Speaker: "3"}, {Speaker: "2"}}
|
||||
Expect(distinctSpeakers(segs)).To(Equal(int32(3)))
|
||||
})
|
||||
|
||||
It("is zero with no segments", func() {
|
||||
Expect(distinctSpeakers(nil)).To(BeZero())
|
||||
})
|
||||
|
||||
// The response field is documented as the count of distinct labels in
|
||||
// `segments`, which is not the model's capacity: Sortformer v2 can label
|
||||
// four speakers whatever the clip actually contains.
|
||||
It("reports what the segments contain, not the model capacity", func() {
|
||||
f := &fakeDiarStream{segs: []cDiarSegment{
|
||||
{StartTime: 0, EndTime: 1, Speaker: 1},
|
||||
{StartTime: 1, EndTime: 2, Speaker: 2},
|
||||
{StartTime: 2, EndTime: 3, Speaker: 1},
|
||||
}}
|
||||
res, err := diarizePCM(f.opener(), []float32{1}, 16000, nil)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(res.GetNumSpeakers()).To(Equal(int32(2)))
|
||||
})
|
||||
|
||||
It("reports no speakers when the runtime found no segments", func() {
|
||||
f := &fakeDiarStream{}
|
||||
res, err := diarizePCM(f.opener(), []float32{1}, 16000, nil)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(res.Segments).To(BeEmpty())
|
||||
Expect(res.GetNumSpeakers()).To(BeZero())
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("collectSegments", func() {
|
||||
It("skips the fill entirely when there is nothing to collect", func() {
|
||||
f := &fakeDiarStream{}
|
||||
segs, err := collectSegments(f)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(segs).To(BeEmpty())
|
||||
Expect(f.counts).To(Equal(1))
|
||||
Expect(f.fills).To(BeZero(), "a zero count must not be followed by a fill")
|
||||
})
|
||||
|
||||
It("sizes the buffer from the query and fills it", func() {
|
||||
f := &fakeDiarStream{segs: []cDiarSegment{
|
||||
{StartTime: 0, EndTime: 1, Speaker: 1},
|
||||
{StartTime: 1, EndTime: 2, Speaker: 2},
|
||||
}}
|
||||
segs, err := collectSegments(f)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(segs).To(HaveLen(2))
|
||||
Expect(segs[1].Speaker).To(Equal(int32(2)))
|
||||
Expect(f.counts).To(Equal(1))
|
||||
Expect(f.fills).To(Equal(1))
|
||||
})
|
||||
|
||||
// nemo_speech_diar_segments writes *count and only then rejects a buffer
|
||||
// that is too small, so the rejected call still reports the size to retry
|
||||
// with. Truncating instead would silently drop turns.
|
||||
It("grows the buffer and retries when the count outran the query", func() {
|
||||
f := &fakeDiarStream{
|
||||
segs: []cDiarSegment{
|
||||
{StartTime: 0, EndTime: 1, Speaker: 1},
|
||||
{StartTime: 1, EndTime: 2, Speaker: 2},
|
||||
{StartTime: 2, EndTime: 3, Speaker: 1},
|
||||
},
|
||||
// The query saw two, the fill found three and rejected the buffer.
|
||||
queryCount: 2,
|
||||
growTo: 3,
|
||||
fillErrs: []error{errors.New("capacity too small (need 3)")},
|
||||
}
|
||||
segs, err := collectSegments(f)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(segs).To(HaveLen(3))
|
||||
Expect(f.fills).To(Equal(2))
|
||||
})
|
||||
|
||||
It("propagates a failure that is not about capacity", func() {
|
||||
f := &fakeDiarStream{
|
||||
segs: []cDiarSegment{{StartTime: 0, EndTime: 1, Speaker: 1}},
|
||||
fillErrs: []error{errors.New("boom")},
|
||||
}
|
||||
_, err := collectSegments(f)
|
||||
Expect(err).To(MatchError(ContainSubstring("boom")))
|
||||
Expect(f.fills).To(Equal(1), "a non-capacity failure must not be retried")
|
||||
})
|
||||
|
||||
// make() panics on a length it cannot satisfy, and a panic in an RPC
|
||||
// handler kills the backend process. A count that could only come from an
|
||||
// uninitialised or corrupted size_t must fail the request instead.
|
||||
It("refuses an implausible count rather than trying to allocate it", func() {
|
||||
f := &fakeDiarStream{queryCount: maxDiarSegments + 1}
|
||||
_, err := collectSegments(f)
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(status.Code(err)).To(Equal(codes.Internal))
|
||||
Expect(err.Error()).To(ContainSubstring("ceiling"))
|
||||
Expect(f.fills).To(BeZero(), "nothing must be allocated or filled for a bad count")
|
||||
})
|
||||
|
||||
It("still accepts a count right at the ceiling", func() {
|
||||
// Only the guard is under test, so the fill is scripted to report zero
|
||||
// rather than actually materialising a hundred megabytes of segments.
|
||||
f := &fakeDiarStream{queryCount: maxDiarSegments}
|
||||
segs, err := collectSegments(f)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(segs).To(BeEmpty())
|
||||
Expect(f.fills).To(Equal(1))
|
||||
})
|
||||
|
||||
It("propagates a failed size query", func() {
|
||||
f := &fakeDiarStream{countErr: errors.New("no stream")}
|
||||
_, err := collectSegments(f)
|
||||
Expect(err).To(MatchError(ContainSubstring("no stream")))
|
||||
Expect(f.fills).To(BeZero())
|
||||
})
|
||||
|
||||
// A runtime whose count grew on every attempt would otherwise loop forever
|
||||
// holding engineMu, which blocks the unload too.
|
||||
It("gives up rather than retrying forever", func() {
|
||||
f := &fakeDiarStream{segs: []cDiarSegment{{StartTime: 0, EndTime: 1, Speaker: 1}}}
|
||||
g := &alwaysGrowing{fakeDiarStream: f}
|
||||
_, err := collectSegments(g)
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(status.Code(err)).To(Equal(codes.Internal))
|
||||
Expect(f.fills).To(Equal(diarSegmentsMaxAttempts))
|
||||
})
|
||||
})
|
||||
|
||||
// alwaysGrowing reports a bigger count on every fill, which is the pathological
|
||||
// case the attempt bound exists for.
|
||||
type alwaysGrowing struct {
|
||||
*fakeDiarStream
|
||||
n uint64
|
||||
}
|
||||
|
||||
func (a *alwaysGrowing) fillSegments([]cDiarSegment) (uint64, error) {
|
||||
a.n += 10
|
||||
a.fills++
|
||||
return a.n, errors.New("capacity too small")
|
||||
}
|
||||
|
||||
// The six frame counts are the one part of the create config that no other
|
||||
// check in the tree can see. The layout assertions pin the struct's shape, and
|
||||
// a wrong VALUE keeps that shape exactly, so without these specs deleting a
|
||||
// sentinel is invisible: c_api.cpp applies left_context_frames at >= 0, so a
|
||||
// dropped -1 there silently pins the model's left context to zero.
|
||||
var _ = Describe("diarModelConfig", func() {
|
||||
It("declares its own size so the runtime accepts the fields", func() {
|
||||
Expect(diarModelConfig(0, -1).Size).To(Equal(unsafe.Sizeof(cDiarModelConfig{})))
|
||||
})
|
||||
|
||||
It("carries the model path and the configured device", func() {
|
||||
cfg := diarModelConfig(0xDEADBEEF, 2)
|
||||
Expect(cfg.ModelPath).To(Equal(uintptr(0xDEADBEEF)))
|
||||
Expect(cfg.GPU).To(Equal(int32(2)))
|
||||
})
|
||||
|
||||
It("passes the CPU sentinel through untouched", func() {
|
||||
Expect(diarModelConfig(0, -1).GPU).To(Equal(int32(-1)))
|
||||
})
|
||||
|
||||
// Asserted field by field rather than as a whole struct so a failure names
|
||||
// the sentinel that went missing.
|
||||
It("leaves every frame-geometry override at the negative sentinel", func() {
|
||||
cfg := diarModelConfig(0, -1)
|
||||
Expect(cfg.ChunkFrames).To(Equal(int32(-1)), "chunk_frames")
|
||||
Expect(cfg.RightContextFrames).To(Equal(int32(-1)), "right_context_frames")
|
||||
Expect(cfg.FIFOFrames).To(Equal(int32(-1)), "fifo_frames")
|
||||
Expect(cfg.SpkcacheFrames).To(Equal(int32(-1)), "spkcache_frames")
|
||||
Expect(cfg.UpdatePeriodFrames).To(Equal(int32(-1)), "update_period_frames")
|
||||
|
||||
// Called out on its own because it is the only one of the six the
|
||||
// runtime applies at >= 0: zero here is a valid explicit left context,
|
||||
// not "unset", so this is the field a dropped sentinel actually breaks.
|
||||
Expect(cfg.LeftContextFrames).To(Equal(int32(-1)), "left_context_frames")
|
||||
Expect(cfg.LeftContextFrames).To(BeNumerically("<", 0),
|
||||
"left_context_frames is applied at >= 0, so a non-negative value pins the geometry")
|
||||
})
|
||||
|
||||
// The preset selects the streaming geometry wholesale, so it has to stay
|
||||
// NULL until there is a model to verify a different one against.
|
||||
It("leaves the preset unset", func() {
|
||||
Expect(diarModelConfig(0, -1).Preset).To(BeZero())
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("segmentationConfig", func() {
|
||||
// A request that set nothing must stay NULL on the C side: diar.h documents
|
||||
// NULL as "library defaults", and those defaults are NeMo's callhome-tuned
|
||||
// values for this checkpoint rather than zeros.
|
||||
It("is absent when the request asked for no postprocessing", func() {
|
||||
Expect(segmentationConfig(&pb.DiarizeRequest{})).To(BeNil())
|
||||
})
|
||||
|
||||
It("maps min_duration_on onto the minimum segment duration", func() {
|
||||
cfg := segmentationConfig(&pb.DiarizeRequest{MinDurationOn: 0.4})
|
||||
Expect(cfg).ToNot(BeNil())
|
||||
Expect(cfg.MinDurationSec).To(BeNumerically("~", 0.4, 1e-6))
|
||||
Expect(cfg.MinGapSec).To(BeZero())
|
||||
})
|
||||
|
||||
It("maps min_duration_off onto the gap fill", func() {
|
||||
cfg := segmentationConfig(&pb.DiarizeRequest{MinDurationOff: 0.25})
|
||||
Expect(cfg).ToNot(BeNil())
|
||||
Expect(cfg.MinGapSec).To(BeNumerically("~", 0.25, 1e-6))
|
||||
Expect(cfg.MinDurationSec).To(BeZero())
|
||||
})
|
||||
|
||||
It("declares its own size so the runtime accepts the fields", func() {
|
||||
cfg := segmentationConfig(&pb.DiarizeRequest{MinDurationOn: 0.4})
|
||||
Expect(cfg.Size).To(Equal(sizeofDiarSegmentationConfig()))
|
||||
})
|
||||
|
||||
// The onset/offset hysteresis is not a clustering threshold and Sortformer
|
||||
// has no clustering stage at all, so mapping one onto the other would be an
|
||||
// invented equivalence. It has to stay unset.
|
||||
It("ignores fields this pipeline has no equivalent for", func() {
|
||||
Expect(segmentationConfig(&pb.DiarizeRequest{
|
||||
NumSpeakers: 2,
|
||||
MinSpeakers: 1,
|
||||
MaxSpeakers: 4,
|
||||
ClusteringThreshold: 0.7,
|
||||
IncludeText: true,
|
||||
Threads: 8,
|
||||
})).To(BeNil())
|
||||
})
|
||||
|
||||
It("ignores non-positive values, which the runtime reads as unset", func() {
|
||||
Expect(segmentationConfig(&pb.DiarizeRequest{MinDurationOn: -1, MinDurationOff: 0})).To(BeNil())
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("unsupportedRequestFields", func() {
|
||||
It("is empty for a request this backend can honour in full", func() {
|
||||
Expect(unsupportedRequestFields(&pb.DiarizeRequest{
|
||||
Dst: "/tmp/a.wav",
|
||||
MinDurationOn: 0.4,
|
||||
MinDurationOff: 0.2,
|
||||
})).To(BeEmpty())
|
||||
})
|
||||
|
||||
It("names every field it had to drop", func() {
|
||||
Expect(unsupportedRequestFields(&pb.DiarizeRequest{
|
||||
NumSpeakers: 2,
|
||||
MinSpeakers: 1,
|
||||
MaxSpeakers: 4,
|
||||
ClusteringThreshold: 0.7,
|
||||
IncludeText: true,
|
||||
Threads: 8,
|
||||
})).To(ConsistOf(
|
||||
"num_speakers", "min_speakers", "max_speakers",
|
||||
"clustering_threshold", "include_text", "threads",
|
||||
))
|
||||
})
|
||||
})
|
||||
@@ -1,116 +0,0 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
|
||||
gguf "github.com/gpustack/gguf-parser-go"
|
||||
)
|
||||
|
||||
// auxOnlyArchitectures are converted NeMo components that attach to a primary
|
||||
// model but are never loadable on their own. Pointing a model config at one is
|
||||
// a configuration mistake worth naming explicitly.
|
||||
var auxOnlyArchitectures = map[string]string{
|
||||
"nemo-nano-codec": "a TTS codec, set it with the codec_model option on a magpietts model",
|
||||
"vad": "a VAD model, set it with the vad_model option on an asr model",
|
||||
"pnc": "a punctuation model, set it with the pnc_model option on an asr model",
|
||||
}
|
||||
|
||||
// familyFor maps a GGUF general.architecture value to a model family.
|
||||
//
|
||||
// Unknown architectures resolve to NMT rather than an error: NMT GGUFs come
|
||||
// from llama.cpp's converter and carry an ordinary LLM architecture, so there
|
||||
// is no NeMo-specific string to match. The user selected this backend
|
||||
// explicitly, which is the signal that the model is meant for it.
|
||||
func familyFor(arch string) (family, error) {
|
||||
if reason, ok := auxOnlyArchitectures[arch]; ok {
|
||||
return familyUnknown, fmt.Errorf(
|
||||
"nemo-speech-cpp: %q is %s, not a model that can be loaded directly", arch, reason)
|
||||
}
|
||||
switch arch {
|
||||
case "asr":
|
||||
return familyASR, nil
|
||||
case "sortformer":
|
||||
return familyDiarization, nil
|
||||
case "magpietts":
|
||||
return familyTTS, nil
|
||||
}
|
||||
return familyNMT, nil
|
||||
}
|
||||
|
||||
// ggufArchitecture reads general.architecture from a GGUF file.
|
||||
func ggufArchitecture(path string) (string, error) {
|
||||
f, err := gguf.ParseGGUFFile(path, gguf.UseMMap(), gguf.SkipLargeMetadata())
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("nemo-speech-cpp: parse gguf %q: %w", path, err)
|
||||
}
|
||||
kv, found := f.Header.MetadataKV.Index([]string{"general.architecture"})
|
||||
if found == 0 {
|
||||
return "", fmt.Errorf("nemo-speech-cpp: %q has no general.architecture key", path)
|
||||
}
|
||||
arch := kv["general.architecture"]
|
||||
// ValueString panics on a mistyped key, and a hand-written or half-converted
|
||||
// GGUF is exactly where that happens. This function is the load-time guard;
|
||||
// it reports, it does not take the process down.
|
||||
if arch.ValueType != gguf.GGUFMetadataValueTypeString {
|
||||
return "", fmt.Errorf(
|
||||
"nemo-speech-cpp: %q has a non-string general.architecture (type %v)", path, arch.ValueType)
|
||||
}
|
||||
return arch.ValueString(), nil
|
||||
}
|
||||
|
||||
// discoverTTSAssets fills in codecModel and tokenizerDir when they were not set
|
||||
// explicitly, by scanning the primary GGUF's own directory.
|
||||
//
|
||||
// A missing asset is a hard error rather than a warning: the runtime would
|
||||
// otherwise load and emit garbage audio, which surfaces far from the cause.
|
||||
func discoverTTSAssets(primaryGGUF string, o *loadOptions) error {
|
||||
dir := filepath.Dir(primaryGGUF)
|
||||
|
||||
if o.codecModel == "" {
|
||||
entries, err := os.ReadDir(dir)
|
||||
if err != nil {
|
||||
return fmt.Errorf("nemo-speech-cpp: scan %q for a codec model: %w", dir, err)
|
||||
}
|
||||
for _, e := range entries {
|
||||
if e.IsDir() {
|
||||
continue
|
||||
}
|
||||
name := e.Name()
|
||||
candidate := filepath.Join(dir, name)
|
||||
// Skip the primary model itself: a file called nanocodec-magpie.gguf
|
||||
// would otherwise be selected as its own codec. Compare basenames,
|
||||
// because candidate is Cleaned by filepath.Join while primaryGGUF
|
||||
// arrives as the caller wrote it, so "/models//magpie.gguf" would
|
||||
// slip past a whole-path equality.
|
||||
if name == filepath.Base(primaryGGUF) {
|
||||
continue
|
||||
}
|
||||
if strings.Contains(strings.ToLower(name), "nanocodec") ||
|
||||
strings.Contains(strings.ToLower(name), "nano-codec") {
|
||||
o.codecModel = candidate
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
if o.codecModel == "" {
|
||||
return fmt.Errorf(
|
||||
"nemo-speech-cpp: no NanoCodec GGUF found next to %q, set the codec_model option",
|
||||
primaryGGUF)
|
||||
}
|
||||
|
||||
if o.tokenizerDir == "" {
|
||||
candidate := filepath.Join(dir, "extracted")
|
||||
if st, err := os.Stat(candidate); err == nil && st.IsDir() {
|
||||
o.tokenizerDir = candidate
|
||||
}
|
||||
}
|
||||
if o.tokenizerDir == "" {
|
||||
return fmt.Errorf(
|
||||
"nemo-speech-cpp: no tokenizer directory found next to %q, set the tokenizer_dir option",
|
||||
primaryGGUF)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -1,168 +0,0 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"encoding/binary"
|
||||
"os"
|
||||
"path/filepath"
|
||||
|
||||
gguf "github.com/gpustack/gguf-parser-go"
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
)
|
||||
|
||||
var _ = Describe("familyFor", func() {
|
||||
It("maps the NeMo architectures to their families", func() {
|
||||
for arch, want := range map[string]family{
|
||||
"asr": familyASR,
|
||||
"sortformer": familyDiarization,
|
||||
"magpietts": familyTTS,
|
||||
} {
|
||||
got, err := familyFor(arch)
|
||||
Expect(err).ToNot(HaveOccurred(), "arch %q", arch)
|
||||
Expect(got).To(Equal(want), "arch %q", arch)
|
||||
}
|
||||
})
|
||||
|
||||
It("treats an unknown architecture as NMT", func() {
|
||||
// NMT GGUFs are produced by llama.cpp's converter, so they carry an LLM
|
||||
// architecture such as qwen3 rather than a NeMo-specific string.
|
||||
got, err := familyFor("qwen3")
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(got).To(Equal(familyNMT))
|
||||
})
|
||||
|
||||
It("rejects an auxiliary-only architecture as a primary model", func() {
|
||||
for _, arch := range []string{"nemo-nano-codec", "vad", "pnc"} {
|
||||
_, err := familyFor(arch)
|
||||
Expect(err).To(HaveOccurred(), "arch %q", arch)
|
||||
Expect(err.Error()).To(ContainSubstring(arch))
|
||||
}
|
||||
})
|
||||
})
|
||||
|
||||
// ggufWithArchValue builds a minimal GGUF v3 carrying general.architecture as
|
||||
// its single metadata entry, with the caller's value type and encoded value.
|
||||
func ggufWithArchValue(valueType gguf.GGUFMetadataValueType, value []byte) []byte {
|
||||
const key = "general.architecture"
|
||||
|
||||
var b []byte
|
||||
b = append(b, 'G', 'G', 'U', 'F')
|
||||
b = binary.LittleEndian.AppendUint32(b, 3) // version
|
||||
b = binary.LittleEndian.AppendUint64(b, 0) // tensor count
|
||||
b = binary.LittleEndian.AppendUint64(b, 1) // metadata kv count
|
||||
b = binary.LittleEndian.AppendUint64(b, uint64(len(key)))
|
||||
b = append(b, key...)
|
||||
b = binary.LittleEndian.AppendUint32(b, uint32(valueType))
|
||||
return append(b, value...)
|
||||
}
|
||||
|
||||
// writeGGUFWithUint32Arch writes a minimal GGUF v3 whose single metadata entry
|
||||
// is general.architecture typed UINT32 rather than STRING. Handwritten and
|
||||
// half-converted files really do carry mistyped keys, and the parser hands them
|
||||
// back rather than rejecting them.
|
||||
func writeGGUFWithUint32Arch(path string) {
|
||||
b := ggufWithArchValue(gguf.GGUFMetadataValueTypeUint32, binary.LittleEndian.AppendUint32(nil, 7))
|
||||
ExpectWithOffset(1, os.WriteFile(path, b, 0o600)).To(Succeed())
|
||||
}
|
||||
|
||||
// writeGGUFWithArch writes a minimal GGUF v3 that parses cleanly and reports
|
||||
// arch as its general.architecture. It is the only way to reach the code past
|
||||
// ggufArchitecture in a test, since there are no real NeMo GGUFs to point at.
|
||||
func writeGGUFWithArch(path, arch string) {
|
||||
v := binary.LittleEndian.AppendUint64(nil, uint64(len(arch)))
|
||||
v = append(v, arch...)
|
||||
ExpectWithOffset(1, os.WriteFile(path, ggufWithArchValue(gguf.GGUFMetadataValueTypeString, v), 0o600)).To(Succeed())
|
||||
}
|
||||
|
||||
var _ = Describe("ggufArchitecture", func() {
|
||||
It("returns an error rather than panicking on a file that is not a GGUF", func() {
|
||||
p := filepath.Join(GinkgoT().TempDir(), "not-a-model.gguf")
|
||||
Expect(os.WriteFile(p, []byte("definitely not a gguf header"), 0o600)).To(Succeed())
|
||||
|
||||
_, err := ggufArchitecture(p)
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring(p))
|
||||
})
|
||||
|
||||
It("returns an error rather than panicking when general.architecture is not a string", func() {
|
||||
p := filepath.Join(GinkgoT().TempDir(), "mistyped-arch.gguf")
|
||||
writeGGUFWithUint32Arch(p)
|
||||
|
||||
arch, err := ggufArchitecture(p)
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(arch).To(BeEmpty())
|
||||
Expect(err.Error()).To(ContainSubstring("general.architecture"))
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("discoverTTSAssets", func() {
|
||||
var dir string
|
||||
|
||||
BeforeEach(func() {
|
||||
dir = GinkgoT().TempDir()
|
||||
})
|
||||
|
||||
write := func(name string) string {
|
||||
p := filepath.Join(dir, name)
|
||||
Expect(os.WriteFile(p, []byte("x"), 0o600)).To(Succeed())
|
||||
return p
|
||||
}
|
||||
|
||||
It("finds a sibling nanocodec gguf and extracted dir", func() {
|
||||
primary := write("magpie.f16.gguf")
|
||||
codec := write("nemo-nano-codec-22khz.f16.gguf")
|
||||
Expect(os.Mkdir(filepath.Join(dir, "extracted"), 0o755)).To(Succeed())
|
||||
|
||||
o := loadOptions{}
|
||||
Expect(discoverTTSAssets(primary, &o)).To(Succeed())
|
||||
Expect(o.codecModel).To(Equal(codec))
|
||||
Expect(o.tokenizerDir).To(Equal(filepath.Join(dir, "extracted")))
|
||||
})
|
||||
|
||||
It("never selects the primary gguf as its own codec", func() {
|
||||
// A file named so it would match a naive *.gguf scan.
|
||||
primary := write("nanocodec-magpie.gguf")
|
||||
Expect(os.Mkdir(filepath.Join(dir, "extracted"), 0o755)).To(Succeed())
|
||||
|
||||
o := loadOptions{}
|
||||
err := discoverTTSAssets(primary, &o)
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("codec_model"))
|
||||
})
|
||||
|
||||
It("never selects the primary gguf as its own codec through an uncleaned path", func() {
|
||||
// LocalAI joins the model directory and the model name itself, so a
|
||||
// trailing separator on ModelPath produces a doubled slash here. The
|
||||
// self-codec guard has to survive that.
|
||||
write("nanocodec-magpie.gguf")
|
||||
primary := dir + "//nanocodec-magpie.gguf"
|
||||
Expect(os.Mkdir(filepath.Join(dir, "extracted"), 0o755)).To(Succeed())
|
||||
|
||||
o := loadOptions{}
|
||||
err := discoverTTSAssets(primary, &o)
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(o.codecModel).To(BeEmpty())
|
||||
Expect(err.Error()).To(ContainSubstring("codec_model"))
|
||||
})
|
||||
|
||||
It("does not overwrite explicitly configured paths", func() {
|
||||
primary := write("magpie.f16.gguf")
|
||||
write("nemo-nano-codec.gguf")
|
||||
Expect(os.Mkdir(filepath.Join(dir, "extracted"), 0o755)).To(Succeed())
|
||||
|
||||
o := loadOptions{codecModel: "/explicit/codec.gguf", tokenizerDir: "/explicit/tok"}
|
||||
Expect(discoverTTSAssets(primary, &o)).To(Succeed())
|
||||
Expect(o.codecModel).To(Equal("/explicit/codec.gguf"))
|
||||
Expect(o.tokenizerDir).To(Equal("/explicit/tok"))
|
||||
})
|
||||
|
||||
It("names the missing option key when the tokenizer dir cannot be found", func() {
|
||||
primary := write("magpie.f16.gguf")
|
||||
write("nemo-nano-codec.gguf")
|
||||
|
||||
o := loadOptions{}
|
||||
err := discoverTTSAssets(primary, &o)
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("tokenizer_dir"))
|
||||
})
|
||||
})
|
||||
@@ -1,76 +0,0 @@
|
||||
package main
|
||||
|
||||
// Started internally by LocalAI, one gRPC server per loaded model.
|
||||
//
|
||||
// Binds NVIDIA NeMo-Speech.cpp through purego. The runtime splits its C ABI
|
||||
// across three shared objects: asr (which also exports the diarization
|
||||
// symbols), tts, and nmt. Library names can be overridden with
|
||||
// NEMO_SPEECH_ASR_LIBRARY / _TTS_LIBRARY / _NMT_LIBRARY, mirroring the
|
||||
// PARAKEET_LIBRARY convention in the sibling backends.
|
||||
//
|
||||
// The naming is asymmetric on purpose: upstream links a dedicated
|
||||
// libnemo_speech_asr_c / libnemo_speech_nmt_c around a private C++ core, but
|
||||
// compiles the TTS c_api straight into libnemo_speech_tts and only aliases the
|
||||
// nemo_speech_tts_c CMake target, so there is no libnemo_speech_tts_c on disk.
|
||||
import (
|
||||
"flag"
|
||||
"fmt"
|
||||
"os"
|
||||
"runtime"
|
||||
|
||||
"github.com/ebitengine/purego"
|
||||
grpc "github.com/mudler/LocalAI/pkg/grpc"
|
||||
)
|
||||
|
||||
var addr = flag.String("addr", "localhost:50051", "the address to connect to")
|
||||
|
||||
// libSuffix is the platform's shared-object extension.
|
||||
func libSuffix() string {
|
||||
if runtime.GOOS == "darwin" {
|
||||
return ".dylib"
|
||||
}
|
||||
return ".so"
|
||||
}
|
||||
|
||||
// libraryName resolves an override env var, falling back to the platform name.
|
||||
func libraryName(envVar, base string) string {
|
||||
if v := os.Getenv(envVar); v != "" {
|
||||
return v
|
||||
}
|
||||
return base + libSuffix()
|
||||
}
|
||||
|
||||
func main() {
|
||||
flag.Parse()
|
||||
|
||||
if err := openLibraries(); err != nil {
|
||||
panic(err)
|
||||
}
|
||||
|
||||
if err := grpc.StartServer(*addr, &NemoSpeech{}); err != nil {
|
||||
panic(err)
|
||||
}
|
||||
}
|
||||
|
||||
// openLibraries dlopens the three C ABI shared objects. All three are opened
|
||||
// eagerly so a packaging mistake fails at startup with a clear message rather
|
||||
// than at first inference of one particular family.
|
||||
func openLibraries() error {
|
||||
for _, l := range []struct {
|
||||
env string
|
||||
base string
|
||||
dst *uintptr
|
||||
}{
|
||||
{"NEMO_SPEECH_ASR_LIBRARY", "libnemo_speech_asr_c", &asrLib},
|
||||
{"NEMO_SPEECH_TTS_LIBRARY", "libnemo_speech_tts", &ttsLib},
|
||||
{"NEMO_SPEECH_NMT_LIBRARY", "libnemo_speech_nmt_c", &nmtLib},
|
||||
} {
|
||||
name := libraryName(l.env, l.base)
|
||||
h, err := purego.Dlopen(name, purego.RTLD_NOW|purego.RTLD_GLOBAL)
|
||||
if err != nil {
|
||||
return fmt.Errorf("nemo-speech-cpp: dlopen %q: %w", name, err)
|
||||
}
|
||||
*l.dst = h
|
||||
}
|
||||
return registerSymbols()
|
||||
}
|
||||
@@ -1,273 +0,0 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"runtime"
|
||||
"sync"
|
||||
"unsafe"
|
||||
|
||||
"github.com/mudler/LocalAI/pkg/grpc/base"
|
||||
pb "github.com/mudler/LocalAI/pkg/grpc/proto"
|
||||
"github.com/mudler/xlog"
|
||||
"google.golang.org/grpc/codes"
|
||||
"google.golang.org/grpc/status"
|
||||
)
|
||||
|
||||
// family is the model family selected at load time from the GGUF architecture.
|
||||
type family int
|
||||
|
||||
const (
|
||||
familyUnknown family = iota
|
||||
familyASR
|
||||
familyDiarization
|
||||
familyTTS
|
||||
familyNMT
|
||||
)
|
||||
|
||||
func (f family) String() string {
|
||||
switch f {
|
||||
case familyASR:
|
||||
return "asr"
|
||||
case familyDiarization:
|
||||
return "diarization"
|
||||
case familyTTS:
|
||||
return "tts"
|
||||
case familyNMT:
|
||||
return "nmt"
|
||||
}
|
||||
return "unknown"
|
||||
}
|
||||
|
||||
// NemoSpeech is one loaded model. Exactly one of the handles is non-zero,
|
||||
// matching fam.
|
||||
type NemoSpeech struct {
|
||||
base.SingleThread
|
||||
|
||||
fam family
|
||||
opts loadOptions
|
||||
|
||||
// engineMu guards fam and the handles, and serializes calls into the C
|
||||
// runtime for this model. Its participants today are withEngine and Free;
|
||||
// the per-family RPCs in Tasks 6 to 9 join it by routing through withEngine.
|
||||
engineMu sync.Mutex
|
||||
|
||||
// synth and nmt are shortened rather than spelled out: synthesizer and
|
||||
// translator are the names of the two RPC-side interfaces those handles are
|
||||
// wrapped in (tts.go, nmt.go), and a field sharing a name with an interface in
|
||||
// the same package makes every construction site read as a conversion.
|
||||
recognizer uintptr
|
||||
diarizer uintptr
|
||||
synth uintptr
|
||||
nmt uintptr
|
||||
}
|
||||
|
||||
// cstr allocates a NUL-terminated C string and returns its pointer plus a
|
||||
// release function. The empty string maps to a null pointer because the C API
|
||||
// treats NULL and "" as equivalent for every optional field.
|
||||
//
|
||||
// The address leaves the Go type system as a uintptr, which the collector does
|
||||
// not trace, so the bytes are pinned for as long as C may read them. Pinning is
|
||||
// the only mechanism with a documented guarantee here: the config structs hold
|
||||
// raw addresses, and an unpinned Go allocation is free to be collected (and, in
|
||||
// principle, moved) the moment its last traced reference dies.
|
||||
//
|
||||
// The returned pointer is for C only, and the direction is one-way. Converting
|
||||
// it back to an unsafe.Pointer to read the bytes from Go is checked by checkptr
|
||||
// (which -race turns on) and kills the process with
|
||||
//
|
||||
// fatal error: checkptr: pointer arithmetic result points to invalid allocation
|
||||
//
|
||||
// as soon as the address lands inside a Go allocation, which is exactly what
|
||||
// this produces. C reading it is fine because C is not instrumented; Go reading
|
||||
// it back is not.
|
||||
//
|
||||
// The caller MUST defer the release function immediately, in the same statement
|
||||
// that takes the pointer. Dropping it leaks the pin, which the runtime reports
|
||||
// at the next collection as:
|
||||
//
|
||||
// runtime.Pinner: found leaking pinned pointer; forgot to call Unpin()?
|
||||
//
|
||||
// That is loud and wrong-looking on purpose: the alternative failure mode is C
|
||||
// reading freed memory, which shows up as rare corruption with no trace back
|
||||
// to here.
|
||||
func cstr(s string) (uintptr, func()) {
|
||||
if s == "" {
|
||||
return 0, func() {}
|
||||
}
|
||||
b := append([]byte(s), 0)
|
||||
pin := new(runtime.Pinner)
|
||||
pin.Pin(&b[0])
|
||||
// #nosec G103 -- b is non-empty (s != "" above) and &b[0] is pinned on the
|
||||
// previous line, so the address C receives cannot be collected or moved
|
||||
// until the returned release runs. One-way by construction: the doc comment
|
||||
// above forbids converting this uintptr back, which is what keeps checkptr
|
||||
// (and therefore -race) out of it.
|
||||
return uintptr(unsafe.Pointer(&b[0])), func() {
|
||||
if pin == nil {
|
||||
return
|
||||
}
|
||||
pin.Unpin()
|
||||
pin = nil
|
||||
}
|
||||
}
|
||||
|
||||
// There is deliberately no inverse of cstr in this package. Every C entry point
|
||||
// that returns a string is bound in abi.go with a Go `string` return, which
|
||||
// purego converts from the char* itself, so a hand-rolled reader would have no
|
||||
// production caller and would exist only as an unsafe helper waiting to be
|
||||
// pointed at the wrong kind of address. Reach for purego's conversion instead;
|
||||
// if a future symbol genuinely needs the raw char* (to tell NULL from ""), bind
|
||||
// it as uintptr at that call site, where the ownership can be reasoned about.
|
||||
|
||||
// requireFamily gates an RPC on the family selected at load time. Returning
|
||||
// Unimplemented rather than a nil dereference means a misconfigured model YAML
|
||||
// produces a message a user can act on.
|
||||
//
|
||||
// Callers must already hold engineMu: Free writes n.fam under it, so an
|
||||
// unlocked read here is a data race. Use withEngine rather than calling this
|
||||
// directly.
|
||||
func (n *NemoSpeech) requireFamily(want family) error {
|
||||
if n.fam != want {
|
||||
return status.Errorf(codes.Unimplemented,
|
||||
"nemo-speech-cpp: this model was loaded as %s, not %s", n.fam, want)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// withEngine runs fn holding engineMu, having first checked the family.
|
||||
//
|
||||
// Every RPC must go through this rather than calling requireFamily on its own.
|
||||
// pkg/grpc/server.go takes the backend lock around each RPC but calls Free
|
||||
// without it, so a teardown can land mid-request. Checking the family and then
|
||||
// making the C calls that trust it under two separate acquisitions leaves a
|
||||
// window in which Free destroys the handle, and the request goes on to use a
|
||||
// zeroed one.
|
||||
func (n *NemoSpeech) withEngine(want family, fn func() error) error {
|
||||
n.engineMu.Lock()
|
||||
defer n.engineMu.Unlock()
|
||||
|
||||
if err := n.requireFamily(want); err != nil {
|
||||
return err
|
||||
}
|
||||
return fn()
|
||||
}
|
||||
|
||||
func (n *NemoSpeech) Load(opts *pb.ModelOptions) error {
|
||||
modelFile := opts.GetModelFile()
|
||||
if modelFile == "" {
|
||||
return errors.New("nemo-speech-cpp: ModelFile is required")
|
||||
}
|
||||
|
||||
// Free writes fam and the handles under engineMu and runs without the
|
||||
// backend lock that serialises the RPCs (pkg/grpc/server.go), so the
|
||||
// load-side writes to those same fields need the same protection: without
|
||||
// it this is the write-side half of the race withEngine closed on the read
|
||||
// side. n.opts is in here too, since the loaders read it.
|
||||
//
|
||||
// The loaders called below must NOT take engineMu themselves; sync.Mutex is
|
||||
// not reentrant and this is why.
|
||||
n.engineMu.Lock()
|
||||
defer n.engineMu.Unlock()
|
||||
|
||||
n.opts = parseOptions(opts.GetOptions(), opts.GetModelPath())
|
||||
|
||||
arch, err := ggufArchitecture(modelFile)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
fam, err := familyFor(arch)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
xlog.Info("nemo-speech-cpp: loading model", "arch", arch, "family", fam.String())
|
||||
|
||||
// fam is committed only once the family-specific loader has succeeded.
|
||||
// requireFamily is the gate every RPC goes through, so a half-loaded model
|
||||
// that kept its family would route requests at a handle that was never
|
||||
// created.
|
||||
switch fam {
|
||||
case familyASR:
|
||||
err = n.loadASR(modelFile)
|
||||
case familyDiarization:
|
||||
err = n.loadDiarizer(modelFile)
|
||||
case familyTTS:
|
||||
if err = discoverTTSAssets(modelFile, &n.opts); err == nil {
|
||||
err = n.loadTTS(modelFile)
|
||||
}
|
||||
case familyNMT:
|
||||
err = n.loadNMT(modelFile)
|
||||
default:
|
||||
err = fmt.Errorf("nemo-speech-cpp: unhandled family for architecture %q", arch)
|
||||
}
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
n.fam = fam
|
||||
return nil
|
||||
}
|
||||
|
||||
// Free destroys the runtime handle created at load time.
|
||||
//
|
||||
// base.SingleThread.Free is a no-op that derived backends are expected to
|
||||
// override, and every family here owns C memory that only its own destroy
|
||||
// entry point can release, so without this an unloaded model leaks a whole
|
||||
// acoustic model. Clearing fam as well means an RPC that races the unload is
|
||||
// refused by the gate rather than handed a dangling handle, but that only holds
|
||||
// for callers that took engineMu, which today means callers that went through
|
||||
// withEngine.
|
||||
func (n *NemoSpeech) Free() error {
|
||||
n.engineMu.Lock()
|
||||
defer n.engineMu.Unlock()
|
||||
|
||||
// Guarded on the handle, not on fam: a load that failed part way through
|
||||
// leaves fam unset, and the destroy functions are nil pointers until
|
||||
// openLibraries has bound them.
|
||||
// Each is tested independently rather than switched on: the one-handle
|
||||
// invariant is an invariant, and if it ever broke, a switch would silently
|
||||
// leak the others.
|
||||
if n.recognizer != 0 {
|
||||
ASRDestroy(n.recognizer)
|
||||
n.recognizer = 0
|
||||
}
|
||||
if n.diarizer != 0 {
|
||||
DiarDestroy(n.diarizer)
|
||||
n.diarizer = 0
|
||||
}
|
||||
if n.synth != 0 {
|
||||
TTSDestroy(n.synth)
|
||||
n.synth = 0
|
||||
}
|
||||
if n.nmt != 0 {
|
||||
NMTDestroy(n.nmt)
|
||||
n.nmt = 0
|
||||
}
|
||||
n.fam = familyUnknown
|
||||
return nil
|
||||
}
|
||||
|
||||
// The loaders are one per family: loadASR in asr.go, loadDiarizer in diar.go,
|
||||
// loadTTS in tts.go and loadNMT in nmt.go. Each populates its config structs
|
||||
// from n.opts and stores the handle in the matching field.
|
||||
//
|
||||
// Locking protocol, in both directions:
|
||||
//
|
||||
// - Every RPC must hold engineMu across its family check AND its C calls,
|
||||
// which means wrapping its body in withEngine. Free runs without the
|
||||
// backend lock (pkg/grpc/server.go:1019), so anything that checks the
|
||||
// family and then releases the lock before calling C can have the handle
|
||||
// destroyed underneath it. asr.go's AudioTranscription is the worked
|
||||
// example: even the audio decode sits inside the closure, because the
|
||||
// backend already serialises RPCs through base.SingleThread and so the
|
||||
// wider hold costs nothing.
|
||||
// - A loader must NOT take engineMu. Load holds it across the whole switch,
|
||||
// and sync.Mutex is not reentrant, so locking in a loader deadlocks.
|
||||
//
|
||||
// One consequence the streaming RPCs have to plan around: a stream whose body
|
||||
// is wrapped in withEngine holds engineMu for the WHOLE stream, so Free blocks
|
||||
// until the stream ends rather than tearing the handle out from under it. That
|
||||
// is the behaviour we want (a half-closed stream over a destroyed recognizer
|
||||
// has no good outcome), but it means an unload waits on a client that has
|
||||
// stopped sending, so a streaming loop must have its own way out: honour the
|
||||
// request context and stop on it, rather than blocking forever on the next
|
||||
// chunk.
|
||||
@@ -1,13 +0,0 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
)
|
||||
|
||||
func TestNemoSpeech(t *testing.T) {
|
||||
RegisterFailHandler(Fail)
|
||||
RunSpecs(t, "nemo-speech-cpp Backend Suite")
|
||||
}
|
||||
@@ -1,260 +0,0 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"runtime"
|
||||
"sync"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
"google.golang.org/grpc/codes"
|
||||
"google.golang.org/grpc/status"
|
||||
|
||||
pb "github.com/mudler/LocalAI/pkg/grpc/proto"
|
||||
)
|
||||
|
||||
var _ = Describe("requireFamily", func() {
|
||||
It("accepts the loaded family", func() {
|
||||
n := &NemoSpeech{fam: familyASR}
|
||||
Expect(n.requireFamily(familyASR)).To(Succeed())
|
||||
})
|
||||
|
||||
It("rejects a mismatched family with Unimplemented and names both", func() {
|
||||
n := &NemoSpeech{fam: familyTTS}
|
||||
err := n.requireFamily(familyASR)
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(status.Code(err)).To(Equal(codes.Unimplemented))
|
||||
Expect(err.Error()).To(ContainSubstring("tts"))
|
||||
Expect(err.Error()).To(ContainSubstring("asr"))
|
||||
})
|
||||
|
||||
It("rejects an unloaded model", func() {
|
||||
n := &NemoSpeech{}
|
||||
err := n.requireFamily(familyASR)
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(status.Code(err)).To(Equal(codes.Unimplemented))
|
||||
})
|
||||
|
||||
It("rejects every family when the model is unloaded", func() {
|
||||
n := &NemoSpeech{}
|
||||
for _, f := range []family{familyASR, familyDiarization, familyTTS, familyNMT} {
|
||||
Expect(n.requireFamily(f)).To(HaveOccurred(), "family %s must be gated on an unloaded model", f)
|
||||
}
|
||||
})
|
||||
})
|
||||
|
||||
// The brief's round-trip spec (cstr then a reader) cannot exist: cstr pins a Go
|
||||
// allocation, and converting a uintptr back into a pointer to Go memory is a
|
||||
// checkptr violation that aborts the process under -race. So cstr is asserted
|
||||
// on what is observable without dereferencing its result.
|
||||
var _ = Describe("cstr", func() {
|
||||
It("returns a non-null pointer for a non-empty string", func() {
|
||||
p, free := cstr("hello")
|
||||
defer free()
|
||||
Expect(p).ToNot(BeZero())
|
||||
})
|
||||
|
||||
It("returns a null pointer for the empty string", func() {
|
||||
// The C API documents NULL and "" as equivalent for optional fields, and
|
||||
// passing NULL avoids allocating for every unset option.
|
||||
p, free := cstr("")
|
||||
defer free()
|
||||
Expect(p).To(BeZero())
|
||||
})
|
||||
|
||||
// The pin has to hold for the whole create call, which spans at least one
|
||||
// safepoint. A collection must therefore neither move nor invalidate the
|
||||
// address that C was handed.
|
||||
It("keeps the pointer stable across a garbage collection", func() {
|
||||
p, free := cstr("/models/nemo/parakeet.gguf")
|
||||
defer free()
|
||||
before := p
|
||||
runtime.GC()
|
||||
runtime.GC()
|
||||
Expect(p).To(Equal(before))
|
||||
})
|
||||
|
||||
It("survives releasing more than once", func() {
|
||||
_, free := cstr("twice")
|
||||
free()
|
||||
Expect(free).ToNot(Panic())
|
||||
})
|
||||
|
||||
// A dropped release leaks the pin, and the runtime turns that into a process
|
||||
// abort at some later collection. Nothing can catch it, so this only pins the
|
||||
// contract in prose: release in the same statement that takes the pointer.
|
||||
It("releases without panicking when used as documented", func() {
|
||||
Expect(func() {
|
||||
p, free := cstr("released")
|
||||
defer free()
|
||||
_ = p
|
||||
}).ToNot(Panic())
|
||||
})
|
||||
})
|
||||
|
||||
// pkg/grpc/server.go:1019 calls Free without taking the backend lock every
|
||||
// other RPC holds, so a teardown really can land while a request is in flight.
|
||||
// The family check and the C calls that trust it therefore have to happen under
|
||||
// engineMu together, or Free can destroy the handle in the gap between them.
|
||||
var _ = Describe("engine locking", func() {
|
||||
It("serialises a teardown against an in-flight request", func() {
|
||||
n := &NemoSpeech{fam: familyASR}
|
||||
|
||||
var wg sync.WaitGroup
|
||||
wg.Add(2)
|
||||
go func() {
|
||||
defer GinkgoRecover()
|
||||
defer wg.Done()
|
||||
for i := 0; i < 2000; i++ {
|
||||
// Errors are expected once the teardown wins the race; what must
|
||||
// not happen is an unsynchronised read of the family.
|
||||
_ = n.withEngine(familyASR, func() error { return nil })
|
||||
}
|
||||
}()
|
||||
go func() {
|
||||
defer GinkgoRecover()
|
||||
defer wg.Done()
|
||||
for i := 0; i < 2000; i++ {
|
||||
Expect(n.Free()).To(Succeed())
|
||||
}
|
||||
}()
|
||||
wg.Wait()
|
||||
})
|
||||
|
||||
It("refuses the body when the family does not match, and still unlocks", func() {
|
||||
n := &NemoSpeech{fam: familyTTS}
|
||||
called := false
|
||||
err := n.withEngine(familyASR, func() error {
|
||||
called = true
|
||||
return nil
|
||||
})
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(status.Code(err)).To(Equal(codes.Unimplemented))
|
||||
Expect(called).To(BeFalse())
|
||||
|
||||
// A lock leaked on the rejection path would deadlock the next request
|
||||
// rather than fail it, so prove the mutex is free afterwards.
|
||||
Expect(n.engineMu.TryLock()).To(BeTrue())
|
||||
n.engineMu.Unlock()
|
||||
})
|
||||
|
||||
It("propagates the body's error and still unlocks", func() {
|
||||
n := &NemoSpeech{fam: familyASR}
|
||||
boom := errors.New("boom")
|
||||
Expect(n.withEngine(familyASR, func() error { return boom })).To(MatchError(boom))
|
||||
Expect(n.engineMu.TryLock()).To(BeTrue())
|
||||
n.engineMu.Unlock()
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("Free", func() {
|
||||
// The destroy entry points are nil function values until openLibraries has
|
||||
// bound them, so an unloaded model must not reach them. LocalAI frees every
|
||||
// backend it shuts down, including one whose Load failed.
|
||||
It("is a no-op on a model that was never loaded", func() {
|
||||
n := &NemoSpeech{}
|
||||
Expect(n.Free()).To(Succeed())
|
||||
})
|
||||
|
||||
It("is idempotent", func() {
|
||||
n := &NemoSpeech{}
|
||||
Expect(n.Free()).To(Succeed())
|
||||
Expect(n.Free()).To(Succeed())
|
||||
})
|
||||
|
||||
It("does not reach the runtime for a load that failed part way through", func() {
|
||||
n := &NemoSpeech{}
|
||||
path := filepath.Join(GinkgoT().TempDir(), "broken.gguf")
|
||||
Expect(os.WriteFile(path, []byte("broken"), 0o600)).To(Succeed())
|
||||
Expect(n.Load(&pb.ModelOptions{ModelFile: path})).ToNot(Succeed())
|
||||
Expect(n.Free()).To(Succeed())
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("Load", func() {
|
||||
It("rejects an empty model file", func() {
|
||||
n := &NemoSpeech{}
|
||||
err := n.Load(&pb.ModelOptions{})
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("ModelFile"))
|
||||
})
|
||||
|
||||
It("reports a model file that does not exist", func() {
|
||||
n := &NemoSpeech{}
|
||||
missing := filepath.Join(GinkgoT().TempDir(), "absent.gguf")
|
||||
err := n.Load(&pb.ModelOptions{ModelFile: missing})
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("absent.gguf"))
|
||||
})
|
||||
|
||||
It("reports a file that is not a GGUF", func() {
|
||||
n := &NemoSpeech{}
|
||||
path := filepath.Join(GinkgoT().TempDir(), "notagguf.gguf")
|
||||
Expect(os.WriteFile(path, []byte("this is not a gguf file at all"), 0o600)).To(Succeed())
|
||||
err := n.Load(&pb.ModelOptions{ModelFile: path})
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("nemo-speech-cpp"))
|
||||
})
|
||||
|
||||
// A failed load must not leave a family selected, or the RPC gate would wave
|
||||
// requests through to a nil handle.
|
||||
It("leaves no family selected when the load fails", func() {
|
||||
n := &NemoSpeech{}
|
||||
path := filepath.Join(GinkgoT().TempDir(), "broken.gguf")
|
||||
Expect(os.WriteFile(path, []byte("broken"), 0o600)).To(Succeed())
|
||||
Expect(n.Load(&pb.ModelOptions{ModelFile: path})).ToNot(Succeed())
|
||||
Expect(n.fam).To(Equal(familyUnknown))
|
||||
Expect(n.requireFamily(familyASR)).To(HaveOccurred())
|
||||
})
|
||||
|
||||
// The load path picks a family and only then runs that family's loader, so
|
||||
// there is a window where the family is known and the load still fails.
|
||||
// Committing n.fam before the loader runs would leave the RPC gate open on a
|
||||
// handle that was never created, and pkg/grpc/server.go keeps serving the
|
||||
// instance after a failed LoadModel, so the next request really would reach
|
||||
// it. TTS is the only family whose loader can fail before touching C.
|
||||
It("does not select the family until that family's loader has succeeded", func() {
|
||||
dir := GinkgoT().TempDir()
|
||||
path := filepath.Join(dir, "magpie.f16.gguf")
|
||||
writeGGUFWithArch(path, "magpietts")
|
||||
|
||||
// Self-guard: if the handwritten GGUF ever stops parsing, Load would fail
|
||||
// at ggufArchitecture instead, before a family is ever chosen, and the
|
||||
// assertions below would pass without exercising the ordering at all.
|
||||
Expect(ggufArchitecture(path)).To(Equal("magpietts"))
|
||||
|
||||
// No sibling codec in the directory, so discoverTTSAssets fails after
|
||||
// familyFor has already resolved familyTTS.
|
||||
n := &NemoSpeech{}
|
||||
err := n.Load(&pb.ModelOptions{ModelFile: path})
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("codec_model"))
|
||||
|
||||
Expect(n.fam).To(Equal(familyUnknown))
|
||||
Expect(n.requireFamily(familyTTS)).To(HaveOccurred())
|
||||
})
|
||||
|
||||
It("closes the family gate again after a free", func() {
|
||||
n := &NemoSpeech{fam: familyASR}
|
||||
Expect(n.Free()).To(Succeed())
|
||||
Expect(n.fam).To(Equal(familyUnknown))
|
||||
Expect(n.requireFamily(familyASR)).To(HaveOccurred())
|
||||
})
|
||||
|
||||
It("parses the model options before it touches the model file", func() {
|
||||
// The options are what tell a TTS load where its codec lives, so they have
|
||||
// to be in place before any family-specific loader runs.
|
||||
n := &NemoSpeech{}
|
||||
path := filepath.Join(GinkgoT().TempDir(), "broken.gguf")
|
||||
Expect(os.WriteFile(path, []byte("broken"), 0o600)).To(Succeed())
|
||||
Expect(n.Load(&pb.ModelOptions{
|
||||
ModelFile: path,
|
||||
ModelPath: "/models",
|
||||
Options: []string{"gpu:2", "codec_model:codec.gguf"},
|
||||
})).ToNot(Succeed())
|
||||
Expect(n.opts.gpu).To(Equal(int32(2)))
|
||||
Expect(n.opts.codecModel).To(Equal("/models/codec.gguf"))
|
||||
})
|
||||
})
|
||||
@@ -1,368 +0,0 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"regexp"
|
||||
"runtime"
|
||||
"strings"
|
||||
"unsafe"
|
||||
|
||||
pb "github.com/mudler/LocalAI/pkg/grpc/proto"
|
||||
"github.com/mudler/xlog"
|
||||
"google.golang.org/grpc/codes"
|
||||
"google.golang.org/grpc/status"
|
||||
)
|
||||
|
||||
// pairDirective matches a leading "[src->tgt] " override.
|
||||
//
|
||||
// Each side is an unbounded run of two-letter segments, not one or two of them.
|
||||
// Either side may also be omitted, which keeps the model-level default for it:
|
||||
// resolve_tag accepts a READY pair tag in one field with the other empty
|
||||
// (src/nmt/langpairs.cc:167-172), so "[->en-de]" names a pair for one request.
|
||||
//
|
||||
// The two rules together are what force the unbounded run. A regional code on
|
||||
// its own is only two segments (pt-br, zh-cn, es-us) and would parse under a
|
||||
// stricter pattern; it is the SINGLE-FIELD form of a regional pair that runs to
|
||||
// three (en-zh-cn, en-zh-tw, en-es-us, en-pt-br, pt-br-en, zh-tw-en). And the
|
||||
// failure is not a mis-split: a pattern too short to cover the tag does not
|
||||
// match the directive at all, so the whole bracket survives into the text and
|
||||
// is handed to the model as something to translate.
|
||||
//
|
||||
// The codes are not normalised or validated here. normalize_language_code
|
||||
// lowercases and folds BCP-47 down to a supported base, and is_supported has the
|
||||
// authoritative table; duplicating either would be a second source of truth that
|
||||
// drifts on the next pin bump.
|
||||
var pairDirective = regexp.MustCompile(`^\[\s*([a-zA-Z]{2}(?:-[a-zA-Z]{2})*)?\s*->\s*([a-zA-Z]{2}(?:-[a-zA-Z]{2})*)?\s*\]\s*`)
|
||||
|
||||
// translator is the NMT half of the C API, narrowed to what Predict uses.
|
||||
//
|
||||
// It is an interface for the same reason synthesizer and diarStream are: no
|
||||
// Riva-Translate GGUF is small enough to keep in the tree, so the layer above
|
||||
// the ABI (pair resolution, validation, the text array, the single-chunk stream)
|
||||
// would otherwise have no test at all. A fake here scripts what the C API
|
||||
// returns; it does not pretend to translate anything.
|
||||
type translator interface {
|
||||
// translate returns one translation per input text, in order.
|
||||
translate(texts []string, source, target string) ([]string, error)
|
||||
}
|
||||
|
||||
// cTranslator is the real translator, over one nemo_speech_nmt_translator.
|
||||
type cTranslator struct {
|
||||
handle uintptr
|
||||
}
|
||||
|
||||
// nmtTexts builds the `const char* const* texts` argument and returns it with
|
||||
// the release the caller MUST defer.
|
||||
//
|
||||
// Two levels need pinning, not one. cstr pins each string's bytes, but the array
|
||||
// carrying their addresses is a separate Go allocation holding uintptrs: the
|
||||
// collector neither traces through it nor is obliged to leave it where it is,
|
||||
// and C dereferences it for the whole call. Pinning only the strings would leave
|
||||
// the array itself free to move out from under the runtime.
|
||||
//
|
||||
// An empty element is refused rather than passed on. cstr maps "" to NULL and
|
||||
// src/nmt/c_api.cpp maps a NULL element back to "" (str_or_empty), so a blank
|
||||
// text would come back as a confident translation of nothing rather than an
|
||||
// error.
|
||||
func nmtTexts(texts []string) ([]uintptr, func(), error) {
|
||||
pin := new(runtime.Pinner)
|
||||
// The pin is released first so that the array stops being pinned before the
|
||||
// strings it points at do.
|
||||
frees := []func(){pin.Unpin}
|
||||
release := func() {
|
||||
for _, f := range frees {
|
||||
f()
|
||||
}
|
||||
}
|
||||
|
||||
if len(texts) == 0 {
|
||||
return nil, release, status.Error(codes.InvalidArgument,
|
||||
"nemo-speech-cpp: nothing to translate")
|
||||
}
|
||||
|
||||
ptrs := make([]uintptr, len(texts))
|
||||
for i, t := range texts {
|
||||
if t == "" {
|
||||
return nil, release, status.Error(codes.InvalidArgument,
|
||||
"nemo-speech-cpp: nothing to translate")
|
||||
}
|
||||
p, free := cstr(t)
|
||||
frees = append(frees, free)
|
||||
ptrs[i] = p
|
||||
}
|
||||
pin.Pin(&ptrs[0])
|
||||
return ptrs, release, nil
|
||||
}
|
||||
|
||||
func (t *cTranslator) translate(texts []string, source, target string) ([]string, error) {
|
||||
ptrs, release, err := nmtTexts(texts)
|
||||
if err != nil {
|
||||
release()
|
||||
return nil, err
|
||||
}
|
||||
defer release()
|
||||
|
||||
// source and target cross as Go strings: purego NUL-terminates and copies
|
||||
// them itself for the duration of the call, and c_api.cpp deep-copies both
|
||||
// into std::string before doing anything with them.
|
||||
var result uintptr
|
||||
if st := NMTTranslate(t.handle, &ptrs[0], uint64(len(ptrs)), source, target, &result); st != 0 {
|
||||
// An unsupported language pair arrives here as INVALID_ARGUMENT
|
||||
// (src/nmt/translator.cpp throws std::invalid_argument, which
|
||||
// src/nmt/c_api.cpp's guard maps to it), which statusErrorf turns into
|
||||
// the caller-facing code rather than Internal.
|
||||
return nil, statusErrorf(st, "nemo-speech-cpp: translate: %s", NMTLastError())
|
||||
}
|
||||
defer NMTResultDestroy(result)
|
||||
|
||||
count := NMTResultCount(result)
|
||||
out := make([]string, 0, count)
|
||||
for i := uint64(0); i < count; i++ {
|
||||
out = append(out, NMTResultText(result, i))
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// nmtTranslatorConfig builds the create-time config.
|
||||
//
|
||||
// Extracted from loadNMT so its four adjacent pointer fields can be asserted
|
||||
// against distinct sentinels. Backend, Model, Generation and Pool are all
|
||||
// uintptr and all sit next to each other, so transposing two of them changes
|
||||
// neither the struct's size nor any field's offset: the layout assertions in
|
||||
// abi_test.go are blind to it, and what it produces at runtime is the backend
|
||||
// config being read as the model config.
|
||||
//
|
||||
// Generation and Pool stay NULL, which nmt.h documents as "library defaults":
|
||||
// max_new_tokens (256) and contexts (1) are create-time settings this backend
|
||||
// has no option to fill them from, and PredictOptions carries no per-request
|
||||
// equivalent that a create-time config could honour anyway.
|
||||
//
|
||||
// backend and model are pinned addresses, not Go pointers, and the caller owns
|
||||
// the pins.
|
||||
func nmtTranslatorConfig(backend, model uintptr) cNMTTranslatorConfig {
|
||||
return cNMTTranslatorConfig{
|
||||
Size: unsafe.Sizeof(cNMTTranslatorConfig{}),
|
||||
Backend: backend,
|
||||
Model: model,
|
||||
}
|
||||
}
|
||||
|
||||
// loadNMT creates the Riva-Translate translator.
|
||||
//
|
||||
// This must not take engineMu: Load is its only caller and already holds it.
|
||||
func (n *NemoSpeech) loadNMT(modelFile string) error {
|
||||
// nemo_speech_nmt_create deep-copies the path into a std::string
|
||||
// (src/nmt/c_api.cpp to_config, via str_or_empty) and retains no pointer
|
||||
// afterwards, so pinning for the duration of the create call is both
|
||||
// necessary and sufficient.
|
||||
var pinner runtime.Pinner
|
||||
defer pinner.Unpin()
|
||||
|
||||
pathP, freePath := cstr(modelFile)
|
||||
defer freePath()
|
||||
|
||||
// NCtx is left at 0, which to_config reads as "keep the default" (it applies
|
||||
// the field only when > 0) and which the runtime resolves to 1024 tokens.
|
||||
// That is sized for the sentence-length input Riva-Translate is built for,
|
||||
// and raising it costs one n_ctx-sized KV cache per pooled context, so it
|
||||
// wants a deliberate option rather than a guess made here.
|
||||
model := cNMTModelConfig{Size: unsafe.Sizeof(cNMTModelConfig{}), Path: pathP}
|
||||
// BackendConfig.gpu defaults to 0 in C++ (device 0), not to CPU, and
|
||||
// to_config assigns it unconditionally, so the option's own -1 default is
|
||||
// what keeps an unconfigured model on the CPU.
|
||||
backend := cNMTBackendConfig{Size: unsafe.Sizeof(cNMTBackendConfig{}), GPU: n.opts.gpu}
|
||||
|
||||
cfg := nmtTranslatorConfig(pinPtr(&pinner, &backend), pinPtr(&pinner, &model))
|
||||
|
||||
xlog.Info("nemo-speech-cpp: creating translator",
|
||||
"gpu", n.opts.gpu,
|
||||
"source_language", n.opts.sourceLanguage,
|
||||
"target_language", n.opts.targetLanguage)
|
||||
|
||||
// #nosec G103 -- cfg is a local POD struct borrowed for this call only. Its
|
||||
// Backend and Model members are pinPtr addresses held by the pinner unpinned
|
||||
// on return, Model.Path is the cstr allocation freed by the defer above, and
|
||||
// nemo_speech_nmt_create deep-copies everything it reads.
|
||||
if st := NMTCreate(unsafe.Pointer(&cfg), &n.nmt); st != 0 {
|
||||
return statusErrorf(st, "nemo-speech-cpp: nmt create: %s", NMTLastError())
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// languagePair resolves the languages for one request and returns the text to
|
||||
// translate.
|
||||
//
|
||||
// nemo_speech_nmt_translate takes explicit source and target languages and has
|
||||
// no free-form generation entry point at all, so there is no prompt in the LLM
|
||||
// sense to carry an instruction. The pair therefore comes from the model
|
||||
// options, and a leading "[src->tgt]" directive is the only per-request control
|
||||
// Predict can offer.
|
||||
func (n *NemoSpeech) languagePair(prompt string) (source, target, text string) {
|
||||
source, target = n.opts.sourceLanguage, n.opts.targetLanguage
|
||||
|
||||
m := pairDirective.FindStringSubmatch(prompt)
|
||||
if m == nil {
|
||||
return source, target, strings.TrimSpace(prompt)
|
||||
}
|
||||
// An omitted side keeps the model-level default rather than blanking it.
|
||||
if m[1] != "" {
|
||||
source = m[1]
|
||||
}
|
||||
if m[2] != "" {
|
||||
target = m[2]
|
||||
}
|
||||
// The directive must not survive into the text: the runtime wraps it in a
|
||||
// chat template (src/nmt/langpairs.cc build_prompt), so anything left here is
|
||||
// translated along with the sentence.
|
||||
return source, target, strings.TrimSpace(prompt[len(m[0]):])
|
||||
}
|
||||
|
||||
// unsupportedPredictFields names the PredictOptions fields a caller may have set
|
||||
// that this C API has no way to honour, so they are logged rather than silently
|
||||
// dropped.
|
||||
//
|
||||
// The list is deliberately narrow. Everything nemo_speech_nmt_translate accepts
|
||||
// is in its five arguments: a translator, the texts, and two language codes.
|
||||
// Everything else in PredictOptions is therefore unsupported, and naming all of
|
||||
// it would log on every single request, because LocalAI fills the sampling
|
||||
// defaults in from the model config whether or not the user asked for them.
|
||||
//
|
||||
// So the sampling and decoding knobs (temperature, top_p, top_k, min_p, seed,
|
||||
// tokens, repeat/frequency/presence penalties, mirostat, tfz, typical_p,
|
||||
// stop_prompts, prompt caching, rope scaling, n_draft, logit_bias) are ignored
|
||||
// silently: there is no field for any of them on either side of the ABI.
|
||||
// max_new_tokens and n_ctx exist but are CREATE-time settings on the translator,
|
||||
// not per-request ones, so PredictOptions.Tokens has nowhere to go either.
|
||||
//
|
||||
// What is named here is the structural asks: requests that only make sense
|
||||
// against a general language model, where honouring them partially would be
|
||||
// worse than saying nothing at all.
|
||||
func unsupportedPredictFields(opts *pb.PredictOptions) []string {
|
||||
var out []string
|
||||
if opts.GetGrammar() != "" {
|
||||
out = append(out, "grammar")
|
||||
}
|
||||
if opts.GetTools() != "" {
|
||||
out = append(out, "tools")
|
||||
}
|
||||
if len(opts.GetImages()) > 0 {
|
||||
out = append(out, "images")
|
||||
}
|
||||
if len(opts.GetVideos()) > 0 {
|
||||
out = append(out, "videos")
|
||||
}
|
||||
if len(opts.GetAudios()) > 0 {
|
||||
out = append(out, "audios")
|
||||
}
|
||||
if opts.GetNegativePrompt() != "" {
|
||||
out = append(out, "negative_prompt")
|
||||
}
|
||||
if opts.GetLogprobs() > 0 {
|
||||
out = append(out, "logprobs")
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// translateText runs one translation and returns it.
|
||||
//
|
||||
// The two rejections happen before anything crosses the ABI. An empty text would
|
||||
// otherwise reach the runtime as a NULL element (see nmtTexts), and a missing
|
||||
// target would come back as "unsupported language pair: -> ", which names
|
||||
// neither the option the operator has to set nor the request that failed.
|
||||
func translateText(t translator, source, target, text string) (string, error) {
|
||||
if text == "" {
|
||||
return "", status.Error(codes.InvalidArgument,
|
||||
"nemo-speech-cpp: PredictOptions.prompt is required, it is the text to translate")
|
||||
}
|
||||
if target == "" {
|
||||
return "", status.Error(codes.InvalidArgument,
|
||||
"nemo-speech-cpp: no target language: set the target_language model option, "+
|
||||
"or prefix the prompt with a [src->tgt] directive")
|
||||
}
|
||||
|
||||
out, err := t.translate([]string{text}, source, target)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
// One text in, one translation out. A call that returned OK with none is a
|
||||
// runtime bug, and the empty string it would hand back reaches the user as a
|
||||
// successful but blank completion with nothing anywhere to say why.
|
||||
if len(out) == 0 {
|
||||
return "", status.Error(codes.Internal, "nemo-speech-cpp: translation produced no result")
|
||||
}
|
||||
return out[0], nil
|
||||
}
|
||||
|
||||
// streamTranslation runs one translation and puts the whole of it on out as a
|
||||
// single chunk.
|
||||
//
|
||||
// That is a limit of the C API and not a shortcut taken here.
|
||||
// nemo_speech_nmt_translate has no token callback and no incremental result: it
|
||||
// returns once the decode has finished, with the completed text. There is
|
||||
// nothing finer to stream, and splitting the finished string into fake chunks
|
||||
// would imitate progress that never happened.
|
||||
//
|
||||
// out is not closed here. PredictStream owns it, and closing it in one of two
|
||||
// places depending on how far the request got is how a stream ends up
|
||||
// half-closed.
|
||||
func streamTranslation(t translator, source, target, text string, out chan<- string) error {
|
||||
translated, err := translateText(t, source, target, text)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
out <- translated
|
||||
return nil
|
||||
}
|
||||
|
||||
// resolveRequest is the shared front half of both RPCs: it names what it is
|
||||
// dropping and works out the pair and the text.
|
||||
func (n *NemoSpeech) resolveRequest(opts *pb.PredictOptions) (source, target, text string) {
|
||||
// Logged rather than rejected, for the reason the diarization path logs its
|
||||
// own dropped fields: a caller that asked for something extra still wants the
|
||||
// translation it can have, and a request naming a field this backend drops
|
||||
// should say so where an operator can find it.
|
||||
if dropped := unsupportedPredictFields(opts); len(dropped) > 0 {
|
||||
xlog.Warn("nemo-speech-cpp: ignoring request fields this model has no equivalent for",
|
||||
"fields", dropped)
|
||||
}
|
||||
return n.languagePair(opts.GetPrompt())
|
||||
}
|
||||
|
||||
// Predict translates PredictOptions.Prompt.
|
||||
//
|
||||
// The whole body runs inside withEngine, so the family check and the C calls
|
||||
// that trust the handle happen under a single acquisition of engineMu. See the
|
||||
// handoff notes at the bottom of nemospeech.go: Free runs without the backend
|
||||
// lock, so anything that checks the family and then releases the lock before
|
||||
// calling C can have the handle destroyed underneath it.
|
||||
func (n *NemoSpeech) Predict(opts *pb.PredictOptions) (string, error) {
|
||||
var out string
|
||||
if err := n.withEngine(familyNMT, func() error {
|
||||
source, target, text := n.resolveRequest(opts)
|
||||
s, err := translateText(&cTranslator{handle: n.nmt}, source, target, text)
|
||||
out = s
|
||||
return err
|
||||
}); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// PredictStream translates PredictOptions.Prompt and emits the result on
|
||||
// results.
|
||||
//
|
||||
// results is closed on EVERY path, including the family rejection and a
|
||||
// validation failure, and the close is deferred outside withEngine so that a
|
||||
// rejected family still closes it. This is the LEGACY streaming contract, which
|
||||
// is the opposite of PredictStreamRich's: pkg/grpc/server.go:529 calls this and
|
||||
// then blocks on a drain goroutine that only finishes when the channel closes,
|
||||
// so a channel left open does not fail the request, it hangs the RPC and, with
|
||||
// the backend lock still held, every request queued behind it. The rich variant
|
||||
// is the one whose channel the host closes; this one is not.
|
||||
func (n *NemoSpeech) PredictStream(opts *pb.PredictOptions, results chan string) error {
|
||||
defer close(results)
|
||||
|
||||
return n.withEngine(familyNMT, func() error {
|
||||
source, target, text := n.resolveRequest(opts)
|
||||
return streamTranslation(&cTranslator{handle: n.nmt}, source, target, text, results)
|
||||
})
|
||||
}
|
||||
@@ -1,416 +0,0 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"unsafe"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
"google.golang.org/grpc/codes"
|
||||
"google.golang.org/grpc/status"
|
||||
|
||||
pb "github.com/mudler/LocalAI/pkg/grpc/proto"
|
||||
)
|
||||
|
||||
// fakeTranslator scripts the C API's answer and records what it was asked, so
|
||||
// the layer above the ABI (pair resolution, validation, the single-element text
|
||||
// array) has a test at all. No Riva-Translate GGUF is small enough to keep in
|
||||
// the tree, and this pretends to translate nothing.
|
||||
type fakeTranslator struct {
|
||||
texts []string
|
||||
source, target string
|
||||
calls int
|
||||
|
||||
out []string
|
||||
err error
|
||||
}
|
||||
|
||||
func (f *fakeTranslator) translate(texts []string, source, target string) ([]string, error) {
|
||||
f.calls++
|
||||
f.texts = texts
|
||||
f.source = source
|
||||
f.target = target
|
||||
return f.out, f.err
|
||||
}
|
||||
|
||||
// collectStrings drains ch until it closes and hands back everything it saw.
|
||||
// The host does the same, so a channel this backend forgets to close hangs the
|
||||
// RPC rather than failing it.
|
||||
func collectStrings(ch chan string) chan []string {
|
||||
done := make(chan []string, 1)
|
||||
go func() {
|
||||
var got []string
|
||||
for s := range ch {
|
||||
got = append(got, s)
|
||||
}
|
||||
done <- got
|
||||
}()
|
||||
return done
|
||||
}
|
||||
|
||||
var _ = Describe("languagePair", func() {
|
||||
It("uses the configured pair and returns the prompt unchanged", func() {
|
||||
n := &NemoSpeech{opts: loadOptions{sourceLanguage: "en", targetLanguage: "de"}}
|
||||
src, tgt, text := n.languagePair("hello world")
|
||||
Expect(src).To(Equal("en"))
|
||||
Expect(tgt).To(Equal("de"))
|
||||
Expect(text).To(Equal("hello world"))
|
||||
})
|
||||
|
||||
// nemo_speech_nmt_translate takes explicit languages and has no prompt path,
|
||||
// so an inline directive is the only way a caller can pick a pair per request.
|
||||
It("honours an inline pair directive and strips it from the text", func() {
|
||||
n := &NemoSpeech{opts: loadOptions{sourceLanguage: "en", targetLanguage: "de"}}
|
||||
src, tgt, text := n.languagePair("[en->fr] hello world")
|
||||
Expect(src).To(Equal("en"))
|
||||
Expect(tgt).To(Equal("fr"))
|
||||
Expect(text).To(Equal("hello world"))
|
||||
})
|
||||
|
||||
// The directive has to be gone from what reaches the model: the runtime
|
||||
// wraps the text in a chat template (src/nmt/langpairs.cc build_prompt), so
|
||||
// a leftover "[en->fr]" would be translated along with the sentence.
|
||||
It("leaves no trace of the directive in the translated text", func() {
|
||||
n := &NemoSpeech{opts: loadOptions{targetLanguage: "de"}}
|
||||
_, _, text := n.languagePair("[en->fr] hello world")
|
||||
Expect(text).ToNot(ContainSubstring("["))
|
||||
Expect(text).ToNot(ContainSubstring("->"))
|
||||
Expect(text).ToNot(ContainSubstring("fr"))
|
||||
Expect(text).To(Equal("hello world"))
|
||||
})
|
||||
|
||||
It("leaves an unparseable directive in the text", func() {
|
||||
n := &NemoSpeech{opts: loadOptions{sourceLanguage: "en", targetLanguage: "de"}}
|
||||
src, tgt, text := n.languagePair("[not a directive] hi")
|
||||
Expect(src).To(Equal("en"))
|
||||
Expect(tgt).To(Equal("de"))
|
||||
Expect(text).To(Equal("[not a directive] hi"))
|
||||
})
|
||||
|
||||
It("trims surrounding whitespace from the text", func() {
|
||||
n := &NemoSpeech{opts: loadOptions{sourceLanguage: "en", targetLanguage: "de"}}
|
||||
_, _, text := n.languagePair(" hello ")
|
||||
Expect(text).To(Equal("hello"))
|
||||
})
|
||||
|
||||
// The model's own tags carry region subtags (src/nmt/langpairs.cc: en-zh-cn,
|
||||
// pt-br, es-us), so a directive that only accepted bare two-letter codes
|
||||
// could not name half the pairs the runtime supports.
|
||||
It("accepts a regional code on either side", func() {
|
||||
n := &NemoSpeech{}
|
||||
src, tgt, text := n.languagePair("[pt-br->en] ola")
|
||||
Expect(src).To(Equal("pt-br"))
|
||||
Expect(tgt).To(Equal("en"))
|
||||
Expect(text).To(Equal("ola"))
|
||||
|
||||
src, tgt, _ = n.languagePair("[en->zh-cn] hi")
|
||||
Expect(src).To(Equal("en"))
|
||||
Expect(tgt).To(Equal("zh-cn"))
|
||||
})
|
||||
|
||||
// resolve_tag accepts a ready pair tag in one field with the other empty, so
|
||||
// a directive that names only one side must keep the configured value for the
|
||||
// other rather than blanking it.
|
||||
It("keeps the configured code for a side the directive omits", func() {
|
||||
n := &NemoSpeech{opts: loadOptions{sourceLanguage: "en", targetLanguage: "de"}}
|
||||
src, tgt, text := n.languagePair("[->fr] hello")
|
||||
Expect(src).To(Equal("en"))
|
||||
Expect(tgt).To(Equal("fr"))
|
||||
Expect(text).To(Equal("hello"))
|
||||
|
||||
src, tgt, _ = n.languagePair("[fr->] hello")
|
||||
Expect(src).To(Equal("fr"))
|
||||
Expect(tgt).To(Equal("de"))
|
||||
})
|
||||
|
||||
// resolve_tag (src/nmt/langpairs.cc:167-172) accepts a READY pair tag in one
|
||||
// field with the other empty, and the model's own tags run to three segments
|
||||
// (en-zh-cn, en-zh-tw, en-es-us, en-pt-br). That single-field three-segment
|
||||
// form is the case a two-segment pattern cannot express: it does not merely
|
||||
// mis-split the tag, it fails to match the directive at all, so the whole
|
||||
// bracket survives into the text and is handed to the model as something to
|
||||
// translate.
|
||||
//
|
||||
// Two-segment codes like pt-br and zh-cn are NOT this case; they parse either
|
||||
// way.
|
||||
It("accepts a three-segment pair tag given in one side of the directive", func() {
|
||||
n := &NemoSpeech{opts: loadOptions{targetLanguage: "de"}}
|
||||
src, tgt, text := n.languagePair("[->en-zh-cn] hi")
|
||||
Expect(src).To(BeEmpty())
|
||||
Expect(tgt).To(Equal("en-zh-cn"))
|
||||
Expect(text).To(Equal("hi"))
|
||||
|
||||
src, tgt, text = n.languagePair("[pt-br-en->] hola")
|
||||
Expect(src).To(Equal("pt-br-en"))
|
||||
Expect(tgt).To(Equal("de"))
|
||||
Expect(text).To(Equal("hola"))
|
||||
})
|
||||
|
||||
It("does not treat a bracketed sentence as a directive", func() {
|
||||
n := &NemoSpeech{opts: loadOptions{targetLanguage: "de"}}
|
||||
_, _, text := n.languagePair("[see figure 1] the cat sat")
|
||||
Expect(text).To(Equal("[see figure 1] the cat sat"))
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("nmtTranslatorConfig", func() {
|
||||
// Backend, Model, Generation and Pool are four adjacent same-typed pointers.
|
||||
// Transposing two of them changes neither the struct's size nor any field's
|
||||
// offset, so the layout assertions in abi_test.go cannot see it, and the
|
||||
// failure it produces is the runtime reading the backend config as the model
|
||||
// config. Distinct sentinels are the only thing that catches it.
|
||||
It("wires each pointer into its own field", func() {
|
||||
cfg := nmtTranslatorConfig(0xB, 0xD)
|
||||
Expect(cfg.Backend).To(Equal(uintptr(0xB)))
|
||||
Expect(cfg.Model).To(Equal(uintptr(0xD)))
|
||||
})
|
||||
|
||||
// NULL is what nmt.h documents as "library defaults" for a subsystem config,
|
||||
// and this backend has no option to fill either of them from.
|
||||
It("leaves the generation and pool configs null", func() {
|
||||
cfg := nmtTranslatorConfig(0xB, 0xD)
|
||||
Expect(cfg.Generation).To(BeZero())
|
||||
Expect(cfg.Pool).To(BeZero())
|
||||
})
|
||||
|
||||
// The runtime decides a field is present with HAS_FIELD, which tests the
|
||||
// caller's size against offsetof + sizeof (src/nmt/c_api.cpp), so a config
|
||||
// sent with Size 0 has every field ignored and the model loads from a path
|
||||
// it was never given.
|
||||
It("declares its own size", func() {
|
||||
Expect(nmtTranslatorConfig(0xB, 0xD).Size).To(Equal(unsafe.Sizeof(cNMTTranslatorConfig{})))
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("nmtTexts", func() {
|
||||
It("produces one non-null pointer per text", func() {
|
||||
ptrs, release, err := nmtTexts([]string{"one", "two", "three"})
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
defer release()
|
||||
|
||||
Expect(ptrs).To(HaveLen(3))
|
||||
for i, p := range ptrs {
|
||||
Expect(p).ToNot(BeZero(), "texts[%d] must not be NULL", i)
|
||||
}
|
||||
// Distinct addresses: one buffer reused for every element would make the
|
||||
// runtime translate the last text three times.
|
||||
Expect(ptrs[0]).ToNot(Equal(ptrs[1]))
|
||||
Expect(ptrs[1]).ToNot(Equal(ptrs[2]))
|
||||
})
|
||||
|
||||
// cstr maps "" to NULL and src/nmt/c_api.cpp maps a NULL element back to "",
|
||||
// so a blank text would be answered with a translation of nothing instead of
|
||||
// an error.
|
||||
It("refuses an empty element", func() {
|
||||
_, release, err := nmtTexts([]string{"one", ""})
|
||||
Expect(release).ToNot(BeNil())
|
||||
release()
|
||||
Expect(status.Code(err)).To(Equal(codes.InvalidArgument))
|
||||
})
|
||||
|
||||
It("refuses an empty batch", func() {
|
||||
_, release, err := nmtTexts(nil)
|
||||
Expect(release).ToNot(BeNil())
|
||||
release()
|
||||
Expect(status.Code(err)).To(Equal(codes.InvalidArgument))
|
||||
})
|
||||
|
||||
It("survives releasing more than once", func() {
|
||||
_, release, err := nmtTexts([]string{"once"})
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
release()
|
||||
Expect(release).ToNot(Panic())
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("translateText", func() {
|
||||
It("passes the resolved pair and the text through to the runtime", func() {
|
||||
f := &fakeTranslator{out: []string{"hallo welt"}}
|
||||
got, err := translateText(f, "en", "de", "hello world")
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(got).To(Equal("hallo welt"))
|
||||
Expect(f.texts).To(Equal([]string{"hello world"}))
|
||||
Expect(f.source).To(Equal("en"))
|
||||
Expect(f.target).To(Equal("de"))
|
||||
})
|
||||
|
||||
// A single-pair model is configured with target_language alone, and
|
||||
// resolve_tag accepts a ready tag in one field with the other empty.
|
||||
It("allows an empty source language", func() {
|
||||
f := &fakeTranslator{out: []string{"ciao"}}
|
||||
_, err := translateText(f, "", "en-it", "hi")
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(f.source).To(BeEmpty())
|
||||
})
|
||||
|
||||
It("rejects a missing target language and names the option to set", func() {
|
||||
f := &fakeTranslator{}
|
||||
_, err := translateText(f, "en", "", "hello")
|
||||
Expect(status.Code(err)).To(Equal(codes.InvalidArgument))
|
||||
Expect(err.Error()).To(ContainSubstring("target_language"))
|
||||
Expect(f.calls).To(BeZero())
|
||||
})
|
||||
|
||||
It("rejects an empty text without calling the runtime", func() {
|
||||
f := &fakeTranslator{}
|
||||
_, err := translateText(f, "en", "de", "")
|
||||
Expect(status.Code(err)).To(Equal(codes.InvalidArgument))
|
||||
Expect(f.calls).To(BeZero())
|
||||
})
|
||||
|
||||
It("propagates a runtime failure", func() {
|
||||
boom := errors.New("boom")
|
||||
_, err := translateText(&fakeTranslator{err: boom}, "en", "de", "hello")
|
||||
Expect(err).To(MatchError(boom))
|
||||
})
|
||||
|
||||
// A call that returned OK with no translations is a runtime bug, and the
|
||||
// empty string it would hand back reaches the user as a successful but blank
|
||||
// completion with nothing anywhere to say why.
|
||||
It("refuses a result that carries no translation", func() {
|
||||
_, err := translateText(&fakeTranslator{}, "en", "de", "hello")
|
||||
Expect(status.Code(err)).To(Equal(codes.Internal))
|
||||
})
|
||||
|
||||
It("takes the first translation when the runtime returns several", func() {
|
||||
f := &fakeTranslator{out: []string{"first", "second"}}
|
||||
got, err := translateText(f, "en", "de", "hello")
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(got).To(Equal("first"))
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("Predict", func() {
|
||||
It("refuses a model loaded as another family", func() {
|
||||
n := &NemoSpeech{fam: familyASR}
|
||||
out, err := n.Predict(&pb.PredictOptions{Prompt: "hello"})
|
||||
Expect(status.Code(err)).To(Equal(codes.Unimplemented))
|
||||
Expect(out).To(BeEmpty())
|
||||
|
||||
// A lock leaked on the rejection path deadlocks the next request rather
|
||||
// than failing it.
|
||||
Expect(n.engineMu.TryLock()).To(BeTrue())
|
||||
n.engineMu.Unlock()
|
||||
})
|
||||
|
||||
It("refuses an unloaded model", func() {
|
||||
n := &NemoSpeech{}
|
||||
_, err := n.Predict(&pb.PredictOptions{Prompt: "hello"})
|
||||
Expect(status.Code(err)).To(Equal(codes.Unimplemented))
|
||||
})
|
||||
|
||||
// The validation has to happen before anything crosses the ABI: nothing is
|
||||
// loaded here, so a guard placed after the C call would panic on a nil
|
||||
// function value instead of failing the request.
|
||||
It("rejects an empty prompt before it reaches the runtime", func() {
|
||||
n := &NemoSpeech{fam: familyNMT, opts: loadOptions{targetLanguage: "de"}}
|
||||
var err error
|
||||
Expect(func() {
|
||||
_, err = n.Predict(&pb.PredictOptions{})
|
||||
}).ToNot(Panic())
|
||||
Expect(status.Code(err)).To(Equal(codes.InvalidArgument))
|
||||
})
|
||||
|
||||
It("rejects a request with no target language, before it reaches the runtime", func() {
|
||||
n := &NemoSpeech{fam: familyNMT}
|
||||
var err error
|
||||
Expect(func() {
|
||||
_, err = n.Predict(&pb.PredictOptions{Prompt: "hello"})
|
||||
}).ToNot(Panic())
|
||||
Expect(status.Code(err)).To(Equal(codes.InvalidArgument))
|
||||
Expect(err.Error()).To(ContainSubstring("target_language"))
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("PredictStream", func() {
|
||||
// pkg/grpc/server.go drains this channel from a goroutine and then blocks on
|
||||
// that goroutine finishing, so a channel left open does not fail the request,
|
||||
// it hangs the RPC and every request queued behind the backend lock.
|
||||
It("closes the channel when the family does not match", func() {
|
||||
n := &NemoSpeech{fam: familyTTS}
|
||||
ch := make(chan string)
|
||||
done := collectStrings(ch)
|
||||
|
||||
err := n.PredictStream(&pb.PredictOptions{Prompt: "hello"}, ch)
|
||||
Expect(status.Code(err)).To(Equal(codes.Unimplemented))
|
||||
Expect(<-done).To(BeEmpty())
|
||||
Expect(n.engineMu.TryLock()).To(BeTrue())
|
||||
n.engineMu.Unlock()
|
||||
})
|
||||
|
||||
It("closes the channel when the request is rejected", func() {
|
||||
n := &NemoSpeech{fam: familyNMT}
|
||||
ch := make(chan string)
|
||||
done := collectStrings(ch)
|
||||
|
||||
err := n.PredictStream(&pb.PredictOptions{Prompt: "hello"}, ch)
|
||||
Expect(status.Code(err)).To(Equal(codes.InvalidArgument))
|
||||
Expect(<-done).To(BeEmpty())
|
||||
})
|
||||
|
||||
It("closes the channel on an unloaded model", func() {
|
||||
n := &NemoSpeech{}
|
||||
ch := make(chan string)
|
||||
done := collectStrings(ch)
|
||||
|
||||
Expect(n.PredictStream(&pb.PredictOptions{Prompt: "hi"}, ch)).ToNot(Succeed())
|
||||
Expect(<-done).To(BeEmpty())
|
||||
})
|
||||
|
||||
// The C API has no token callback, so the whole translation is one chunk.
|
||||
// The seam is the only place that can be asserted without a model.
|
||||
It("emits the whole translation as a single chunk", func() {
|
||||
f := &fakeTranslator{out: []string{"hallo welt"}}
|
||||
ch := make(chan string)
|
||||
done := collectStrings(ch)
|
||||
|
||||
Expect(streamTranslation(f, "en", "de", "hello world", ch)).To(Succeed())
|
||||
close(ch)
|
||||
Expect(<-done).To(Equal([]string{"hallo welt"}))
|
||||
})
|
||||
|
||||
It("emits nothing when the translation fails", func() {
|
||||
f := &fakeTranslator{err: errors.New("boom")}
|
||||
ch := make(chan string)
|
||||
done := collectStrings(ch)
|
||||
|
||||
Expect(streamTranslation(f, "en", "de", "hello", ch)).ToNot(Succeed())
|
||||
close(ch)
|
||||
Expect(<-done).To(BeEmpty())
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("unsupportedPredictFields", func() {
|
||||
It("names nothing for a plain translation request", func() {
|
||||
Expect(unsupportedPredictFields(&pb.PredictOptions{Prompt: "hello"})).To(BeEmpty())
|
||||
})
|
||||
|
||||
// The sampling knobs are deliberately absent from this list: LocalAI fills
|
||||
// them in from the model config on every request, so warning about them
|
||||
// would log on every translation and say nothing.
|
||||
It("stays quiet about sampling parameters the runtime has no field for", func() {
|
||||
Expect(unsupportedPredictFields(&pb.PredictOptions{
|
||||
Prompt: "hello",
|
||||
Temperature: 0.7,
|
||||
TopP: 0.9,
|
||||
TopK: 40,
|
||||
Seed: 42,
|
||||
Tokens: 256,
|
||||
})).To(BeEmpty())
|
||||
})
|
||||
|
||||
It("names the asks the C API cannot serve at all", func() {
|
||||
got := unsupportedPredictFields(&pb.PredictOptions{
|
||||
Prompt: "hello",
|
||||
Grammar: "root ::= x",
|
||||
Tools: `[{"type":"function"}]`,
|
||||
Images: []string{"a.png"},
|
||||
Videos: []string{"a.mp4"},
|
||||
Audios: []string{"a.wav"},
|
||||
NegativePrompt: "no",
|
||||
Logprobs: 3,
|
||||
})
|
||||
Expect(got).To(ConsistOf("grammar", "tools", "images", "videos", "audios",
|
||||
"negative_prompt", "logprobs"))
|
||||
})
|
||||
})
|
||||
@@ -1,100 +0,0 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/mudler/xlog"
|
||||
)
|
||||
|
||||
// loadOptions holds the parsed model-level options. Path fields are resolved
|
||||
// against ModelOptions.ModelPath at parse time so every consumer sees an
|
||||
// absolute path.
|
||||
type loadOptions struct {
|
||||
// ASR
|
||||
vadModel string
|
||||
pncModel string
|
||||
diarModel string
|
||||
itnDir string
|
||||
languageCode string
|
||||
|
||||
// TTS
|
||||
codecModel string
|
||||
tokenizerDir string
|
||||
tnDir string
|
||||
|
||||
// NMT
|
||||
sourceLanguage string
|
||||
targetLanguage string
|
||||
|
||||
// gpu is the device index passed to the runtime's backend config.
|
||||
// -1 selects CPU, matching the C API's own sentinel.
|
||||
gpu int32
|
||||
}
|
||||
|
||||
// splitOption splits on the FIRST colon so values may themselves contain one.
|
||||
func splitOption(o string) (key, value string, ok bool) {
|
||||
i := strings.Index(o, ":")
|
||||
if i < 0 {
|
||||
return "", "", false
|
||||
}
|
||||
return strings.TrimSpace(o[:i]), strings.TrimSpace(o[i+1:]), true
|
||||
}
|
||||
|
||||
// resolve makes a relative asset path absolute against the models directory.
|
||||
// Empty stays empty so callers can test for "unset".
|
||||
func resolve(base, p string) string {
|
||||
if p == "" || filepath.IsAbs(p) {
|
||||
return p
|
||||
}
|
||||
return filepath.Join(base, p)
|
||||
}
|
||||
|
||||
// parseOptions reads the backend "key:value" option slice. Unknown keys are
|
||||
// ignored rather than rejected, so a config written for a newer backend still
|
||||
// loads on an older one.
|
||||
func parseOptions(opts []string, modelPath string) loadOptions {
|
||||
o := loadOptions{gpu: -1}
|
||||
for _, oo := range opts {
|
||||
key, value, ok := splitOption(oo)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
switch key {
|
||||
case "vad_model":
|
||||
o.vadModel = resolve(modelPath, value)
|
||||
case "pnc_model":
|
||||
o.pncModel = resolve(modelPath, value)
|
||||
case "diar_model":
|
||||
o.diarModel = resolve(modelPath, value)
|
||||
case "itn_dir":
|
||||
o.itnDir = resolve(modelPath, value)
|
||||
case "language_code":
|
||||
o.languageCode = value
|
||||
case "codec_model":
|
||||
o.codecModel = resolve(modelPath, value)
|
||||
case "tokenizer_dir":
|
||||
o.tokenizerDir = resolve(modelPath, value)
|
||||
case "tn_dir":
|
||||
o.tnDir = resolve(modelPath, value)
|
||||
case "source_language":
|
||||
o.sourceLanguage = value
|
||||
case "target_language":
|
||||
o.targetLanguage = value
|
||||
case "gpu":
|
||||
// An unknown key is ignored for forward compatibility, but a known key
|
||||
// with an unparseable value is a typo, and this one fails expensively:
|
||||
// the model still loads and still produces correct output, just on CPU
|
||||
// and far slower, with nothing anywhere to say why.
|
||||
n, err := strconv.ParseInt(value, 10, 32)
|
||||
if err != nil {
|
||||
xlog.Warn("nemo-speech-cpp: ignoring unparseable option value, falling back to CPU",
|
||||
"key", key, "value", value)
|
||||
continue
|
||||
}
|
||||
o.gpu = int32(n)
|
||||
}
|
||||
}
|
||||
return o
|
||||
}
|
||||
@@ -1,70 +0,0 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
)
|
||||
|
||||
var _ = Describe("parseOptions", func() {
|
||||
It("parses every known key", func() {
|
||||
o := parseOptions([]string{
|
||||
"vad_model:silero.gguf",
|
||||
"pnc_model:pnc.gguf",
|
||||
"diar_model:sortformer.gguf",
|
||||
"itn_dir:tn_configs",
|
||||
"language_code:es-ES",
|
||||
"codec_model:nanocodec.gguf",
|
||||
"tokenizer_dir:extracted",
|
||||
"tn_dir:tn",
|
||||
"source_language:en",
|
||||
"target_language:de",
|
||||
}, "/models")
|
||||
|
||||
Expect(o.vadModel).To(Equal("/models/silero.gguf"))
|
||||
Expect(o.pncModel).To(Equal("/models/pnc.gguf"))
|
||||
Expect(o.diarModel).To(Equal("/models/sortformer.gguf"))
|
||||
Expect(o.itnDir).To(Equal("/models/tn_configs"))
|
||||
Expect(o.languageCode).To(Equal("es-ES"))
|
||||
Expect(o.codecModel).To(Equal("/models/nanocodec.gguf"))
|
||||
Expect(o.tokenizerDir).To(Equal("/models/extracted"))
|
||||
Expect(o.tnDir).To(Equal("/models/tn"))
|
||||
Expect(o.sourceLanguage).To(Equal("en"))
|
||||
Expect(o.targetLanguage).To(Equal("de"))
|
||||
})
|
||||
|
||||
It("leaves absolute paths untouched", func() {
|
||||
o := parseOptions([]string{"vad_model:/abs/silero.gguf"}, "/models")
|
||||
Expect(o.vadModel).To(Equal("/abs/silero.gguf"))
|
||||
})
|
||||
|
||||
It("ignores unknown keys and entries without a separator", func() {
|
||||
o := parseOptions([]string{"nonsense", "unknown_key:value"}, "/models")
|
||||
Expect(o).To(Equal(loadOptions{gpu: -1}))
|
||||
})
|
||||
|
||||
It("trims whitespace around keys and values", func() {
|
||||
o := parseOptions([]string{" language_code : en-US "}, "/models")
|
||||
Expect(o.languageCode).To(Equal("en-US"))
|
||||
})
|
||||
|
||||
It("keeps a value containing a colon intact", func() {
|
||||
// URIs must survive the split on the FIRST colon.
|
||||
o := parseOptions([]string{"tokenizer_dir:/a/b:c"}, "/models")
|
||||
Expect(o.tokenizerDir).To(Equal("/a/b:c"))
|
||||
})
|
||||
|
||||
It("leaves an empty value empty so callers can detect unset", func() {
|
||||
o := parseOptions([]string{"vad_model:"}, "/models")
|
||||
Expect(o.vadModel).To(BeEmpty())
|
||||
})
|
||||
|
||||
It("defaults gpu to -1 meaning CPU", func() {
|
||||
o := parseOptions(nil, "/models")
|
||||
Expect(o.gpu).To(Equal(int32(-1)))
|
||||
})
|
||||
|
||||
It("parses an explicit gpu index", func() {
|
||||
o := parseOptions([]string{"gpu:0"}, "/models")
|
||||
Expect(o.gpu).To(Equal(int32(0)))
|
||||
})
|
||||
})
|
||||
@@ -1,161 +0,0 @@
|
||||
#!/bin/bash
|
||||
#
|
||||
# Bundle the nemo-speech-cpp-grpc binary, the five nemo_speech shared objects,
|
||||
# the text-normalization stack on a WITH_NORM build, the core runtime libs
|
||||
# (libc/libstdc++/libgomp + ld.so) and the GPU runtime for the active BUILD_TYPE
|
||||
# so the package is self-contained. Mirrors backend/go/whisper/package.sh;
|
||||
# run.sh routes the (CGO_ENABLED=0) binary through lib/ld.so so the packaged
|
||||
# libc is used instead of the host's.
|
||||
#
|
||||
# Five, not three: ASR and NMT each ship a thin _c ABI shim plus the
|
||||
# implementation DSO it depends on, while TTS ships one object with no _c
|
||||
# suffix at all.
|
||||
|
||||
set -e
|
||||
|
||||
CURDIR=$(dirname "$(realpath "$0")")
|
||||
REPO_ROOT="${CURDIR}/../../.."
|
||||
|
||||
mkdir -p "$CURDIR/package/lib"
|
||||
|
||||
cp -avf "$CURDIR/nemo-speech-cpp-grpc" "$CURDIR/package/"
|
||||
cp -avf "$CURDIR/run.sh" "$CURDIR/package/"
|
||||
|
||||
# The runtime ships three C ABI shared objects, not one. All three are
|
||||
# required: main.go dlopens them eagerly, so a package missing any of them
|
||||
# fails at startup. ASR and NMT expose the ABI through a dedicated _c library;
|
||||
# TTS compiles its c_api into libnemo_speech_tts itself and has no _c variant,
|
||||
# hence the asymmetric list. purego.Dlopen resolves them via the
|
||||
# NEMO_SPEECH_*_LIBRARY paths that run.sh points at lib/.
|
||||
#
|
||||
# libnemo_speech_asr and libnemo_speech_nmt are in the list because the matching
|
||||
# _c shims carry a DT_NEEDED on them: dlopen of the shim fails without the
|
||||
# implementation DSO alongside it.
|
||||
for lib in libnemo_speech_asr_c libnemo_speech_asr libnemo_speech_tts libnemo_speech_nmt_c libnemo_speech_nmt; do
|
||||
cp -avf "$CURDIR"/${lib}.so* "$CURDIR/package/lib/" 2>/dev/null || true
|
||||
cp -avf "$CURDIR"/${lib}*.dylib "$CURDIR/package/lib/" 2>/dev/null || true
|
||||
if ! ls "$CURDIR"/package/lib/${lib}.* >/dev/null 2>&1; then
|
||||
echo "ERROR: ${lib} shared library not found in $CURDIR, run 'make' first" >&2
|
||||
exit 1
|
||||
fi
|
||||
done
|
||||
|
||||
# Text normalization (WITH_NORM=ON, Linux only) links Sparrowhawk and OpenFST
|
||||
# into libnemo_speech_asr.so. Those live in a project-local prefix that the
|
||||
# Makefile stages here, so anything staged that is not a nemo_speech object is
|
||||
# part of that stack. Absent on a WITH_NORM=OFF build, which is why this is a
|
||||
# glob that tolerates no matches rather than a required list.
|
||||
shopt -s nullglob
|
||||
for so in "$CURDIR"/*.so "$CURDIR"/*.so.* "$CURDIR"/*.dylib; do
|
||||
case "$(basename "$so")" in
|
||||
libnemo_speech_*) continue ;;
|
||||
esac
|
||||
cp -avf "$so" "$CURDIR/package/lib/"
|
||||
done
|
||||
shopt -u nullglob
|
||||
|
||||
# Detect architecture and copy the core runtime libs the shared objects link
|
||||
# against, plus the matching dynamic loader as lib/ld.so.
|
||||
source "$CURDIR/../../../scripts/build/package-system-libs.sh" "$CURDIR/package/lib" ""
|
||||
|
||||
# Dependency-closure guard.
|
||||
#
|
||||
# The lists above are maintained by hand, and the WITH_NORM build in particular
|
||||
# pulls in transitive dependencies nobody enumerated: Sparrowhawk drags in
|
||||
# protobuf, re2 and absl, none of which package-system-libs.sh provides. Rather
|
||||
# than hard-code that set, walk the DT_NEEDED entries of everything staged and
|
||||
# copy whatever is still unresolved. On a WITH_NORM=OFF build the closure is
|
||||
# already complete, so this copies nothing.
|
||||
#
|
||||
# Skipped deliberately: the core runtime set that package-system-libs.sh owns,
|
||||
# and the GPU stack that package-gpu-libs.sh owns.
|
||||
shopt -s nullglob
|
||||
staged_libs=("$CURDIR"/package/lib/*.so*)
|
||||
shopt -u nullglob
|
||||
|
||||
if [ "$(uname)" != "Darwin" ] && [ "${#staged_libs[@]}" -gt 0 ]; then
|
||||
# No silent skip. If the closure cannot be checked, the package cannot be
|
||||
# shown to be complete, and shipping an unverified one is the failure this
|
||||
# guard exists to prevent.
|
||||
if command -v readelf >/dev/null 2>&1; then
|
||||
read_needed() { readelf -d "$1" 2>/dev/null | sed -n 's/.*(NEEDED).*\[\(.*\)\]/\1/p'; }
|
||||
elif command -v objdump >/dev/null 2>&1; then
|
||||
read_needed() { objdump -p "$1" 2>/dev/null | awk '$1 == "NEEDED" { print $2 }'; }
|
||||
else
|
||||
echo "ERROR: neither readelf nor objdump is available, so the dependency" >&2
|
||||
echo " closure of ${#staged_libs[@]} staged libraries cannot be verified." >&2
|
||||
echo " Install binutils in the build image; refusing to ship an" >&2
|
||||
echo " unverified package." >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
is_provided() {
|
||||
case "$1" in
|
||||
ld-linux*|libc.so.6|libstdc++.so.6|libgcc_s.so.1|libm.so.6|libgomp.so.1) return 0 ;;
|
||||
libdl.so.2|librt.so.1|libpthread.so.0) return 0 ;;
|
||||
libcuda*|libcudart*|libcublas*|libcublasLt*|libnvrtc*|libnvidia*) return 0 ;;
|
||||
libamdhip*|libhsa*|librocm*|libze_*|libOpenCL*|libvulkan*) return 0 ;;
|
||||
esac
|
||||
[ -e "$CURDIR/package/lib/$1" ]
|
||||
}
|
||||
|
||||
# Walk until the staged set stops growing. The glob below expands once per
|
||||
# pass, so each pass advances the closure by exactly one dependency level;
|
||||
# a copied library can itself pull in new dependencies.
|
||||
#
|
||||
# CLOSURE_MAX_PASSES is a runaway guard, not a depth limit. Exhausting it
|
||||
# means the walk never converged and the package is therefore incomplete,
|
||||
# which has to fail the build: a fixed pass count that just falls out of the
|
||||
# loop would silently ship a package missing its deepest libraries, and
|
||||
# libnemo_speech_asr -> sparrowhawk -> protobuf -> absl already runs several
|
||||
# levels deep.
|
||||
CLOSURE_MAX_PASSES="${CLOSURE_MAX_PASSES:-64}"
|
||||
converged=0
|
||||
for (( pass=1; pass<=CLOSURE_MAX_PASSES; pass++ )); do
|
||||
missing=0
|
||||
for so in "$CURDIR"/package/lib/*.so*; do
|
||||
[ -f "$so" ] || continue
|
||||
for need in $(read_needed "$so"); do
|
||||
# Written as an if rather than "is_provided && continue" so a
|
||||
# false return cannot trip set -e via the AND-list exit status.
|
||||
if is_provided "$need"; then
|
||||
continue
|
||||
fi
|
||||
# Resolve against the staging dir first, then the system loader.
|
||||
src="$(LD_LIBRARY_PATH="$CURDIR:$CURDIR/package/lib:${LD_LIBRARY_PATH:-}" \
|
||||
ldd "$so" 2>/dev/null | awk -v n="$need" '$1 == n { print $3 }' | head -1)"
|
||||
if [ -z "$src" ] || [ ! -e "$src" ]; then
|
||||
echo "ERROR: $(basename "$so") needs $need and it could not be resolved." >&2
|
||||
echo " The packaged backend would fail to dlopen at runtime." >&2
|
||||
exit 1
|
||||
fi
|
||||
cp -aLvf "$src" "$CURDIR/package/lib/$need"
|
||||
missing=1
|
||||
done
|
||||
done
|
||||
if [ "$missing" -eq 0 ]; then
|
||||
converged=1
|
||||
break
|
||||
fi
|
||||
done
|
||||
|
||||
if [ "$converged" -ne 1 ]; then
|
||||
echo "ERROR: the dependency closure was still growing after" >&2
|
||||
echo " $CLOSURE_MAX_PASSES passes, so the package is incomplete and" >&2
|
||||
echo " would fail to dlopen at runtime. Refusing to ship it." >&2
|
||||
exit 1
|
||||
fi
|
||||
fi
|
||||
|
||||
# Package GPU libraries (CUDA/ROCm/Intel/Vulkan loader + ICDs + drivers)
|
||||
# based on BUILD_TYPE so the backend can reach the GPU without the runtime
|
||||
# base image shipping those drivers.
|
||||
GPU_LIB_SCRIPT="${REPO_ROOT}/scripts/build/package-gpu-libs.sh"
|
||||
if [ -f "$GPU_LIB_SCRIPT" ]; then
|
||||
echo "Packaging GPU libraries for BUILD_TYPE=${BUILD_TYPE:-cpu}..."
|
||||
source "$GPU_LIB_SCRIPT" "$CURDIR/package/lib"
|
||||
package_gpu_libs
|
||||
fi
|
||||
|
||||
echo "Packaging completed successfully"
|
||||
ls -liah "$CURDIR/package/" "$CURDIR/package/lib/"
|
||||
@@ -1,28 +0,0 @@
|
||||
#!/bin/bash
|
||||
set -e
|
||||
|
||||
CURDIR=$(dirname "$(realpath "$0")")
|
||||
|
||||
# The runtime splits its C ABI across three shared objects, so each gets its
|
||||
# own override variable. main.go reads exactly these names.
|
||||
if [ "$(uname)" = "Darwin" ]; then
|
||||
export DYLD_LIBRARY_PATH="$CURDIR/lib:"$CURDIR":${DYLD_LIBRARY_PATH:-}"
|
||||
export NEMO_SPEECH_ASR_LIBRARY="$CURDIR/lib/libnemo_speech_asr_c.dylib"
|
||||
export NEMO_SPEECH_TTS_LIBRARY="$CURDIR/lib/libnemo_speech_tts.dylib"
|
||||
export NEMO_SPEECH_NMT_LIBRARY="$CURDIR/lib/libnemo_speech_nmt_c.dylib"
|
||||
else
|
||||
export LD_LIBRARY_PATH="$CURDIR/lib:"$CURDIR":${LD_LIBRARY_PATH:-}"
|
||||
export NEMO_SPEECH_ASR_LIBRARY="$CURDIR/lib/libnemo_speech_asr_c.so"
|
||||
export NEMO_SPEECH_TTS_LIBRARY="$CURDIR/lib/libnemo_speech_tts.so"
|
||||
export NEMO_SPEECH_NMT_LIBRARY="$CURDIR/lib/libnemo_speech_nmt_c.so"
|
||||
fi
|
||||
|
||||
# If a self-contained ld.so was packaged, route through it so the
|
||||
# packaged libc / libstdc++ are used instead of the host's (matches the
|
||||
# whisper backend's runtime layout). Linux only.
|
||||
if [ -f "$CURDIR/lib/ld.so" ]; then
|
||||
echo "Using lib/ld.so"
|
||||
exec "$CURDIR/lib/ld.so" "$CURDIR/nemo-speech-cpp-grpc" "$@"
|
||||
fi
|
||||
|
||||
exec "$CURDIR/nemo-speech-cpp-grpc" "$@"
|
||||
@@ -1,102 +0,0 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"google.golang.org/grpc/codes"
|
||||
"google.golang.org/grpc/status"
|
||||
)
|
||||
|
||||
// The status values every C entry point in this backend returns.
|
||||
//
|
||||
// There is no single C enum to mirror. asr.h:43-49, tts.h:38-44 and nmt.h:40-45
|
||||
// each declare their own, and diar.h has none of its own at all: it includes
|
||||
// asr.h and types every diarization function as nemo_speech_asr_status
|
||||
// (diar.h:18). The names the three surfaces share carry the same numbers:
|
||||
//
|
||||
// value asr.h tts.h nmt.h
|
||||
// 0 NEMO_SPEECH_ASR_OK NEMO_SPEECH_TTS_OK NEMO_SPEECH_NMT_OK
|
||||
// 1 NEMO_SPEECH_ASR_ERROR_INVALID_... NEMO_SPEECH_TTS_ERROR_INVALID_... NEMO_SPEECH_NMT_ERROR_INVALID_...
|
||||
// 2 NEMO_SPEECH_ASR_ERROR_OUT_OF_MEM. NEMO_SPEECH_TTS_ERROR_OUT_OF_MEM. NEMO_SPEECH_NMT_ERROR_OUT_OF_MEM.
|
||||
// 3 NEMO_SPEECH_ASR_ERROR_RUNTIME NEMO_SPEECH_TTS_ERROR_RUNTIME NEMO_SPEECH_NMT_ERROR_RUNTIME
|
||||
// 4 NEMO_SPEECH_ASR_ERROR_CANCELLED NEMO_SPEECH_TTS_ERROR_CANCELLED (not declared)
|
||||
//
|
||||
// The one divergence is 4, and it is an absence rather than a disagreement. ASR
|
||||
// and TTS both drive a consumer callback that can ask for the work to stop, and
|
||||
// cancellation is what they report when it does; nemo_speech_nmt_translate takes
|
||||
// no callback and returns only when the decode has finished, so the NMT surface
|
||||
// has no cancellation to name. That is why one table can serve all three: 4 is
|
||||
// not some other NMT status that would be mislabelled, it is a value the NMT
|
||||
// surface never produces.
|
||||
//
|
||||
// Recheck this table after an upstream pin bump. A status added to one header
|
||||
// and not the others is exactly the shape of change that would break the single
|
||||
// mapping, and nothing in the build or the linker can see it: purego binds by
|
||||
// name, and the return value is a bare int32 on the Go side.
|
||||
const (
|
||||
statusOK int32 = 0
|
||||
statusInvalidArgument int32 = 1
|
||||
statusOutOfMemory int32 = 2
|
||||
statusRuntime int32 = 3
|
||||
statusCancelled int32 = 4
|
||||
)
|
||||
|
||||
// statusCode maps a C status onto the gRPC code the caller should be told.
|
||||
//
|
||||
// What the mapping is really carrying is whose mistake the failure was.
|
||||
// INVALID_ARGUMENT is what every guard in src/{asr,tts,nmt}/c_api.cpp returns
|
||||
// for a std::invalid_argument from the runtime, and the things that throw it are
|
||||
// requests: an unknown voice_name (src/tts/synthesizer.cpp), an unsupported
|
||||
// language pair (src/nmt/translator.cpp), an out-of-range sample rate. Reporting
|
||||
// those as Internal turns a 400 into a 500 and sends the user hunting for a
|
||||
// broken model or a broken backend instead of fixing the request.
|
||||
//
|
||||
// OUT_OF_MEMORY is a resource limit rather than a defect, which is what
|
||||
// ResourceExhausted means, and it is the one failure a client can sensibly
|
||||
// retry later or retry smaller. CANCELLED is the consumer having stopped
|
||||
// listening, which is not a failure of this backend at all: the streaming sinks
|
||||
// return false once their client is gone (see ttsDeliverPCM), and the runtime
|
||||
// turns that into status 4.
|
||||
//
|
||||
// RUNTIME, and anything a future pin adds that this table has not been taught,
|
||||
// stay Internal. An unrecognised status is precisely the case where the backend
|
||||
// does not know whose fault it was, and Internal is the honest answer.
|
||||
func statusCode(st int32) codes.Code {
|
||||
switch st {
|
||||
case statusOK:
|
||||
return codes.OK
|
||||
case statusInvalidArgument:
|
||||
return codes.InvalidArgument
|
||||
case statusOutOfMemory:
|
||||
return codes.ResourceExhausted
|
||||
case statusCancelled:
|
||||
return codes.Canceled
|
||||
case statusRuntime:
|
||||
// Named rather than folded into the default so this switch reads as the
|
||||
// whole enum. A status the table has never heard of is a different thing
|
||||
// from a runtime error even though both answer Internal, and a reader
|
||||
// checking the mapping against the headers should not have to work out
|
||||
// which arm RUNTIME lands in.
|
||||
return codes.Internal
|
||||
default:
|
||||
return codes.Internal
|
||||
}
|
||||
}
|
||||
|
||||
// statusErrorf builds the gRPC error for a failed C call.
|
||||
//
|
||||
// Every C call site in this backend goes through this rather than through
|
||||
// status.Errorf directly, and that is the whole point of it existing: the
|
||||
// mapping used to be written out at exactly one of sixteen call sites, so the
|
||||
// same backend answered an unsupported language pair with InvalidArgument and an
|
||||
// unknown TTS voice, which is the same class of caller mistake against the same
|
||||
// process, with Internal.
|
||||
//
|
||||
// The OK guard is not defensive noise. status.Errorf(codes.OK, ...) returns a
|
||||
// nil error, so a call site that built its error without first checking the
|
||||
// status would report a hard C failure as a successful request with no
|
||||
// diagnostic anywhere. Returning Internal instead keeps that mistake loud.
|
||||
func statusErrorf(st int32, format string, args ...any) error {
|
||||
if st == statusOK {
|
||||
return status.Errorf(codes.Internal, format, args...)
|
||||
}
|
||||
return status.Errorf(statusCode(st), format, args...)
|
||||
}
|
||||
@@ -1,109 +0,0 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"os"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
"google.golang.org/grpc/codes"
|
||||
"google.golang.org/grpc/status"
|
||||
|
||||
pb "github.com/mudler/LocalAI/pkg/grpc/proto"
|
||||
)
|
||||
|
||||
var _ = Describe("C status mapping", func() {
|
||||
// The whole enum, so a status that quietly moves to a different code is
|
||||
// visible here rather than in a bug report about an HTTP 500. The names on
|
||||
// the left are transcribed from asr.h:43-49, tts.h:38-44 and nmt.h:40-45;
|
||||
// see the table in status.go for how the three surfaces line up.
|
||||
DescribeTable("maps each declared C status onto a gRPC code",
|
||||
func(st int32, want codes.Code) {
|
||||
Expect(statusCode(st)).To(Equal(want))
|
||||
},
|
||||
Entry("OK", statusOK, codes.OK),
|
||||
Entry("INVALID_ARGUMENT", statusInvalidArgument, codes.InvalidArgument),
|
||||
Entry("OUT_OF_MEMORY", statusOutOfMemory, codes.ResourceExhausted),
|
||||
Entry("RUNTIME", statusRuntime, codes.Internal),
|
||||
Entry("CANCELLED", statusCancelled, codes.Canceled),
|
||||
)
|
||||
|
||||
// A pin bump that adds a status this table has never been taught must not
|
||||
// guess. Internal is the honest answer when the backend does not know whose
|
||||
// mistake the failure was.
|
||||
DescribeTable("reports an unknown status as Internal",
|
||||
func(st int32) {
|
||||
Expect(statusCode(st)).To(Equal(codes.Internal))
|
||||
},
|
||||
Entry("one past the last declared value", int32(5)),
|
||||
Entry("far past it", int32(99)),
|
||||
Entry("negative", int32(-1)),
|
||||
)
|
||||
|
||||
It("carries the mapped code and the formatted message into the error", func() {
|
||||
err := statusErrorf(statusInvalidArgument, "nemo-speech-cpp: %s: %d", "synthesize", 7)
|
||||
Expect(status.Code(err)).To(Equal(codes.InvalidArgument))
|
||||
Expect(err.Error()).To(ContainSubstring("nemo-speech-cpp: synthesize: 7"))
|
||||
})
|
||||
|
||||
// status.Errorf(codes.OK, ...) returns nil, so a call site that built its
|
||||
// error without checking the status first would turn a hard C failure into a
|
||||
// silent success with no diagnostic anywhere.
|
||||
It("never returns nil, not even for OK", func() {
|
||||
err := statusErrorf(statusOK, "nemo-speech-cpp: should not happen")
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(status.Code(err)).To(Equal(codes.Internal))
|
||||
})
|
||||
})
|
||||
|
||||
// These drive real C statuses out of the real shared objects, one per family,
|
||||
// rather than asserting the Go mapping against itself.
|
||||
//
|
||||
// A NULL handle is the one failure every surface can be provoked into without a
|
||||
// model: nemo_speech_asr_recognize_f32 and nemo_speech_nmt_translate check the
|
||||
// handle up front, nemo_speech_tts_synthesize_text does the same, and the
|
||||
// diarization stream entry points throw std::invalid_argument for a dead stream,
|
||||
// which src/asr/c_api.cpp's guard maps to the same status. All four are the
|
||||
// caller's mistake, and the point of the exercise is that all four now come back
|
||||
// as InvalidArgument instead of Internal.
|
||||
var _ = Describe("C status mapping at the call sites", func() {
|
||||
BeforeEach(func() {
|
||||
if !librariesPresent() {
|
||||
if requireLibs() {
|
||||
cwd, _ := os.Getwd()
|
||||
Fail("NEMO_SPEECH_REQUIRE_LIBS=1 but the shared libraries are not in " + cwd +
|
||||
": these specs are the ABI defence and must not be skipped." +
|
||||
" Run make -C backend/go/nemo-speech-cpp stage-libs")
|
||||
}
|
||||
Skip("shared libraries not built, run make in backend/go/nemo-speech-cpp")
|
||||
}
|
||||
Expect(openLibraries()).To(Succeed())
|
||||
})
|
||||
|
||||
It("reports an ASR INVALID_ARGUMENT as InvalidArgument", func() {
|
||||
opts := ASRRecognitionOptionsDef()
|
||||
// Non-empty PCM on purpose: recognizeF32 rejects an empty slice itself,
|
||||
// which would prove nothing about what the C side returned.
|
||||
_, err := recognizeF32(0, &opts, []float32{0, 0, 0}, 16000)
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(status.Code(err)).To(Equal(codes.InvalidArgument))
|
||||
})
|
||||
|
||||
It("reports a diarization INVALID_ARGUMENT as InvalidArgument", func() {
|
||||
err := (&cDiarStream{}).finish()
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(status.Code(err)).To(Equal(codes.InvalidArgument))
|
||||
})
|
||||
|
||||
It("reports a TTS INVALID_ARGUMENT as InvalidArgument", func() {
|
||||
s := &cSynthesizer{}
|
||||
err := s.synthesize(&pb.TTSRequest{Text: "hello"}, "en", func([]byte) bool { return true })
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(status.Code(err)).To(Equal(codes.InvalidArgument))
|
||||
})
|
||||
|
||||
It("reports an NMT INVALID_ARGUMENT as InvalidArgument", func() {
|
||||
_, err := (&cTranslator{}).translate([]string{"hello"}, "en", "de")
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(status.Code(err)).To(Equal(codes.InvalidArgument))
|
||||
})
|
||||
})
|
||||
@@ -1,600 +0,0 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"math"
|
||||
"os"
|
||||
"runtime"
|
||||
"strconv"
|
||||
"sync"
|
||||
"unsafe"
|
||||
|
||||
"github.com/ebitengine/purego"
|
||||
laudio "github.com/mudler/LocalAI/pkg/audio"
|
||||
pb "github.com/mudler/LocalAI/pkg/grpc/proto"
|
||||
"github.com/mudler/xlog"
|
||||
"google.golang.org/grpc/codes"
|
||||
"google.golang.org/grpc/status"
|
||||
)
|
||||
|
||||
// The backend-preference enum from include/nemo_speech/tts.h. The C type is an
|
||||
// enum, which this toolchain lays out as int32, so the values are written here
|
||||
// rather than inferred.
|
||||
const (
|
||||
ttsBackendAuto int32 = 0
|
||||
ttsBackendCPU int32 = 1
|
||||
)
|
||||
|
||||
// maxWAVDataBytes is the largest PCM payload a RIFF WAV can describe.
|
||||
//
|
||||
// Both size fields in the header are uint32, so a longer payload would not
|
||||
// merely be unusual, it would wrap and produce a file whose header disagrees
|
||||
// with its contents. At 22.05 kHz mono 16-bit that ceiling is about 27 hours of
|
||||
// speech, so nothing real is being refused.
|
||||
//
|
||||
// Typed int64 rather than left untyped so the comparison below is the same one
|
||||
// on every architecture: an untyped constant this large does not fit in a
|
||||
// 32-bit int and would not compile there at all.
|
||||
const maxWAVDataBytes int64 = math.MaxUint32 - laudio.WAVHeaderSize
|
||||
|
||||
// wavStreamingSize is the placeholder both size fields carry while the total
|
||||
// length is still unknown. Players read it as "stream until the socket closes".
|
||||
const wavStreamingSize = 0xFFFFFFFF
|
||||
|
||||
// ttsSink receives one PCM chunk, already copied into Go memory.
|
||||
//
|
||||
// It returns false to cancel the synthesis in progress: that is the C
|
||||
// callback's only way to stop work early, and the runtime turns it into
|
||||
// NEMO_SPEECH_TTS_ERROR_CANCELLED.
|
||||
type ttsSink func(pcm []byte) bool
|
||||
|
||||
// ttsSinkTable maps the user_data value handed to C back to the Go sink the
|
||||
// chunk belongs to.
|
||||
//
|
||||
// A single "current sink" pointer would be enough for one model, since every
|
||||
// RPC holds that model's engineMu for its whole body. It is not enough for the
|
||||
// process: engineMu is per-NemoSpeech, one backend process can hold several
|
||||
// loaded models, and the callback below is shared by all of them, so two TTS
|
||||
// models synthesizing at once would overwrite each other's sink. The id is what
|
||||
// keeps them apart.
|
||||
//
|
||||
// The id is an integer and never a Go pointer. user_data crosses into C as a
|
||||
// void*, which the collector does not trace, so a Go pointer parked there would
|
||||
// have exactly the lifetime problem cstr documents.
|
||||
type ttsSinkTable struct {
|
||||
mu sync.Mutex
|
||||
next uintptr
|
||||
sinks map[uintptr]ttsSink
|
||||
}
|
||||
|
||||
var ttsSinks = &ttsSinkTable{sinks: map[uintptr]ttsSink{}}
|
||||
|
||||
// register adds sink and returns its id together with the release the caller
|
||||
// MUST defer. Ids start at 1 so a zeroed or stale user_data cannot resolve to
|
||||
// somebody else's sink.
|
||||
func (t *ttsSinkTable) register(sink ttsSink) (uintptr, func()) {
|
||||
t.mu.Lock()
|
||||
defer t.mu.Unlock()
|
||||
|
||||
t.next++
|
||||
id := t.next
|
||||
t.sinks[id] = sink
|
||||
return id, func() {
|
||||
t.mu.Lock()
|
||||
defer t.mu.Unlock()
|
||||
delete(t.sinks, id)
|
||||
}
|
||||
}
|
||||
|
||||
// lookup returns the sink for id, or nil once it has been released.
|
||||
func (t *ttsSinkTable) lookup(id uintptr) ttsSink {
|
||||
t.mu.Lock()
|
||||
defer t.mu.Unlock()
|
||||
|
||||
return t.sinks[id]
|
||||
}
|
||||
|
||||
var (
|
||||
ttsCallbackOnce sync.Once
|
||||
ttsCallbackFn uintptr
|
||||
)
|
||||
|
||||
// ttsPCMCallback returns the C function pointer the runtime drives PCM through,
|
||||
// compiling it on first use.
|
||||
//
|
||||
// Exactly one is ever created per process, and that is a hard requirement
|
||||
// rather than a tidiness argument. purego.NewCallback writes into a fixed table
|
||||
// of maxCB = 2000 entries (purego/syscall_sysv.go) and never releases an entry,
|
||||
// so a callback compiled per request panics the whole backend process with
|
||||
// "purego: the maximum number of callbacks has been reached" on the 2001st
|
||||
// synthesis. Per model load is not safe either: a server that swaps models
|
||||
// reaches the same ceiling, just later and even less predictably. Routing every
|
||||
// synthesis through one callback plus a user_data id is what keeps the count at
|
||||
// one for the life of the process.
|
||||
func ttsPCMCallback() uintptr {
|
||||
ttsCallbackOnce.Do(func() { ttsCallbackFn = purego.NewCallback(ttsDeliverPCM) })
|
||||
return ttsCallbackFn
|
||||
}
|
||||
|
||||
// ttsDeliverPCM is the body of that callback: nemo_speech_tts_pcm_callback,
|
||||
// which the runtime invokes on its own thread for each chunk it produces.
|
||||
//
|
||||
// The bytes are copied rather than aliased. The pointer addresses a std::string
|
||||
// the runtime owns and reuses for the next chunk (src/tts/c_api.cpp
|
||||
// make_callback), so a slice over it would be rewritten under the consumer as
|
||||
// soon as this returns.
|
||||
func ttsDeliverPCM(pcm unsafe.Pointer, nBytes uint64, userData uintptr) bool {
|
||||
// The table's lock is released before the sink runs, which matters because
|
||||
// TTSStream's sink blocks on a channel send until its client drains it.
|
||||
// Holding the lock across that would stall every other model's callback
|
||||
// behind one slow consumer.
|
||||
sink := ttsSinks.lookup(userData)
|
||||
if sink == nil {
|
||||
// The request that registered this sink has already returned, so there
|
||||
// is nowhere to put the audio. false cancels rather than letting the
|
||||
// runtime synthesize to completion into a consumer that stopped
|
||||
// listening.
|
||||
return false
|
||||
}
|
||||
// c_api.cpp filters empty chunks before calling us, so this is belt and
|
||||
// braces: unsafe.Slice on a null pointer is what it protects against.
|
||||
if pcm == nil || nBytes == 0 {
|
||||
return true
|
||||
}
|
||||
|
||||
buf := make([]byte, nBytes)
|
||||
// #nosec G103 -- pcm and nBytes are the C-owned buffer and its length from
|
||||
// one callback invocation, both null/zero-checked above. The slice is read
|
||||
// only, its length is the length the runtime declared for that buffer, and it
|
||||
// is copied into Go memory here and never retained past this return.
|
||||
copy(buf, unsafe.Slice((*byte)(pcm), nBytes))
|
||||
return sink(buf)
|
||||
}
|
||||
|
||||
// synthesizer is the TTS half of the C API, narrowed to what the two RPCs use.
|
||||
//
|
||||
// It is an interface for the same reason asrSession and diarStream are: no
|
||||
// MagpieTTS GGUF is small enough to keep in the tree, so the logic layered on
|
||||
// top of the ABI (validation, the WAV framing, chunk ordering) would otherwise
|
||||
// have no test at all. The seam is at the ABI, not at the model: a fake here
|
||||
// scripts what the C API emits, it does not pretend to synthesize anything.
|
||||
type synthesizer interface {
|
||||
// sampleRate is the rate the PCM chunks arrive at.
|
||||
sampleRate() int32
|
||||
// synthesize maps req onto the runtime's per-request options and runs one
|
||||
// synthesis, handing each chunk to sink as it is produced.
|
||||
synthesize(req *pb.TTSRequest, defaultLanguage string, sink ttsSink) error
|
||||
}
|
||||
|
||||
// cSynthesizer is the real synthesizer, over one nemo_speech_tts_synthesizer.
|
||||
type cSynthesizer struct {
|
||||
handle uintptr
|
||||
}
|
||||
|
||||
func (s *cSynthesizer) sampleRate() int32 { return TTSSampleRate(s.handle) }
|
||||
|
||||
func (s *cSynthesizer) synthesize(req *pb.TTSRequest, defaultLanguage string, sink ttsSink) error {
|
||||
// Started from the runtime's own defaults, not from a zero struct: every
|
||||
// numeric field here is sentinel-sensitive (speaker/seed < 0, steps/top_k
|
||||
// <= 0 all mean "use the synthesizer's value"), and a zeroed struct would
|
||||
// read as speaker 0, seed 0 and zero decoding steps.
|
||||
opts := TTSSynthesisOptionsDefault()
|
||||
|
||||
// A per-request language wins over the model-level default; both may be
|
||||
// empty, which the runtime resolves to the synthesizer's own default.
|
||||
language := req.GetLanguage()
|
||||
if language == "" {
|
||||
language = defaultLanguage
|
||||
}
|
||||
langP, freeLang := cstr(language)
|
||||
defer freeLang()
|
||||
opts.LanguageCode = langP
|
||||
|
||||
speaker, voiceName := resolveSpeaker(req.GetVoice())
|
||||
opts.Speaker = speaker
|
||||
voiceP, freeVoice := cstr(voiceName)
|
||||
defer freeVoice()
|
||||
opts.VoiceName = voiceP
|
||||
|
||||
applySynthesisParams(&opts, req.GetParams())
|
||||
|
||||
id, release := ttsSinks.register(sink)
|
||||
defer release()
|
||||
|
||||
// stats_out is NULL: nemo_speech_tts_synthesis_stats is 300-odd bytes of
|
||||
// timing detail with nowhere to go on either RPC, and the C API documents
|
||||
// NULL as the way to decline it.
|
||||
// #nosec G103 -- opts is a local POD struct borrowed for this call only. Its
|
||||
// two uintptr members (LanguageCode, VoiceName) are cstr allocations pinned
|
||||
// by the defers above, and this entry point is synchronous, so it returns
|
||||
// before those pins are released even though the callbacks run off-thread.
|
||||
st := TTSSynthesizeText(s.handle, unsafe.Pointer(&opts), req.GetText(), ttsPCMCallback(), id, nil)
|
||||
if st != 0 {
|
||||
// An unknown voice_name arrives here as INVALID_ARGUMENT
|
||||
// (src/tts/synthesizer.cpp throws std::invalid_argument, which
|
||||
// src/tts/c_api.cpp's guard maps to it), and a consumer that stopped
|
||||
// reading arrives as CANCELLED. Neither is this backend's failure, so
|
||||
// neither goes out as Internal.
|
||||
return statusErrorf(st, "nemo-speech-cpp: synthesize: %s", TTSLastError())
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// resolveSpeaker splits a request's voice into the two fields the C API has for
|
||||
// it: a speaker index and a voice name.
|
||||
//
|
||||
// nemo_speech_tts_synthesis_options.voice_name is documented as ignored
|
||||
// whenever speaker >= 0, and src/tts/synthesizer.cpp only calls resolve_speaker
|
||||
// when options.speaker is negative, so the two are alternatives and never a
|
||||
// pair. A named voice must therefore leave the index at -1 or the name is
|
||||
// silently dropped.
|
||||
//
|
||||
// The numeric split cannot change what the runtime picks: resolve_speaker parses
|
||||
// a numeric voice_name itself, so anything this function passes through as a
|
||||
// name and that happens to be a number lands on the same speaker anyway. What it
|
||||
// must not do is let a NEGATIVE number through as an index. "-1" is not a
|
||||
// speaker, it is the sentinel for "use the default", and treating it as an index
|
||||
// would turn a request naming an invalid voice into one that quietly synthesizes
|
||||
// in the default voice instead of being rejected.
|
||||
func resolveSpeaker(voice string) (int32, string) {
|
||||
if voice == "" {
|
||||
return -1, ""
|
||||
}
|
||||
if idx, err := strconv.ParseInt(voice, 10, 32); err == nil && idx >= 0 {
|
||||
return int32(idx), ""
|
||||
}
|
||||
return -1, voice
|
||||
}
|
||||
|
||||
// applySynthesisParams maps TTSRequest.params onto the runtime's per-request
|
||||
// options.
|
||||
//
|
||||
// Only the five knobs nemo_speech_tts_synthesis_options actually has are read.
|
||||
// An unset or unparseable value leaves the field alone rather than resetting it:
|
||||
// the struct arrives carrying the runtime's defaults, and params is documented
|
||||
// as "unset leaves the backend's configured defaults".
|
||||
//
|
||||
// The sentinels are the reason each write is guarded rather than unconditional.
|
||||
// src/tts/magpietts/runtime.cpp takes the request's seed only when it is >= 0
|
||||
// and its steps and top_k only when they are > 0, so writing a parsed 0 or a
|
||||
// negative would not merely be ignored, it would erase the option's meaning for
|
||||
// a caller who passed "0" expecting something.
|
||||
//
|
||||
// temperature and cfg_scale each need their override flag set as well. The
|
||||
// runtime reads the float only when the flag is true and otherwise falls back to
|
||||
// the synthesizer's config, so a temperature written without its flag is
|
||||
// silently discarded.
|
||||
func applySynthesisParams(o *cTTSSynthesisOptions, params map[string]string) {
|
||||
if len(params) == 0 {
|
||||
return
|
||||
}
|
||||
if v, ok := parseInt32Param(params["seed"]); ok && v >= 0 {
|
||||
o.Seed = v
|
||||
}
|
||||
if v, ok := parseInt32Param(params["steps"]); ok && v > 0 {
|
||||
o.Steps = v
|
||||
}
|
||||
if v, ok := parseInt32Param(params["top_k"]); ok && v > 0 {
|
||||
o.TopK = v
|
||||
}
|
||||
if v, ok := parseFloat32Param(params["temperature"]); ok {
|
||||
o.Temperature = v
|
||||
o.OverrideTemperature = true
|
||||
}
|
||||
if v, ok := parseFloat32Param(params["cfg_scale"]); ok {
|
||||
o.CFGScale = v
|
||||
o.OverrideCFGScale = true
|
||||
}
|
||||
}
|
||||
|
||||
// parseInt32Param reads one params entry. ok is false for an absent or
|
||||
// unparseable value, which the caller reads as "leave the default".
|
||||
func parseInt32Param(v string) (int32, bool) {
|
||||
if v == "" {
|
||||
return 0, false
|
||||
}
|
||||
n, err := strconv.ParseInt(v, 10, 32)
|
||||
if err != nil {
|
||||
xlog.Warn("nemo-speech-cpp: ignoring unparseable TTS parameter", "value", v)
|
||||
return 0, false
|
||||
}
|
||||
return int32(n), true
|
||||
}
|
||||
|
||||
func parseFloat32Param(v string) (float32, bool) {
|
||||
if v == "" {
|
||||
return 0, false
|
||||
}
|
||||
f, err := strconv.ParseFloat(v, 32)
|
||||
if err != nil {
|
||||
xlog.Warn("nemo-speech-cpp: ignoring unparseable TTS parameter", "value", v)
|
||||
return 0, false
|
||||
}
|
||||
return float32(f), true
|
||||
}
|
||||
|
||||
// ttsModelConfig builds the create-time model config.
|
||||
//
|
||||
// Extracted from loadTTS and asserted field by field because three adjacent
|
||||
// members of nemo_speech_tts_model_config are same-typed paths. Swapping two of
|
||||
// them changes neither the struct's size nor any field's offset, so the layout
|
||||
// assertions in abi_test.go cannot see it, and the failure it produces is the
|
||||
// runtime loading the codec as the acoustic model.
|
||||
//
|
||||
// Every argument is a C pointer from cstr, not a Go string, and the caller owns
|
||||
// the releases. tnDir may be null: text_normalizer_model_dir is optional and an
|
||||
// empty one leaves the text unchanged.
|
||||
func ttsModelConfig(magpieModel, codecModel, tokenizerDir, tnDir uintptr) cTTSModelConfig {
|
||||
return cTTSModelConfig{
|
||||
Size: unsafe.Sizeof(cTTSModelConfig{}),
|
||||
MagpieModel: magpieModel,
|
||||
CodecModel: codecModel,
|
||||
TokenizerModelDir: tokenizerDir,
|
||||
TextNormalizerModelDir: tnDir,
|
||||
}
|
||||
}
|
||||
|
||||
// ttsRuntimeBackend maps the backend's gpu option onto the TTS runtime's
|
||||
// backend preference.
|
||||
//
|
||||
// nemo_speech_tts_runtime_config has no device index at all, only a three-way
|
||||
// AUTO/CPU/CUDA preference, so a gpu option naming a particular device cannot be
|
||||
// honoured and AUTO is the honest answer for it. A negative gpu is different: it
|
||||
// is the option's documented "CPU" across this whole backend (asr.h: "-1 = CPU")
|
||||
// and it is also the default, so it has to pin the preference rather than leave
|
||||
// the runtime free to pick CUDA.
|
||||
func ttsRuntimeBackend(gpu int32) int32 {
|
||||
if gpu < 0 {
|
||||
return ttsBackendCPU
|
||||
}
|
||||
return ttsBackendAuto
|
||||
}
|
||||
|
||||
// loadTTS creates the MagpieTTS synthesizer.
|
||||
//
|
||||
// It runs after discoverTTSAssets, so codecModel and tokenizerDir are already
|
||||
// resolved and non-empty; tnDir stays optional.
|
||||
//
|
||||
// This must not take engineMu: Load is its only caller and already holds it.
|
||||
func (n *NemoSpeech) loadTTS(modelFile string) error {
|
||||
// nemo_speech_tts_create deep-copies every const char* into a std::string
|
||||
// (src/tts/c_api.cpp, via str_or_empty) and keeps no pointer afterwards, so
|
||||
// pinning across the create call is both necessary and sufficient.
|
||||
var pinner runtime.Pinner
|
||||
defer pinner.Unpin()
|
||||
|
||||
magpieP, freeMagpie := cstr(modelFile)
|
||||
defer freeMagpie()
|
||||
codecP, freeCodec := cstr(n.opts.codecModel)
|
||||
defer freeCodec()
|
||||
tokenizerP, freeTokenizer := cstr(n.opts.tokenizerDir)
|
||||
defer freeTokenizer()
|
||||
tnP, freeTN := cstr(n.opts.tnDir)
|
||||
defer freeTN()
|
||||
|
||||
model := ttsModelConfig(magpieP, codecP, tokenizerP, tnP)
|
||||
|
||||
rt := TTSRuntimeConfigDefault()
|
||||
backend := ttsRuntimeBackend(n.opts.gpu)
|
||||
rt.LTBackend = backend
|
||||
rt.SamplingBackend = backend
|
||||
// The codec is a separate graph with its own placement, so a CPU-only
|
||||
// request has to say so here too or it would still try to run on the GPU.
|
||||
rt.CodecCPU = backend == ttsBackendCPU
|
||||
|
||||
langP, freeLang := cstr(n.opts.languageCode)
|
||||
defer freeLang()
|
||||
|
||||
cfg := cTTSSynthesizerConfig{
|
||||
Size: unsafe.Sizeof(cTTSSynthesizerConfig{}),
|
||||
Model: pinPtr(&pinner, &model),
|
||||
Runtime: pinPtr(&pinner, &rt),
|
||||
DefaultLanguageCode: langP,
|
||||
}
|
||||
|
||||
xlog.Info("nemo-speech-cpp: creating synthesizer",
|
||||
"gpu", n.opts.gpu,
|
||||
"codec", n.opts.codecModel,
|
||||
"tokenizer", n.opts.tokenizerDir,
|
||||
"text_normalizer", n.opts.tnDir != "")
|
||||
|
||||
// Compiled before the handle exists so that a full callback table fails the
|
||||
// load, where the operator can see it, rather than the first synthesis.
|
||||
ttsPCMCallback()
|
||||
|
||||
// #nosec G103 -- cfg is a local POD struct borrowed for this call only. Model
|
||||
// and Runtime are pinPtr addresses held by the pinner unpinned on return, the
|
||||
// paths they carry are cstr allocations freed by the defers above, and
|
||||
// nemo_speech_tts_create deep-copies every string it reads.
|
||||
if st := TTSCreate(unsafe.Pointer(&cfg), &n.synth); st != 0 {
|
||||
return statusErrorf(st, "nemo-speech-cpp: tts create: %s", TTSLastError())
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// validateTTSRequest rejects what the runtime would reject, before anything
|
||||
// crosses the ABI, and names the fields this backend drops.
|
||||
//
|
||||
// Empty text is checked here rather than left to the C side for the error code:
|
||||
// src/tts/synthesizer.cpp throws "text is required", which arrives as a status
|
||||
// this layer would otherwise report as Internal, and an empty prompt is a client
|
||||
// mistake, not a backend failure.
|
||||
//
|
||||
// instructions is logged rather than rejected, for the reason the diarization
|
||||
// path logs its own dropped fields: a caller that asked for an expressive style
|
||||
// still wants the audio it can have, and a request naming something this backend
|
||||
// silently ignores should say so where an operator can find it. There is nothing
|
||||
// to map it onto, because MagpieTTS conditions on a speaker, not on a prose
|
||||
// style description: nemo_speech_tts_synthesis_options has speaker and
|
||||
// voice_name and no free-text field at all.
|
||||
func validateTTSRequest(req *pb.TTSRequest) error {
|
||||
if req.GetText() == "" {
|
||||
return status.Error(codes.InvalidArgument, "nemo-speech-cpp: TTSRequest.text is required")
|
||||
}
|
||||
if req.GetInstructions() != "" {
|
||||
xlog.Warn("nemo-speech-cpp: ignoring TTSRequest.instructions, this model has no equivalent")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// outputSampleRate reads the rate the synthesizer emits at.
|
||||
//
|
||||
// A non-positive rate is refused rather than passed on. nemo_speech_tts_sample_rate
|
||||
// answers 0 for a null handle, and a WAV header carrying 0 is not a slightly
|
||||
// wrong file, it is one no player can decode and one whose duration is
|
||||
// undefined.
|
||||
func outputSampleRate(s synthesizer) (uint32, error) {
|
||||
rate := s.sampleRate()
|
||||
if rate <= 0 {
|
||||
return 0, status.Error(codes.Internal,
|
||||
"nemo-speech-cpp: the synthesizer reported no sample rate")
|
||||
}
|
||||
return uint32(rate), nil
|
||||
}
|
||||
|
||||
// wavFile frames PCM as a complete WAV: a header with real sizes, then the
|
||||
// samples.
|
||||
//
|
||||
// pcm is little-endian signed 16-bit mono, which is what the runtime's callback
|
||||
// delivers, and is exactly what pkg/audio's header describes, so nothing is
|
||||
// converted on the way through.
|
||||
func wavFile(pcm []byte, sampleRate uint32) ([]byte, error) {
|
||||
if int64(len(pcm)) > maxWAVDataBytes {
|
||||
return nil, status.Errorf(codes.Internal,
|
||||
"nemo-speech-cpp: synthesis produced %d bytes, more than a WAV header can describe", len(pcm))
|
||||
}
|
||||
|
||||
var buf bytes.Buffer
|
||||
// #nosec G115 -- len(pcm) is checked against maxWAVDataBytes (MaxUint32 minus
|
||||
// the header) immediately above, so the narrowing to uint32 cannot wrap.
|
||||
h := laudio.NewWAVHeaderWithRate(uint32(len(pcm)), sampleRate)
|
||||
if err := h.Write(&buf); err != nil {
|
||||
return nil, status.Errorf(codes.Internal, "nemo-speech-cpp: write WAV header: %v", err)
|
||||
}
|
||||
buf.Write(pcm)
|
||||
return buf.Bytes(), nil
|
||||
}
|
||||
|
||||
// streamingWAVHeader is the first chunk of a streamed synthesis: the same header
|
||||
// with both sizes left unknown, since the total length is not known until the
|
||||
// synthesis ends.
|
||||
//
|
||||
// NewWAVHeaderWithRate derives ChunkSize from the payload length, so the RIFF
|
||||
// size has to be overwritten as well: 36 + 0xFFFFFFFF wraps to 35, which is a
|
||||
// smaller number than the header itself.
|
||||
func streamingWAVHeader(sampleRate uint32) []byte {
|
||||
h := laudio.NewWAVHeaderWithRate(wavStreamingSize, sampleRate)
|
||||
h.ChunkSize = wavStreamingSize
|
||||
|
||||
var buf bytes.Buffer
|
||||
// Write only fails on the writer, and bytes.Buffer does not fail.
|
||||
_ = h.Write(&buf)
|
||||
return buf.Bytes()
|
||||
}
|
||||
|
||||
// synthesizeWAV runs one synthesis and writes the whole result to dst.
|
||||
func synthesizeWAV(s synthesizer, req *pb.TTSRequest, defaultLanguage string) error {
|
||||
if err := validateTTSRequest(req); err != nil {
|
||||
return err
|
||||
}
|
||||
if req.GetDst() == "" {
|
||||
return status.Error(codes.InvalidArgument,
|
||||
"nemo-speech-cpp: TTSRequest.dst (output path) is required")
|
||||
}
|
||||
|
||||
// Read before the synthesis rather than after: it is what the header is
|
||||
// built from, and failing on a bad handle here costs nothing, where failing
|
||||
// after costs the whole synthesis.
|
||||
rate, err := outputSampleRate(s)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
var pcm []byte
|
||||
err = s.synthesize(req, defaultLanguage, func(chunk []byte) bool {
|
||||
pcm = append(pcm, chunk...)
|
||||
return true
|
||||
})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
// A synthesis that returned OK having emitted nothing is a runtime bug, but
|
||||
// the file it would produce is a valid empty WAV, which reaches the user as
|
||||
// silence with no error anywhere.
|
||||
if len(pcm) == 0 {
|
||||
return status.Error(codes.Internal, "nemo-speech-cpp: synthesis produced no audio")
|
||||
}
|
||||
|
||||
out, err := wavFile(pcm, rate)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := os.WriteFile(req.GetDst(), out, 0o600); err != nil {
|
||||
return status.Errorf(codes.Internal, "nemo-speech-cpp: write %q: %v", req.GetDst(), err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// streamWAV runs one synthesis and emits a WAV header followed by each PCM
|
||||
// chunk as the runtime produces it.
|
||||
//
|
||||
// The header is the backend's job, not the caller's: pkg/grpc/server.go only
|
||||
// ever sets Reply.Audio on this path and never Reply.Message, and
|
||||
// core/backend/tts.go's own header branch is keyed on Message, so a backend that
|
||||
// emitted bare PCM would stream something no client could decode. sherpa-onnx
|
||||
// and magpie-tts-cpp both do the same.
|
||||
//
|
||||
// out is not closed here. TTSStream owns it, and closing it in one of two places
|
||||
// depending on how far the request got is how a stream ends up half-closed.
|
||||
func streamWAV(s synthesizer, req *pb.TTSRequest, defaultLanguage string, out chan<- []byte) error {
|
||||
if err := validateTTSRequest(req); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
rate, err := outputSampleRate(s)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
out <- streamingWAVHeader(rate)
|
||||
|
||||
return s.synthesize(req, defaultLanguage, func(chunk []byte) bool {
|
||||
out <- chunk
|
||||
return true
|
||||
})
|
||||
}
|
||||
|
||||
// TTS synthesizes req.Text and writes a WAV to req.Dst.
|
||||
//
|
||||
// The whole body runs inside withEngine, so the family check and the C calls
|
||||
// that trust the handle happen under a single acquisition of engineMu. See the
|
||||
// handoff notes at the bottom of nemospeech.go: Free runs without the backend
|
||||
// lock, so anything that checks the family and then releases the lock before
|
||||
// calling C can have the handle destroyed underneath it.
|
||||
func (n *NemoSpeech) TTS(req *pb.TTSRequest) error {
|
||||
return n.withEngine(familyTTS, func() error {
|
||||
return synthesizeWAV(&cSynthesizer{handle: n.synth}, req, n.opts.languageCode)
|
||||
})
|
||||
}
|
||||
|
||||
// TTSStream synthesizes req.Text and emits the audio on results as it is
|
||||
// produced.
|
||||
//
|
||||
// results is closed on EVERY path, including the family rejection and a
|
||||
// validation failure, and the close is deferred outside withEngine so that a
|
||||
// rejected family still closes it. pkg/grpc/server.go drains this channel from a
|
||||
// goroutine and then blocks on that goroutine finishing, so a channel left open
|
||||
// does not fail the request, it hangs the RPC and, because the backend lock is
|
||||
// still held, every request behind it.
|
||||
//
|
||||
// Holding engineMu for the whole stream is deliberate and is the consequence
|
||||
// documented on the locking protocol: an unload waits for the stream to end
|
||||
// rather than destroying the synthesizer underneath it. There is no unbounded
|
||||
// wait here, because unlike the ASR streams this one is driven by the runtime
|
||||
// and ends when the text does, not when a client decides to stop sending.
|
||||
func (n *NemoSpeech) TTSStream(req *pb.TTSRequest, results chan []byte) error {
|
||||
defer close(results)
|
||||
|
||||
return n.withEngine(familyTTS, func() error {
|
||||
return streamWAV(&cSynthesizer{handle: n.synth}, req, n.opts.languageCode, results)
|
||||
})
|
||||
}
|
||||
@@ -1,727 +0,0 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"encoding/binary"
|
||||
"errors"
|
||||
"go/ast"
|
||||
"go/parser"
|
||||
"go/token"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sync"
|
||||
"unsafe"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
"google.golang.org/grpc/codes"
|
||||
"google.golang.org/grpc/status"
|
||||
|
||||
laudio "github.com/mudler/LocalAI/pkg/audio"
|
||||
pb "github.com/mudler/LocalAI/pkg/grpc/proto"
|
||||
)
|
||||
|
||||
// puregoCallbackTableSize is the hard ceiling purego compiles callbacks into:
|
||||
// maxCB in purego/syscall_sysv.go, which panics rather than growing once it is
|
||||
// full and never releases an entry. Read off the module source for v0.10.0
|
||||
// rather than assumed, because the whole point of the specs below is that
|
||||
// exceeding it kills the process.
|
||||
const puregoCallbackTableSize = 2000
|
||||
|
||||
// fakeSynthesizer scripts what the TTS C API emits for one synthesis.
|
||||
//
|
||||
// There is no MagpieTTS GGUF in the tree, so this is the only way the logic on
|
||||
// top of the ABI (validation, WAV framing, chunk ordering, channel closure)
|
||||
// gets tested at all. It fakes the C contract, not the model: chunks is
|
||||
// whatever nemo_speech_tts_synthesize_text would have handed the callback.
|
||||
type fakeSynthesizer struct {
|
||||
rate int32
|
||||
chunks [][]byte
|
||||
err error
|
||||
|
||||
calls int
|
||||
gotReq *pb.TTSRequest
|
||||
gotLang string
|
||||
cancelled bool
|
||||
}
|
||||
|
||||
func (f *fakeSynthesizer) sampleRate() int32 { return f.rate }
|
||||
|
||||
func (f *fakeSynthesizer) synthesize(req *pb.TTSRequest, defaultLanguage string, sink ttsSink) error {
|
||||
f.calls++
|
||||
f.gotReq = req
|
||||
f.gotLang = defaultLanguage
|
||||
for _, c := range f.chunks {
|
||||
if !sink(c) {
|
||||
f.cancelled = true
|
||||
break
|
||||
}
|
||||
}
|
||||
return f.err
|
||||
}
|
||||
|
||||
var _ = Describe("resolveSpeaker", func() {
|
||||
It("passes a numeric voice through as a speaker index", func() {
|
||||
idx, name := resolveSpeaker("3")
|
||||
Expect(idx).To(Equal(int32(3)))
|
||||
Expect(name).To(BeEmpty())
|
||||
})
|
||||
|
||||
It("passes a named voice through as a name with no index", func() {
|
||||
// voice_name is ignored whenever speaker >= 0 (tts.h, and
|
||||
// synthesizer.cpp only calls resolve_speaker for a negative speaker), so
|
||||
// a named voice must leave the index negative or the name is dropped.
|
||||
idx, name := resolveSpeaker("Aria")
|
||||
Expect(idx).To(Equal(int32(-1)))
|
||||
Expect(name).To(Equal("Aria"))
|
||||
})
|
||||
|
||||
It("leaves both unset for an empty voice so the synthesizer default wins", func() {
|
||||
idx, name := resolveSpeaker("")
|
||||
Expect(idx).To(Equal(int32(-1)))
|
||||
Expect(name).To(BeEmpty())
|
||||
})
|
||||
|
||||
// A negative number is the C API's sentinel for "use the default", not a
|
||||
// speaker. Passing it through as an index would turn a request naming an
|
||||
// invalid voice into one that quietly synthesizes in the default voice.
|
||||
// Handed on as a name instead, resolve_speaker rejects it.
|
||||
It("does not let a negative number become a speaker index", func() {
|
||||
idx, name := resolveSpeaker("-1")
|
||||
Expect(idx).To(Equal(int32(-1)))
|
||||
Expect(name).To(Equal("-1"))
|
||||
})
|
||||
|
||||
It("treats a non-numeric voice that merely starts with digits as a name", func() {
|
||||
idx, name := resolveSpeaker("3-alpha")
|
||||
Expect(idx).To(Equal(int32(-1)))
|
||||
Expect(name).To(Equal("3-alpha"))
|
||||
})
|
||||
|
||||
It("keeps speaker 0 addressable", func() {
|
||||
// 0 is a real speaker index, and the only sentinel here is < 0.
|
||||
idx, name := resolveSpeaker("0")
|
||||
Expect(idx).To(BeZero())
|
||||
Expect(name).To(BeEmpty())
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("applySynthesisParams", func() {
|
||||
// The struct the runtime hands out: speaker/seed/steps/top_k all -1, the
|
||||
// overrides off. Written literally rather than taken from
|
||||
// TTSSynthesisOptionsDefault so the specs run without the shared libraries.
|
||||
defaults := func() cTTSSynthesisOptions {
|
||||
return cTTSSynthesisOptions{
|
||||
Size: unsafe.Sizeof(cTTSSynthesisOptions{}),
|
||||
Speaker: -1,
|
||||
Seed: -1,
|
||||
Steps: -1,
|
||||
TopK: -1,
|
||||
}
|
||||
}
|
||||
|
||||
It("leaves every default alone for an absent params map", func() {
|
||||
o := defaults()
|
||||
applySynthesisParams(&o, nil)
|
||||
Expect(o).To(Equal(defaults()))
|
||||
})
|
||||
|
||||
It("leaves every default alone for an empty params map", func() {
|
||||
o := defaults()
|
||||
applySynthesisParams(&o, map[string]string{})
|
||||
Expect(o).To(Equal(defaults()))
|
||||
})
|
||||
|
||||
It("maps the five knobs the C options struct actually has", func() {
|
||||
o := defaults()
|
||||
applySynthesisParams(&o, map[string]string{
|
||||
"seed": "42",
|
||||
"steps": "12",
|
||||
"top_k": "80",
|
||||
"temperature": "0.7",
|
||||
"cfg_scale": "1.5",
|
||||
})
|
||||
Expect(o.Seed).To(Equal(int32(42)))
|
||||
Expect(o.Steps).To(Equal(int32(12)))
|
||||
Expect(o.TopK).To(Equal(int32(80)))
|
||||
Expect(o.Temperature).To(BeNumerically("~", 0.7, 1e-6))
|
||||
Expect(o.CFGScale).To(BeNumerically("~", 1.5, 1e-6))
|
||||
})
|
||||
|
||||
// magpietts/runtime.cpp reads options.temperature only when
|
||||
// override_temperature is true and otherwise falls back to the
|
||||
// synthesizer's config, so a temperature written without its flag is
|
||||
// silently discarded and the request looks like it was honoured.
|
||||
It("sets the override flag with the temperature", func() {
|
||||
o := defaults()
|
||||
applySynthesisParams(&o, map[string]string{"temperature": "0.4"})
|
||||
Expect(o.OverrideTemperature).To(BeTrue())
|
||||
Expect(o.OverrideCFGScale).To(BeFalse(), "cfg_scale was not asked for")
|
||||
})
|
||||
|
||||
It("sets the override flag with the cfg scale", func() {
|
||||
o := defaults()
|
||||
applySynthesisParams(&o, map[string]string{"cfg_scale": "2"})
|
||||
Expect(o.OverrideCFGScale).To(BeTrue())
|
||||
Expect(o.OverrideTemperature).To(BeFalse(), "temperature was not asked for")
|
||||
})
|
||||
|
||||
It("keeps the defaults when a value cannot be parsed", func() {
|
||||
o := defaults()
|
||||
applySynthesisParams(&o, map[string]string{
|
||||
"seed": "many",
|
||||
"steps": "",
|
||||
"top_k": "8.5",
|
||||
"temperature": "warm",
|
||||
"cfg_scale": "-",
|
||||
})
|
||||
Expect(o).To(Equal(defaults()))
|
||||
})
|
||||
|
||||
// The runtime takes a request's seed only when it is >= 0 and its steps and
|
||||
// top_k only when they are > 0. Writing a parsed 0 or a negative would not
|
||||
// be ignored downstream, it would erase the sentinel that means "use the
|
||||
// synthesizer's value".
|
||||
It("refuses values that would erase a sentinel", func() {
|
||||
o := defaults()
|
||||
applySynthesisParams(&o, map[string]string{
|
||||
"seed": "-5",
|
||||
"steps": "0",
|
||||
"top_k": "0",
|
||||
})
|
||||
Expect(o.Seed).To(Equal(int32(-1)))
|
||||
Expect(o.Steps).To(Equal(int32(-1)))
|
||||
Expect(o.TopK).To(Equal(int32(-1)))
|
||||
})
|
||||
|
||||
It("keeps seed 0, which is a real seed", func() {
|
||||
o := defaults()
|
||||
applySynthesisParams(&o, map[string]string{"seed": "0"})
|
||||
Expect(o.Seed).To(BeZero())
|
||||
})
|
||||
|
||||
// TTSRequest carries fields with no equivalent in
|
||||
// nemo_speech_tts_synthesis_options. They must not be smuggled in through a
|
||||
// param name that happens to match.
|
||||
It("ignores params the C options struct has no field for", func() {
|
||||
o := defaults()
|
||||
applySynthesisParams(&o, map[string]string{
|
||||
"top_p": "0.9",
|
||||
"repetition_penalty": "1.1",
|
||||
"speed": "1.2",
|
||||
"instructions": "cheerful",
|
||||
})
|
||||
Expect(o).To(Equal(defaults()))
|
||||
})
|
||||
})
|
||||
|
||||
// The PCM callback is the one resource in this backend with a hard, silent,
|
||||
// process-wide ceiling: purego compiles each into a fixed table of 2000 entries
|
||||
// and never releases one, so a callback built per request takes the whole
|
||||
// backend process down with a panic after 2000 syntheses. Nothing about a
|
||||
// handful of manual calls shows that.
|
||||
var _ = Describe("ttsPCMCallback", func() {
|
||||
It("compiles a usable callback", func() {
|
||||
Expect(ttsPCMCallback()).ToNot(BeZero())
|
||||
})
|
||||
|
||||
It("compiles exactly one callback however many times it is asked", func() {
|
||||
first := ttsPCMCallback()
|
||||
|
||||
// One more than the table holds: a callback compiled per call panics
|
||||
// with "purego: the maximum number of callbacks has been reached"
|
||||
// before this loop ends, which is precisely the production failure.
|
||||
for i := 0; i <= puregoCallbackTableSize; i++ {
|
||||
Expect(ttsPCMCallback()).To(Equal(first),
|
||||
"call %d returned a different callback, so a new one was compiled", i)
|
||||
}
|
||||
})
|
||||
|
||||
// A source-level assertion, deliberately, because the failure it guards
|
||||
// against is invisible from inside the process: the way a per-request
|
||||
// callback gets reintroduced is by someone calling purego.NewCallback at the
|
||||
// synthesis site instead of going through ttsPCMCallback, and no in-process
|
||||
// spec can reach that call without a MagpieTTS GGUF to synthesize with.
|
||||
// Funnelling every compile through one accessor is what the whole design
|
||||
// rests on, so the single call site is the invariant worth pinning.
|
||||
It("compiles callbacks from exactly one place in the TTS path", func() {
|
||||
fset := token.NewFileSet()
|
||||
file, err := parser.ParseFile(fset, "tts.go", nil, 0)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
// Counted over the syntax tree rather than by grepping the text: the
|
||||
// doc comment on ttsPCMCallback names purego.NewCallback too, and a
|
||||
// spec that cannot tell an explanation from a call would be pinning the
|
||||
// prose.
|
||||
var sites []string
|
||||
ast.Inspect(file, func(n ast.Node) bool {
|
||||
call, ok := n.(*ast.CallExpr)
|
||||
if !ok {
|
||||
return true
|
||||
}
|
||||
sel, ok := call.Fun.(*ast.SelectorExpr)
|
||||
if !ok || sel.Sel.Name != "NewCallback" {
|
||||
return true
|
||||
}
|
||||
if pkg, ok := sel.X.(*ast.Ident); ok && pkg.Name == "purego" {
|
||||
sites = append(sites, fset.Position(call.Pos()).String())
|
||||
}
|
||||
return true
|
||||
})
|
||||
Expect(sites).To(HaveLen(1),
|
||||
"every callback must be compiled through ttsPCMCallback, which memoises it")
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("the PCM sink table", func() {
|
||||
It("routes a chunk to the sink registered for that id", func() {
|
||||
var got []byte
|
||||
id, release := ttsSinks.register(func(pcm []byte) bool {
|
||||
got = pcm
|
||||
return true
|
||||
})
|
||||
defer release()
|
||||
|
||||
src := []byte{1, 2, 3, 4}
|
||||
Expect(ttsDeliverPCM(unsafe.Pointer(&src[0]), uint64(len(src)), id)).To(BeTrue())
|
||||
Expect(got).To(Equal([]byte{1, 2, 3, 4}))
|
||||
})
|
||||
|
||||
// The pointer addresses a std::string the runtime reuses for the next
|
||||
// chunk, so a slice over it would be rewritten under the consumer.
|
||||
It("copies the chunk out of the runtime's buffer", func() {
|
||||
var got []byte
|
||||
id, release := ttsSinks.register(func(pcm []byte) bool {
|
||||
got = pcm
|
||||
return true
|
||||
})
|
||||
defer release()
|
||||
|
||||
src := []byte{9, 8, 7}
|
||||
Expect(ttsDeliverPCM(unsafe.Pointer(&src[0]), uint64(len(src)), id)).To(BeTrue())
|
||||
src[0], src[1], src[2] = 0, 0, 0
|
||||
Expect(got).To(Equal([]byte{9, 8, 7}))
|
||||
})
|
||||
|
||||
It("gives each registration its own id", func() {
|
||||
idA, releaseA := ttsSinks.register(func([]byte) bool { return true })
|
||||
defer releaseA()
|
||||
idB, releaseB := ttsSinks.register(func([]byte) bool { return true })
|
||||
defer releaseB()
|
||||
|
||||
Expect(idA).ToNot(Equal(idB))
|
||||
Expect(idA).ToNot(BeZero(), "id 0 is what a zeroed user_data would carry")
|
||||
Expect(idB).ToNot(BeZero())
|
||||
})
|
||||
|
||||
// Two models synthesizing at once share one callback, and engineMu is
|
||||
// per-model, so nothing serialises them against each other.
|
||||
It("keeps concurrent sinks apart", func() {
|
||||
var mu sync.Mutex
|
||||
got := map[uintptr][]byte{}
|
||||
|
||||
var wg sync.WaitGroup
|
||||
for i := range 16 {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer GinkgoRecover()
|
||||
defer wg.Done()
|
||||
|
||||
src := []byte{byte(i)}
|
||||
var mine []byte
|
||||
id, release := ttsSinks.register(func(pcm []byte) bool {
|
||||
mine = pcm
|
||||
return true
|
||||
})
|
||||
defer release()
|
||||
|
||||
Expect(ttsDeliverPCM(unsafe.Pointer(&src[0]), 1, id)).To(BeTrue())
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
got[id] = mine
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
|
||||
Expect(got).To(HaveLen(16))
|
||||
for id, pcm := range got {
|
||||
Expect(pcm).To(HaveLen(1), "sink %d received the wrong chunk", id)
|
||||
}
|
||||
})
|
||||
|
||||
// A released id means the request has returned. Answering true would leave
|
||||
// the runtime synthesizing into nothing while the RPC that owns the lock
|
||||
// waits for it.
|
||||
It("cancels the synthesis when the sink is gone", func() {
|
||||
id, release := ttsSinks.register(func([]byte) bool { return true })
|
||||
release()
|
||||
|
||||
src := []byte{1}
|
||||
Expect(ttsDeliverPCM(unsafe.Pointer(&src[0]), 1, id)).To(BeFalse())
|
||||
})
|
||||
|
||||
It("cancels for a user_data that was never registered", func() {
|
||||
src := []byte{1}
|
||||
Expect(ttsDeliverPCM(unsafe.Pointer(&src[0]), 1, 0)).To(BeFalse())
|
||||
})
|
||||
|
||||
It("accepts an empty chunk without touching the pointer", func() {
|
||||
id, release := ttsSinks.register(func([]byte) bool {
|
||||
Fail("an empty chunk must not reach the sink")
|
||||
return true
|
||||
})
|
||||
defer release()
|
||||
|
||||
Expect(ttsDeliverPCM(nil, 0, id)).To(BeTrue())
|
||||
})
|
||||
|
||||
It("passes the sink's cancellation back to the runtime", func() {
|
||||
id, release := ttsSinks.register(func([]byte) bool { return false })
|
||||
defer release()
|
||||
|
||||
src := []byte{1}
|
||||
Expect(ttsDeliverPCM(unsafe.Pointer(&src[0]), 1, id)).To(BeFalse())
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("ttsModelConfig", func() {
|
||||
// Three adjacent same-typed path fields: swapping two changes neither the
|
||||
// struct size nor any offset, so abi_test.go's layout assertions cannot see
|
||||
// it and the runtime would load the codec as the acoustic model.
|
||||
It("assigns each path to its own field", func() {
|
||||
cfg := ttsModelConfig(1, 2, 3, 4)
|
||||
Expect(cfg.MagpieModel).To(Equal(uintptr(1)))
|
||||
Expect(cfg.CodecModel).To(Equal(uintptr(2)))
|
||||
Expect(cfg.TokenizerModelDir).To(Equal(uintptr(3)))
|
||||
Expect(cfg.TextNormalizerModelDir).To(Equal(uintptr(4)))
|
||||
})
|
||||
|
||||
// A config sent with the wrong size has every field past it ignored by
|
||||
// HAS_FIELD, and the model loads with defaults instead of failing.
|
||||
It("declares the size the runtime validates against", func() {
|
||||
Expect(ttsModelConfig(1, 2, 3, 4).Size).To(Equal(unsafe.Sizeof(cTTSModelConfig{})))
|
||||
})
|
||||
|
||||
It("leaves an unset text normalizer null", func() {
|
||||
Expect(ttsModelConfig(1, 2, 3, 0).TextNormalizerModelDir).To(BeZero())
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("ttsRuntimeBackend", func() {
|
||||
// -1 is this backend's documented "CPU" everywhere (asr.h: "-1 = CPU") and
|
||||
// it is also the default, so it has to pin the preference rather than leave
|
||||
// the runtime free to pick CUDA.
|
||||
It("pins CPU for a negative gpu option", func() {
|
||||
Expect(ttsRuntimeBackend(-1)).To(Equal(ttsBackendCPU))
|
||||
})
|
||||
|
||||
// nemo_speech_tts_runtime_config has no device index at all, so a request
|
||||
// for a particular device cannot be honoured and AUTO is the honest answer.
|
||||
It("leaves the choice to the runtime when a device was named", func() {
|
||||
Expect(ttsRuntimeBackend(0)).To(Equal(ttsBackendAuto))
|
||||
Expect(ttsRuntimeBackend(3)).To(Equal(ttsBackendAuto))
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("WAV framing", func() {
|
||||
// 16-bit mono little-endian, the format the runtime's callback delivers.
|
||||
pcm := []byte{0x01, 0x00, 0xff, 0x7f, 0x00, 0x80}
|
||||
|
||||
It("writes a header the audio helpers can read back", func() {
|
||||
out, err := wavFile(pcm, 22050)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
body, rate := laudio.ParseWAV(out)
|
||||
Expect(rate).To(Equal(22050))
|
||||
Expect(body).To(Equal(pcm))
|
||||
})
|
||||
|
||||
It("describes the payload it actually carries", func() {
|
||||
out, err := wavFile(pcm, 22050)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(out).To(HaveLen(laudio.WAVHeaderSize + len(pcm)))
|
||||
|
||||
Expect(string(out[0:4])).To(Equal("RIFF"))
|
||||
Expect(string(out[8:12])).To(Equal("WAVE"))
|
||||
Expect(binary.LittleEndian.Uint32(out[4:8])).To(Equal(uint32(36 + len(pcm))))
|
||||
Expect(binary.LittleEndian.Uint32(out[40:44])).To(Equal(uint32(len(pcm))))
|
||||
Expect(binary.LittleEndian.Uint16(out[22:24])).To(Equal(uint16(1)), "mono")
|
||||
Expect(binary.LittleEndian.Uint16(out[34:36])).To(Equal(uint16(16)), "16-bit")
|
||||
Expect(binary.LittleEndian.Uint32(out[24:28])).To(Equal(uint32(22050)))
|
||||
// byte rate = sample rate * block align, and a wrong one plays back at
|
||||
// the wrong speed in players that trust it.
|
||||
Expect(binary.LittleEndian.Uint32(out[28:32])).To(Equal(uint32(22050 * 2)))
|
||||
})
|
||||
|
||||
It("carries whatever rate the synthesizer reported", func() {
|
||||
out, err := wavFile(pcm, 44100)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
_, rate := laudio.ParseWAV(out)
|
||||
Expect(rate).To(Equal(44100))
|
||||
})
|
||||
|
||||
Describe("the streaming header", func() {
|
||||
It("is a complete header on its own", func() {
|
||||
h := streamingWAVHeader(22050)
|
||||
Expect(h).To(HaveLen(laudio.WAVHeaderSize))
|
||||
Expect(string(h[0:4])).To(Equal("RIFF"))
|
||||
Expect(string(h[8:12])).To(Equal("WAVE"))
|
||||
Expect(binary.LittleEndian.Uint32(h[24:28])).To(Equal(uint32(22050)))
|
||||
})
|
||||
|
||||
// NewWAVHeaderWithRate derives ChunkSize from the payload length, so
|
||||
// leaving it alone would write 36 + 0xFFFFFFFF, which wraps to 35: a
|
||||
// RIFF size smaller than the header itself.
|
||||
It("leaves both sizes unknown rather than wrapping", func() {
|
||||
h := streamingWAVHeader(22050)
|
||||
Expect(binary.LittleEndian.Uint32(h[4:8])).To(Equal(uint32(0xFFFFFFFF)))
|
||||
Expect(binary.LittleEndian.Uint32(h[40:44])).To(Equal(uint32(0xFFFFFFFF)))
|
||||
})
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("synthesizeWAV", func() {
|
||||
var dst string
|
||||
|
||||
BeforeEach(func() {
|
||||
dst = filepath.Join(GinkgoT().TempDir(), "out.wav")
|
||||
})
|
||||
|
||||
It("writes one WAV holding every chunk in order", func() {
|
||||
s := &fakeSynthesizer{rate: 22050, chunks: [][]byte{{1, 0}, {2, 0}, {3, 0}}}
|
||||
Expect(synthesizeWAV(s, &pb.TTSRequest{Text: "hello", Dst: dst}, "")).To(Succeed())
|
||||
|
||||
out, err := os.ReadFile(dst)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
body, rate := laudio.ParseWAV(out)
|
||||
Expect(rate).To(Equal(22050))
|
||||
Expect(body).To(Equal([]byte{1, 0, 2, 0, 3, 0}))
|
||||
})
|
||||
|
||||
It("hands the request and the model default language to the runtime", func() {
|
||||
s := &fakeSynthesizer{rate: 22050, chunks: [][]byte{{1, 0}}}
|
||||
req := &pb.TTSRequest{Text: "hello", Dst: dst, Voice: "Aria"}
|
||||
Expect(synthesizeWAV(s, req, "it-IT")).To(Succeed())
|
||||
Expect(s.gotReq).To(Equal(req))
|
||||
Expect(s.gotLang).To(Equal("it-IT"))
|
||||
})
|
||||
|
||||
It("rejects an empty text before it reaches the runtime", func() {
|
||||
s := &fakeSynthesizer{rate: 22050}
|
||||
err := synthesizeWAV(s, &pb.TTSRequest{Dst: dst}, "")
|
||||
Expect(status.Code(err)).To(Equal(codes.InvalidArgument))
|
||||
Expect(s.calls).To(BeZero())
|
||||
})
|
||||
|
||||
// instructions has no equivalent in nemo_speech_tts_synthesis_options, which
|
||||
// conditions on a speaker rather than a prose style. Dropping it must not
|
||||
// fail the request: the caller still wants the audio it can have.
|
||||
It("synthesizes anyway for a request carrying instructions it cannot honour", func() {
|
||||
s := &fakeSynthesizer{rate: 22050, chunks: [][]byte{{1, 0}}}
|
||||
instructions := "speak cheerfully"
|
||||
Expect(synthesizeWAV(s, &pb.TTSRequest{
|
||||
Text: "hello",
|
||||
Dst: dst,
|
||||
Instructions: &instructions,
|
||||
}, "")).To(Succeed())
|
||||
Expect(dst).To(BeAnExistingFile())
|
||||
})
|
||||
|
||||
It("rejects a request with no destination", func() {
|
||||
s := &fakeSynthesizer{rate: 22050}
|
||||
err := synthesizeWAV(s, &pb.TTSRequest{Text: "hello"}, "")
|
||||
Expect(status.Code(err)).To(Equal(codes.InvalidArgument))
|
||||
Expect(s.calls).To(BeZero())
|
||||
})
|
||||
|
||||
// A zero rate is what a null handle reports. The file it would produce is
|
||||
// undecodable, and the synthesis that produced it would be wasted.
|
||||
It("refuses to write a file at an unusable sample rate", func() {
|
||||
s := &fakeSynthesizer{rate: 0, chunks: [][]byte{{1, 0}}}
|
||||
err := synthesizeWAV(s, &pb.TTSRequest{Text: "hello", Dst: dst}, "")
|
||||
Expect(status.Code(err)).To(Equal(codes.Internal))
|
||||
Expect(s.calls).To(BeZero())
|
||||
Expect(dst).ToNot(BeAnExistingFile())
|
||||
})
|
||||
|
||||
It("propagates a synthesis failure and writes nothing", func() {
|
||||
boom := errors.New("boom")
|
||||
s := &fakeSynthesizer{rate: 22050, chunks: [][]byte{{1, 0}}, err: boom}
|
||||
Expect(synthesizeWAV(s, &pb.TTSRequest{Text: "hello", Dst: dst}, "")).To(MatchError(boom))
|
||||
Expect(dst).ToNot(BeAnExistingFile())
|
||||
})
|
||||
|
||||
// An empty WAV is a valid file, so this would otherwise reach the user as
|
||||
// silence with no error anywhere.
|
||||
It("fails rather than write a silent file when nothing was produced", func() {
|
||||
s := &fakeSynthesizer{rate: 22050}
|
||||
err := synthesizeWAV(s, &pb.TTSRequest{Text: "hello", Dst: dst}, "")
|
||||
Expect(status.Code(err)).To(Equal(codes.Internal))
|
||||
Expect(dst).ToNot(BeAnExistingFile())
|
||||
})
|
||||
|
||||
It("reports a destination it cannot write", func() {
|
||||
s := &fakeSynthesizer{rate: 22050, chunks: [][]byte{{1, 0}}}
|
||||
bad := filepath.Join(GinkgoT().TempDir(), "no-such-dir", "out.wav")
|
||||
err := synthesizeWAV(s, &pb.TTSRequest{Text: "hello", Dst: bad}, "")
|
||||
Expect(status.Code(err)).To(Equal(codes.Internal))
|
||||
Expect(err.Error()).To(ContainSubstring("out.wav"))
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("streamWAV", func() {
|
||||
// drain collects everything streamWAV emits. The channel is buffered
|
||||
// because streamWAV sends inline, so an unbuffered one would deadlock the
|
||||
// spec rather than fail it.
|
||||
drain := func(s synthesizer, req *pb.TTSRequest) ([][]byte, error) {
|
||||
out := make(chan []byte, 16)
|
||||
err := streamWAV(s, req, "", out)
|
||||
close(out)
|
||||
|
||||
var got [][]byte
|
||||
for c := range out {
|
||||
got = append(got, c)
|
||||
}
|
||||
return got, err
|
||||
}
|
||||
|
||||
It("emits the header first, then each chunk as it arrives", func() {
|
||||
s := &fakeSynthesizer{rate: 22050, chunks: [][]byte{{1, 0}, {2, 0}}}
|
||||
got, err := drain(s, &pb.TTSRequest{Text: "hello"})
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
Expect(got).To(HaveLen(3))
|
||||
Expect(got[0]).To(Equal(streamingWAVHeader(22050)))
|
||||
Expect(got[1]).To(Equal([]byte{1, 0}))
|
||||
Expect(got[2]).To(Equal([]byte{2, 0}))
|
||||
})
|
||||
|
||||
// pkg/grpc/server.go only ever sets Reply.Audio, so core/backend's own
|
||||
// header branch (keyed on Reply.Message) never runs and a backend that
|
||||
// emitted bare PCM would stream something no client could decode.
|
||||
It("owns the header rather than leaving it to the caller", func() {
|
||||
s := &fakeSynthesizer{rate: 44100, chunks: [][]byte{{1, 0}}}
|
||||
got, err := drain(s, &pb.TTSRequest{Text: "hello"})
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(string(got[0][0:4])).To(Equal("RIFF"))
|
||||
Expect(binary.LittleEndian.Uint32(got[0][24:28])).To(Equal(uint32(44100)))
|
||||
})
|
||||
|
||||
It("rejects an empty text before emitting anything", func() {
|
||||
s := &fakeSynthesizer{rate: 22050, chunks: [][]byte{{1, 0}}}
|
||||
got, err := drain(s, &pb.TTSRequest{})
|
||||
Expect(status.Code(err)).To(Equal(codes.InvalidArgument))
|
||||
Expect(got).To(BeEmpty())
|
||||
Expect(s.calls).To(BeZero())
|
||||
})
|
||||
|
||||
It("emits no header at an unusable sample rate", func() {
|
||||
s := &fakeSynthesizer{rate: 0, chunks: [][]byte{{1, 0}}}
|
||||
got, err := drain(s, &pb.TTSRequest{Text: "hello"})
|
||||
Expect(status.Code(err)).To(Equal(codes.Internal))
|
||||
Expect(got).To(BeEmpty())
|
||||
})
|
||||
|
||||
It("propagates a synthesis failure after the chunks it did emit", func() {
|
||||
boom := errors.New("boom")
|
||||
s := &fakeSynthesizer{rate: 22050, chunks: [][]byte{{1, 0}}, err: boom}
|
||||
got, err := drain(s, &pb.TTSRequest{Text: "hello"})
|
||||
Expect(err).To(MatchError(boom))
|
||||
Expect(got).To(HaveLen(2))
|
||||
})
|
||||
|
||||
// streamWAV must not close the channel: TTSStream owns it, and closing in
|
||||
// one of two places depending on how far the request got is how a stream
|
||||
// ends up double-closed.
|
||||
It("leaves the channel open for its caller to close", func() {
|
||||
s := &fakeSynthesizer{rate: 22050, chunks: [][]byte{{1, 0}}}
|
||||
out := make(chan []byte, 4)
|
||||
Expect(streamWAV(s, &pb.TTSRequest{Text: "hello"}, "", out)).To(Succeed())
|
||||
Expect(func() { close(out) }).ToNot(Panic())
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("the TTS RPCs", func() {
|
||||
It("refuses TTS on a model loaded as another family", func() {
|
||||
n := &NemoSpeech{fam: familyASR}
|
||||
err := n.TTS(&pb.TTSRequest{Text: "hello", Dst: "/tmp/out.wav"})
|
||||
Expect(status.Code(err)).To(Equal(codes.Unimplemented))
|
||||
})
|
||||
|
||||
It("refuses TTS on an unloaded model", func() {
|
||||
n := &NemoSpeech{}
|
||||
Expect(status.Code(n.TTS(&pb.TTSRequest{Text: "hello", Dst: "/tmp/out.wav"}))).
|
||||
To(Equal(codes.Unimplemented))
|
||||
})
|
||||
|
||||
It("releases the engine lock after a refusal", func() {
|
||||
n := &NemoSpeech{fam: familyASR}
|
||||
Expect(n.TTS(&pb.TTSRequest{Text: "hello", Dst: "/tmp/out.wav"})).ToNot(Succeed())
|
||||
Expect(n.engineMu.TryLock()).To(BeTrue())
|
||||
n.engineMu.Unlock()
|
||||
})
|
||||
|
||||
// pkg/grpc/server.go drains this channel from a goroutine and then blocks
|
||||
// on that goroutine finishing, so a channel left open does not fail the
|
||||
// request, it hangs the RPC with the backend lock still held. Every exit
|
||||
// path has to close it.
|
||||
Describe("TTSStream channel closure", func() {
|
||||
// streamed runs TTSStream the way the server does and returns once the
|
||||
// channel has been closed, so a spec that hangs is a real hang.
|
||||
streamed := func(n *NemoSpeech, req *pb.TTSRequest) ([][]byte, error) {
|
||||
ch := make(chan []byte, 16)
|
||||
done := make(chan [][]byte, 1)
|
||||
go func() {
|
||||
defer GinkgoRecover()
|
||||
var got [][]byte
|
||||
for c := range ch {
|
||||
got = append(got, c)
|
||||
}
|
||||
done <- got
|
||||
}()
|
||||
|
||||
err := n.TTSStream(req, ch)
|
||||
var got [][]byte
|
||||
Eventually(done).Should(Receive(&got))
|
||||
return got, err
|
||||
}
|
||||
|
||||
It("closes the channel when the family does not match", func() {
|
||||
n := &NemoSpeech{fam: familyASR}
|
||||
got, err := streamed(n, &pb.TTSRequest{Text: "hello"})
|
||||
Expect(status.Code(err)).To(Equal(codes.Unimplemented))
|
||||
Expect(got).To(BeEmpty())
|
||||
})
|
||||
|
||||
It("closes the channel when the model was never loaded", func() {
|
||||
n := &NemoSpeech{}
|
||||
_, err := streamed(n, &pb.TTSRequest{Text: "hello"})
|
||||
Expect(status.Code(err)).To(Equal(codes.Unimplemented))
|
||||
})
|
||||
|
||||
// familyTTS with a zero handle: validation has to reject this before
|
||||
// anything reaches the C entry points, which are nil function values
|
||||
// until openLibraries has bound them.
|
||||
It("closes the channel when the request is rejected", func() {
|
||||
n := &NemoSpeech{fam: familyTTS}
|
||||
got, err := streamed(n, &pb.TTSRequest{})
|
||||
Expect(status.Code(err)).To(Equal(codes.InvalidArgument))
|
||||
Expect(got).To(BeEmpty())
|
||||
})
|
||||
|
||||
It("releases the engine lock afterwards", func() {
|
||||
n := &NemoSpeech{fam: familyTTS}
|
||||
_, err := streamed(n, &pb.TTSRequest{})
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(n.engineMu.TryLock()).To(BeTrue())
|
||||
n.engineMu.Unlock()
|
||||
})
|
||||
})
|
||||
|
||||
// The same guard on the offline path: a rejected request must not reach a
|
||||
// nil C function through a zero handle.
|
||||
It("rejects an invalid TTS request without touching the runtime", func() {
|
||||
n := &NemoSpeech{fam: familyTTS}
|
||||
Expect(status.Code(n.TTS(&pb.TTSRequest{Dst: "/tmp/out.wav"}))).To(Equal(codes.InvalidArgument))
|
||||
Expect(status.Code(n.TTS(&pb.TTSRequest{Text: "hello"}))).To(Equal(codes.InvalidArgument))
|
||||
})
|
||||
})
|
||||
@@ -7,18 +7,8 @@ GO_TAGS?=
|
||||
JOBS?=$(shell nproc --ignore=1)
|
||||
|
||||
# qwentts.cpp version
|
||||
#
|
||||
# Held at 35ebe537 rather than tracking latest: abab6b3 hangs in synthesis.
|
||||
# TTS() never returns from the native call, so tests-qwen3-tts-cpp goes from
|
||||
# ~5 minutes to the 20 minute Go test timeout. Reproduced on master on
|
||||
# 2026-08-01 and again on re-run, and the bump PR (#11241) was merged with
|
||||
# this same check already red.
|
||||
#
|
||||
# The regression is in 35ebe537..abab6b3, three upstream commits whose only
|
||||
# functional change is 26dd8adb, "predictor: unroll the frame into one cgraph
|
||||
# and sample in standard ops". Restore the bump once that is fixed upstream.
|
||||
QWEN3TTS_REPO?=https://github.com/ServeurpersoCom/qwentts.cpp
|
||||
QWEN3TTS_CPP_VERSION?=35ebe5376b82a0a59d008586d55bbe623d449011
|
||||
QWEN3TTS_CPP_VERSION?=abab6b3bf317cfa1b788efce1d25f4f9239395ad
|
||||
SO_TARGET?=libgoqwen3ttscpp.so
|
||||
|
||||
CMAKE_ARGS+=-DBUILD_SHARED_LIBS=OFF
|
||||
|
||||
@@ -8,7 +8,7 @@ JOBS?=$(shell nproc --ignore=1)
|
||||
|
||||
# stablediffusion.cpp (ggml)
|
||||
STABLEDIFFUSION_GGML_REPO?=https://github.com/leejet/stable-diffusion.cpp
|
||||
STABLEDIFFUSION_GGML_VERSION?=c6beeef35526c6dc94b74a7fb69f9d2e6a2a7a12
|
||||
STABLEDIFFUSION_GGML_VERSION?=e31a86ce9110b11a98bd5990c329093244c2d1e3
|
||||
|
||||
CMAKE_ARGS+=-DGGML_MAX_NAME=128
|
||||
|
||||
|
||||
@@ -11,30 +11,7 @@ JOBS?=$(shell nproc --ignore=1 2>/dev/null || sysctl -n hw.ncpu 2>/dev/null || e
|
||||
|
||||
# vllm.cpp version
|
||||
VLLM_CPP_REPO?=https://github.com/mudler/vllm.cpp
|
||||
VLLM_CPP_VERSION?=0757cac231ecd571a83c4fd2f50805c9251fc225
|
||||
|
||||
# MLX GEMM provider (darwin/metal only; see the metal branch below for why).
|
||||
# Consumed as the prebuilt pip wheel: building MLX from source needs `xcrun
|
||||
# metal`, i.e. a full Xcode the macOS runners do not have, while the wheel ships
|
||||
# include/, lib/libmlx.dylib and the compiled mlx.metallib ready to link.
|
||||
#
|
||||
# DEFAULT ON, but ONLY because VLLM_CPP_VERSION above is pinned at or past
|
||||
# vllm.cpp 89c46aeb, which SHAPE-GATES the provider to prefill. The ordering is
|
||||
# load-bearing, not incidental:
|
||||
#
|
||||
# pin >= 89c46aeb, MLX on -> 99.1% of MLX-LM (gated: prefill only)
|
||||
# pin < 89c46aeb, MLX on -> ~51% (ungated: it also takes decode)
|
||||
#
|
||||
# MLX's steel GEMM wins prefill (537 ms TTFT against 602) and loses decode badly,
|
||||
# because the provider pays an mx::eval sync plus an output memcpy per call and
|
||||
# decode makes ~112 calls per TOKEN. Ungated it does both; gated it does only the
|
||||
# good half. So if this pin is ever moved BACKWARDS, this default must go with it.
|
||||
VLLM_CPP_MLX?=on
|
||||
MLX_VERSION?=0.29.4
|
||||
MLX_VENV?=$(abspath ./mlx-venv)
|
||||
# Resolved lazily (recursive `=`, not `:=`): the glob only matches once the venv
|
||||
# target has run, and the interpreter version in the path varies per runner.
|
||||
MLX_ROOT=$(shell echo $(MLX_VENV)/lib/python*/site-packages/mlx)
|
||||
VLLM_CPP_VERSION?=9e1c9025ae61167a3335454d7cc0de6093c21845
|
||||
|
||||
# The backend consumes only the stable C ABI (libvllm + include/vllm.h), so the
|
||||
# server, examples and tests of the engine are never built here.
|
||||
@@ -72,23 +49,6 @@ else ifeq ($(BUILD_TYPE),vulkan)
|
||||
CMAKE_ARGS+=-DVLLM_CPP_VULKAN=ON -DVLLM_CPP_CUDA=OFF
|
||||
else ifeq ($(BUILD_TYPE),metal)
|
||||
CMAKE_ARGS+=-DVLLM_CPP_METAL=ON
|
||||
# The optional MLX GEMM provider. vllm.cpp keeps it OFF by default because it
|
||||
# is a ~19 MB libmlx.dylib plus a ~105 MB mlx.metallib, and upstream's
|
||||
# position is that it must earn that cost by measurement. It does, on the
|
||||
# only hardware this build targets: measured on an Apple M4 against the
|
||||
# native MSL GEMM in the SAME binary (arms toggled by
|
||||
# VT_OP_PROVIDER_DISABLE=mlx), Qwen3-1.7B-bf16 p=512 g=128, it is 1.5x to
|
||||
# 2.2x aggregate throughput and 2x to 3x faster TTFT, at equal peak memory
|
||||
# and bit-identical output on every parity shape. See vllm.cpp
|
||||
# docs/BENCHMARKS.md "MLX GEMM provider A/B on Apple M4".
|
||||
#
|
||||
# MLX delegates the dense GEMM ONLY: kPagedAttention stays vllm.cpp's own
|
||||
# kernel, because MLX has no paged-KV primitive at all.
|
||||
#
|
||||
# Set VLLM_CPP_MLX=off for a Metal build without it (smaller image, slower).
|
||||
ifeq ($(VLLM_CPP_MLX),on)
|
||||
MLX_ENABLED=1
|
||||
endif
|
||||
else
|
||||
CMAKE_ARGS+=-DVLLM_CPP_CUDA=OFF
|
||||
endif
|
||||
@@ -96,12 +56,6 @@ endif
|
||||
UNAME_S := $(shell uname -s)
|
||||
ifeq ($(UNAME_S),Darwin)
|
||||
LIB=libvllm.dylib
|
||||
# Apple Clang diagnoses a pair of constant-folded array bounds in the Metal
|
||||
# build as a GNU extension. Disable that diagnostic for both Objective-C and
|
||||
# C++ because vllm.cpp appends target-local -Werror after these global flags.
|
||||
CMAKE_ARGS+=-DCMAKE_CXX_FLAGS=-Wno-gnu-folding-constant
|
||||
CMAKE_ARGS+=-DCMAKE_OBJC_FLAGS=-Wno-gnu-folding-constant
|
||||
CMAKE_ARGS+=-DCMAKE_OBJCXX_FLAGS=-Wno-gnu-folding-constant
|
||||
else
|
||||
LIB=libvllm.so
|
||||
endif
|
||||
@@ -114,54 +68,10 @@ sources/vllm.cpp:
|
||||
git fetch --depth 1 origin $(VLLM_CPP_VERSION) && \
|
||||
git checkout FETCH_HEAD
|
||||
|
||||
ifeq ($(MLX_ENABLED),1)
|
||||
# A stamp FILE, not a phony target: a phony prerequisite is always "newer" than
|
||||
# $(LIB) and would re-link libvllm on every invocation. Keyed on the version so
|
||||
# a MLX_VERSION bump reinstalls instead of silently reusing the old wheel.
|
||||
MLX_STAMP=$(MLX_VENV)/.mlx-$(MLX_VERSION).stamp
|
||||
MLX_CMAKE_ARGS=-DVLLM_CPP_MLX=ON -DMLX_ROOT=$(MLX_ROOT)
|
||||
|
||||
$(MLX_STAMP):
|
||||
@if [ ! -x "$(MLX_VENV)/bin/pip" ]; then \
|
||||
python3 -m venv "$(MLX_VENV)" || { echo "vllm-cpp: python3 with venv is required to build the MLX provider; pass VLLM_CPP_MLX=off to build Metal without it" >&2; exit 1; }; \
|
||||
fi
|
||||
"$(MLX_VENV)"/bin/pip install --quiet --disable-pip-version-check "mlx==$(MLX_VERSION)"
|
||||
@# Resolved in the SHELL, not by $(MLX_ROOT): make expands a whole recipe
|
||||
@# before running its first line, so the glob would still be unmatched here.
|
||||
@# Every later use (the cmake args, package.sh) expands after this target has
|
||||
@# completed, where $(MLX_ROOT) does resolve.
|
||||
@root=$$(echo "$(MLX_VENV)"/lib/python*/site-packages/mlx); \
|
||||
test -f "$$root/lib/libmlx.dylib" -a -f "$$root/include/mlx/array.h" || \
|
||||
{ echo "vllm-cpp: mlx==$(MLX_VERSION) did not provide lib/libmlx.dylib + include/mlx/array.h under $$root" >&2; exit 1; }
|
||||
touch $@
|
||||
else
|
||||
MLX_STAMP=
|
||||
MLX_CMAKE_ARGS=
|
||||
endif
|
||||
|
||||
# govllmcpp.go mirrors vllm.h by hand, and the only guard against the two
|
||||
# drifting apart is the vllm_abi_version check inside registerLib - which fires
|
||||
# at runtime, on the user's machine, taking down every model load (issue
|
||||
# #11379). Compare the two here instead, so moving VLLM_CPP_VERSION past the
|
||||
# mirrors turns the build red while the header is still around to diff.
|
||||
abi-check: sources/vllm.cpp
|
||||
@engine=$$(sed -n 's/^#define VLLM_ABI_VERSION \([0-9][0-9]*\).*/\1/p' sources/vllm.cpp/include/vllm.h); \
|
||||
backend=$$(sed -n 's/^const abiVersion = \([0-9][0-9]*\).*/\1/p' govllmcpp.go); \
|
||||
if [ -z "$$engine" ] || [ -z "$$backend" ]; then \
|
||||
echo "vllm-cpp: cannot read the ABI version (engine='$$engine' backend='$$backend')" >&2; exit 1; \
|
||||
fi; \
|
||||
if [ "$$engine" != "$$backend" ]; then \
|
||||
echo "vllm-cpp: ABI mismatch: vllm.cpp $(VLLM_CPP_VERSION) is v$$engine, govllmcpp.go mirrors v$$backend." >&2; \
|
||||
echo " Update the struct mirrors and abiVersion in govllmcpp.go (and the offsets in vllmcpp_test.go) to v$$engine." >&2; \
|
||||
exit 1; \
|
||||
fi; \
|
||||
echo "vllm-cpp: ABI v$$engine matches the pinned engine"
|
||||
|
||||
$(LIB): sources/vllm.cpp $(MLX_STAMP)
|
||||
$(MAKE) abi-check
|
||||
$(LIB): sources/vllm.cpp
|
||||
mkdir -p build && \
|
||||
cd build && \
|
||||
cmake ../sources/vllm.cpp $(CMAKE_ARGS) $(MLX_CMAKE_ARGS) && \
|
||||
cmake ../sources/vllm.cpp $(CMAKE_ARGS) && \
|
||||
cmake --build . --config Release -j$(JOBS) --target vllm_shared
|
||||
cp -fL build/$(LIB) ./$(LIB)
|
||||
|
||||
@@ -169,18 +79,16 @@ vllm-cpp: main.go govllmcpp.go backend.go options.go $(LIB)
|
||||
CGO_ENABLED=0 $(GOCMD) build -tags "$(GO_TAGS)" -o vllm-cpp ./
|
||||
|
||||
package: vllm-cpp
|
||||
MLX_ROOT="$(MLX_ROOT)" bash package.sh
|
||||
bash package.sh
|
||||
|
||||
build: package
|
||||
|
||||
clean: purge
|
||||
rm -rf libvllm.so libvllm.dylib package sources/vllm.cpp vllm-cpp "$(MLX_VENV)"
|
||||
rm -rf libvllm.so libvllm.dylib package sources/vllm.cpp vllm-cpp
|
||||
|
||||
purge:
|
||||
rm -rf build
|
||||
|
||||
.PHONY: abi-check
|
||||
|
||||
.NOTPARALLEL:
|
||||
|
||||
# The unit specs are pure Go (struct mirrors, option mapping, load
|
||||
|
||||
@@ -6,7 +6,7 @@ safetensors + GGUF loading, CUDA / CPU / Metal / Vulkan) with no Python at
|
||||
inference time.
|
||||
|
||||
The backend dlopens the engine's stable C ABI (`libvllm`, `include/vllm.h`,
|
||||
ABI v10) through purego:
|
||||
ABI v2) through purego:
|
||||
|
||||
- `Load` -> `vllm_engine_load`: accepts a `.gguf` file or a HF-style model
|
||||
directory (`config.json` + safetensors). `context_size` maps to
|
||||
@@ -29,12 +29,6 @@ ABI v10) through purego:
|
||||
LocalAI's Go-side grammar-constrained tool calling; JSON-schema / regex /
|
||||
choice constraints are also exposed by the ABI.
|
||||
|
||||
The struct mirrors in `govllmcpp.go` are hand-written against one ABI version,
|
||||
and the engine refuses to load against any other. Moving `VLLM_CPP_VERSION` in
|
||||
the Makefile therefore means updating `abiVersion` plus the mirrors (and their
|
||||
offsets in `vllmcpp_test.go`) in the same change; `make abi-check` compares the
|
||||
pinned header against the bindings and the library build runs it first.
|
||||
|
||||
Model config example:
|
||||
|
||||
```yaml
|
||||
@@ -47,50 +41,5 @@ options:
|
||||
- max_num_seqs:16
|
||||
```
|
||||
|
||||
## Apple Silicon: the MLX GEMM provider (ON by default, gated to prefill)
|
||||
|
||||
`BUILD_TYPE=metal` builds vllm.cpp's MLX provider for the dense GEMM
|
||||
(`VLLM_CPP_MLX=on`, the default here). It is on because upstream now SHAPE-GATES
|
||||
it to prefill; it was briefly off in this branch's history, and that was correct
|
||||
at the time for an ungated provider.
|
||||
|
||||
The gate matters more than the flag. MLX's steel GEMM wins prefill but loses
|
||||
decode, because the provider pays an `mx::eval` synchronisation plus an output
|
||||
memcpy on every call and decode makes ~112 calls *per token*. Measured on an
|
||||
Apple M4, Qwen3-1.7B-bf16 warm at p=512 g=128:
|
||||
|
||||
| configuration | prefill TTFT | warm throughput |
|
||||
|---|--:|--:|
|
||||
| MLX **gated to prefill** (pin >= 89c46aeb) | **524.5 ms** | **24.37 tok/s, 97.6% of MLX-LM** |
|
||||
| MLX ungated (older pins) | 537 ms | 12.7 tok/s |
|
||||
| MLX off | 602 ms | 23.9 tok/s, 95.9% |
|
||||
|
||||
Ratios are against an MLX-LM baseline measured INTERLEAVED with ours over four
|
||||
ABBA blocks (its spread 0.34%, ours 0.12%). An earlier revision of this file
|
||||
claimed 99.1%; that used a two-run MLX-LM baseline containing an outlier and
|
||||
overstated us by about 1.5 points.
|
||||
|
||||
**`VLLM_CPP_VERSION` and this flag are coupled.** Moving the pin back before
|
||||
`89c46aeb` while leaving `VLLM_CPP_MLX=on` would take the middle row — roughly
|
||||
half throughput. If you roll the pin back, roll the default back with it.
|
||||
|
||||
One caveat: MLX's GEMM is not bit-identical to the native kernel, so an MLX build
|
||||
produces a different greedy sequence than a non-MLX one. That is a property of the
|
||||
provider, not of the gate, and it predates this packaging. Full disposition in
|
||||
vllm.cpp `docs/BENCHMARKS.md`.
|
||||
|
||||
Build knobs:
|
||||
|
||||
- `VLLM_CPP_MLX=off` builds Metal without the provider: ~124 MB smaller, and
|
||||
96.4% of MLX-LM instead of 99.1%.
|
||||
- `MLX_VERSION` pins the wheel (default `0.29.4`). MLX is consumed as the
|
||||
prebuilt pip wheel because building it from source needs `xcrun metal`, i.e. a
|
||||
full Xcode the macOS runners do not have.
|
||||
|
||||
Packaging vendors `libmlx.dylib`, `mlx.metallib` and MLX's MIT license into
|
||||
`package/lib/`, and rewrites `libvllm.dylib`'s rpath to `@loader_path/lib`
|
||||
(re-signing it, since `install_name_tool` invalidates the signature). The
|
||||
metallib must stay beside `libmlx.dylib`: MLX looks for it there.
|
||||
|
||||
Testing: `make test` runs the unit specs; export `VLLM_CPP_MODEL=<model>` (and
|
||||
optionally `VLLM_CPP_LIBRARY=<libvllm path>`) to enable the e2e specs.
|
||||
|
||||
@@ -109,16 +109,6 @@ func (v *VllmCpp) Load(opts *pb.ModelOptions) error {
|
||||
|
||||
v.opts = parseOptions(opts)
|
||||
|
||||
// A DFlash draft is a second checkpoint the engine opens by path, and the
|
||||
// engine never downloads one. Resolve it against LocalAI's models directory
|
||||
// now so a repo-id spelling works, and so a missing draft fails here with an
|
||||
// actionable message rather than as an HF-cache miss inside the load.
|
||||
resolvedSpec, err := resolveDraftModelPath(v.opts.speculativeConfig, opts.ModelPath)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
v.opts.speculativeConfig = resolvedSpec
|
||||
|
||||
mp := defaultModelParams()
|
||||
if v.opts.blockSize > 0 {
|
||||
mp.BlockSize = v.opts.blockSize
|
||||
@@ -126,62 +116,34 @@ func (v *VllmCpp) Load(opts *pb.ModelOptions) error {
|
||||
if v.opts.numBlocks > 0 {
|
||||
mp.NumBlocks = v.opts.numBlocks
|
||||
}
|
||||
// Sequence-length precedence, narrowest source last: context_size is the
|
||||
// generic LocalAI knob every backend honours, max_model_len is the
|
||||
// vLLM-specific one, and engine_args.max_model_len is the explicit
|
||||
// vllm-cpp override.
|
||||
if opts.ContextSize > 0 {
|
||||
mp.MaxModelLen = opts.ContextSize
|
||||
}
|
||||
if opts.MaxModelLen > 0 {
|
||||
mp.MaxModelLen = opts.MaxModelLen
|
||||
}
|
||||
if v.opts.maxModelLen > 0 {
|
||||
mp.MaxModelLen = v.opts.maxModelLen
|
||||
}
|
||||
if v.opts.maxNumSeqs > 0 {
|
||||
mp.MaxNumSeqs = v.opts.maxNumSeqs
|
||||
}
|
||||
if v.opts.maxNumBatchedTokens > 0 {
|
||||
mp.MaxNumBatchedTokens = v.opts.maxNumBatchedTokens
|
||||
}
|
||||
mp.EnablePrefixCaching = v.opts.enablePrefixCaching
|
||||
mp.EnableJumpForward = v.opts.enableJumpForward
|
||||
|
||||
// Every string below is borrowed by C for the duration of the load call
|
||||
// only (the library copies what it keeps), so the backing slices just have
|
||||
// to outlive vllmEngineLoad - hence the single KeepAlive after it.
|
||||
modelC := cString(model)
|
||||
mp.ModelPath = uintptr(unsafe.Pointer(&modelC[0])) // #nosec G103 -- borrowed by C for the load call only
|
||||
keep := [][]byte{modelC}
|
||||
setStr := func(dst *uintptr, s string) {
|
||||
if s == "" {
|
||||
return
|
||||
}
|
||||
b := cString(s)
|
||||
keep = append(keep, b)
|
||||
*dst = uintptr(unsafe.Pointer(&b[0])) // #nosec G103 -- borrowed by C for the load call only
|
||||
var toolParserC, reasoningParserC []byte
|
||||
if v.opts.toolParser != "" {
|
||||
toolParserC = cString(v.opts.toolParser)
|
||||
mp.ToolParser = uintptr(unsafe.Pointer(&toolParserC[0])) // #nosec G103 -- borrowed by C for the load call only
|
||||
}
|
||||
if v.opts.reasoningParser != "" {
|
||||
reasoningParserC = cString(v.opts.reasoningParser)
|
||||
mp.ReasoningParser = uintptr(unsafe.Pointer(&reasoningParserC[0])) // #nosec G103 -- borrowed by C for the load call only
|
||||
}
|
||||
setStr(&mp.ToolParser, v.opts.toolParser)
|
||||
setStr(&mp.ReasoningParser, v.opts.reasoningParser)
|
||||
setStr(&mp.SpeculativeConfig, v.opts.speculativeConfig)
|
||||
setStr(&mp.KVTransferConfig, v.opts.kvTransferConfig)
|
||||
setStr(&mp.SchedulingPolicy, v.opts.schedulingPolicy)
|
||||
setStr(&mp.TokenizerConfigPath, v.opts.tokenizerConfigPath)
|
||||
|
||||
xlog.Info("[vllm-cpp] Load", "model", model, "engine", vllmVersion(),
|
||||
"blockSize", mp.BlockSize, "numBlocks", mp.NumBlocks,
|
||||
"maxModelLen", mp.MaxModelLen, "maxNumSeqs", mp.MaxNumSeqs,
|
||||
"maxNumBatchedTokens", mp.MaxNumBatchedTokens,
|
||||
"prefixCaching", triStateName(mp.EnablePrefixCaching),
|
||||
"jumpForward", triStateName(mp.EnableJumpForward),
|
||||
"schedulingPolicy", v.opts.schedulingPolicy,
|
||||
"speculativeConfig", v.opts.speculativeConfig,
|
||||
"kvTransferConfig", v.opts.kvTransferConfig)
|
||||
"maxModelLen", mp.MaxModelLen, "maxNumSeqs", mp.MaxNumSeqs)
|
||||
|
||||
var engine uintptr
|
||||
rc := vllmEngineLoad(unsafe.Pointer(&mp), unsafe.Pointer(&engine)) // #nosec G103 -- POD out-params
|
||||
runtime.KeepAlive(keep)
|
||||
runtime.KeepAlive(modelC)
|
||||
runtime.KeepAlive(toolParserC)
|
||||
runtime.KeepAlive(reasoningParserC)
|
||||
if rc != vllmOK {
|
||||
return fmt.Errorf("vllm-cpp: engine load failed: %s", vllmLastError())
|
||||
}
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
package main
|
||||
|
||||
// purego bindings for the vllm.cpp stable C ABI (include/vllm.h, ABI v10).
|
||||
// purego bindings for the vllm.cpp stable C ABI (include/vllm.h, ABI v2).
|
||||
//
|
||||
// The structs below are hand-mirrored PODs of the C declarations, with
|
||||
// explicit padding so the Go layout matches the C layout on linux/darwin
|
||||
@@ -17,65 +17,29 @@ import (
|
||||
"github.com/ebitengine/purego"
|
||||
)
|
||||
|
||||
// abiVersion is the VLLM_ABI_VERSION this file mirrors (vllm.h). It must track
|
||||
// the header of the VLLM_CPP_VERSION pinned in the Makefile: the build checks
|
||||
// the two against each other, because a mismatch is only caught at runtime by
|
||||
// registerLib, where it takes the backend down on every load (issue #11379).
|
||||
const abiVersion = 10
|
||||
|
||||
// The ABI's tri-state toggles (enable_prefix_caching ABI v7,
|
||||
// enable_jump_forward ABI v10) share one encoding: 0 is NOT "off", it is
|
||||
// "defer" - to the model capability for prefix caching, to the environment for
|
||||
// jump forward. Only 2 is an explicit off.
|
||||
const (
|
||||
triStateDefer int32 = 0
|
||||
triStateOn int32 = 1
|
||||
triStateOff int32 = 2
|
||||
)
|
||||
|
||||
// triStateName renders a tri-state for the load log line, where "0" would
|
||||
// otherwise read as "off" rather than "whatever the default resolves to".
|
||||
func triStateName(state int32) string {
|
||||
switch state {
|
||||
case triStateOn:
|
||||
return "on"
|
||||
case triStateOff:
|
||||
return "off"
|
||||
default:
|
||||
return "model-default"
|
||||
}
|
||||
}
|
||||
// abiVersion is the VLLM_ABI_VERSION this file mirrors (vllm.h).
|
||||
const abiVersion = 5
|
||||
|
||||
// vllm_status (vllm.h).
|
||||
const (
|
||||
vllmOK = 0
|
||||
)
|
||||
|
||||
// cModelParams mirrors vllm_model_params. The int32 fields sit in pairs so the
|
||||
// interior needs no padding on LP64, but the struct is 8-aligned (it holds
|
||||
// pointers) and ends on a lone int32, so the trailing pad is explicit. Offsets
|
||||
// and total size are asserted in vllmcpp_test.go.
|
||||
// cModelParams mirrors vllm_model_params.
|
||||
type cModelParams struct {
|
||||
ModelPath uintptr // const char*
|
||||
TokenizerConfigPath uintptr // const char*; NULL = <model_dir>/... (ABI v9)
|
||||
TokenizerConfigPath uintptr // const char*
|
||||
BlockSize int32
|
||||
NumBlocks int32
|
||||
MaxModelLen int32
|
||||
MaxNumSeqs int32
|
||||
ToolParser uintptr // const char*; NULL = auto-detect (ABI v4)
|
||||
ReasoningParser uintptr // const char*; NULL = auto-detect (ABI v5)
|
||||
SpeculativeConfig uintptr // const char* JSON; NULL = no speculation (ABI v6)
|
||||
EnablePrefixCaching int32 // tri-state 0/1/2 (ABI v7)
|
||||
MaxNumBatchedTokens int32 // <= 0 = per-arch default (ABI v9)
|
||||
SchedulingPolicy uintptr // const char*; NULL = "fcfs" (ABI v9)
|
||||
KVTransferConfig uintptr // const char* JSON; NULL = no connector (ABI v9)
|
||||
EnableJumpForward int32 // tri-state 0/1/2 (ABI v10)
|
||||
_ [4]byte // trailing pad to the struct's 8-byte alignment
|
||||
}
|
||||
|
||||
// cSamplingParams mirrors vllm_sampling_params (structured fields included).
|
||||
// Padding matches the C compiler's: the uint64 seed is 8-aligned, and each
|
||||
// pointer following an int32 is 8-aligned.
|
||||
// cSamplingParams mirrors vllm_sampling_params (ABI v2, structured fields
|
||||
// included). Padding matches the C compiler's: the uint64 seed is 8-aligned,
|
||||
// and each pointer following an int32 is 8-aligned.
|
||||
type cSamplingParams struct {
|
||||
Temperature float32
|
||||
TopP float32
|
||||
@@ -101,12 +65,6 @@ type cSamplingParams struct {
|
||||
StructuredGrammar uintptr // const char*
|
||||
StructuredJSONObject int32
|
||||
_ [4]byte
|
||||
// ABI v8 tail. LocalAI installs no custom logits processor, but the fields
|
||||
// MUST be mirrored: the C side reads them off the pointer we hand it, so a
|
||||
// Go struct that stopped at StructuredJSONObject would have the engine read
|
||||
// 16 bytes past our allocation and call whatever garbage sat there.
|
||||
LogitsProcessor uintptr // vllm_logits_processor; NULL = none
|
||||
LogitsProcessorUserData uintptr // void*
|
||||
}
|
||||
|
||||
// cCompletion mirrors vllm_completion.
|
||||
|
||||
@@ -1,80 +1,30 @@
|
||||
package main
|
||||
|
||||
// Load-time engine configuration, from two config surfaces:
|
||||
//
|
||||
// - `engine_args:` (ModelOptions.EngineArgs, a JSON object) is the canonical
|
||||
// one. Keys are spelled exactly as vLLM's own CLI flags, so a config written
|
||||
// against vLLM works verbatim here - `speculative_config` and
|
||||
// `kv_transfer_config` in particular take the same JSON documents vLLM's
|
||||
// --speculative-config / --kv-transfer-config accept, and are handed to the
|
||||
// engine unparsed.
|
||||
// - `options:` (the free-form "key:value" list) is the older surface this
|
||||
// backend shipped with. It is still honoured so existing configs keep
|
||||
// working; engine_args wins on any key set in both.
|
||||
//
|
||||
// Anything unrecognised is ignored rather than fatal: the engine validates the
|
||||
// documents it is given and reports a precise error at load, and a config that
|
||||
// also carries knobs for a different backend must not fail the load here.
|
||||
// Engine-sizing knobs carried through the model config's free-form
|
||||
// `options:` list ("key:value" entries), mirroring how the other in-house
|
||||
// backends pass engine-specific settings that have no proto field.
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
"path"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
pb "github.com/mudler/LocalAI/pkg/grpc/proto"
|
||||
"github.com/mudler/xlog"
|
||||
)
|
||||
|
||||
type loadOptions struct {
|
||||
blockSize int32 // KV block size (tokens/block); engine default 32.
|
||||
numBlocks int32 // KV blocks to allocate; engine default 256.
|
||||
maxNumSeqs int32 // max concurrent sequences; engine default 8.
|
||||
// Max sequence length. Also settable through the model config's
|
||||
// context_size / max_model_len; see Load for the precedence.
|
||||
maxModelLen int32
|
||||
// Per-step chunked-prefill token budget (ABI v9). 0 = the engine's
|
||||
// bounded per-arch default.
|
||||
maxNumBatchedTokens int32
|
||||
// Automatic prefix caching tri-state (ABI v7): 0 = the model-capability
|
||||
// default, 1 = force on, 2 = force off.
|
||||
enablePrefixCaching int32
|
||||
// Jump-forward decoding tri-state (ABI v10), SGLang's grammar-speed subset:
|
||||
// 0 = defer to the environment (VT_ENABLE_JUMP_FORWARD, default off),
|
||||
// 1 = force on, 2 = force off.
|
||||
enableJumpForward int32
|
||||
// Scheduler admission policy (ABI v9): "" = fcfs, else fcfs|priority|lpm.
|
||||
schedulingPolicy string
|
||||
// Engine-side parser selection (ABI v4/v5). Empty = the engine
|
||||
// auto-detects from the chat template; "none" disables the reasoning
|
||||
// split; unknown names fail the first chat call.
|
||||
toolParser string
|
||||
reasoningParser string
|
||||
// Speculative decoding (ABI v6), as vLLM's --speculative-config JSON:
|
||||
// {"method":"mtp"|"dflash"|"ngram", ...}. Empty = no speculation.
|
||||
speculativeConfig string
|
||||
// External KV connector / LMCache (ABI v9), as vLLM's --kv-transfer-config
|
||||
// JSON. Empty = no connector.
|
||||
kvTransferConfig string
|
||||
// Override for the tokenizer_config.json the chat template is read from
|
||||
// (ABI v9). Empty = <model_dir>/tokenizer_config.json.
|
||||
tokenizerConfigPath string
|
||||
}
|
||||
|
||||
func parseOptions(opts *pb.ModelOptions) loadOptions {
|
||||
lo := loadOptions{}
|
||||
applyOptionsList(&lo, opts.GetOptions())
|
||||
applyEngineArgs(&lo, opts.GetEngineArgs())
|
||||
return lo
|
||||
}
|
||||
|
||||
// applyOptionsList reads the legacy free-form "key:value" list. strings.Cut
|
||||
// splits on the FIRST colon only, so a JSON object value survives intact.
|
||||
func applyOptionsList(lo *loadOptions, options []string) {
|
||||
for _, o := range options {
|
||||
for _, o := range opts.GetOptions() {
|
||||
k, v, found := strings.Cut(o, ":")
|
||||
if !found {
|
||||
continue
|
||||
@@ -86,211 +36,13 @@ func applyOptionsList(lo *loadOptions, options []string) {
|
||||
lo.numBlocks = parseInt32(v, lo.numBlocks)
|
||||
case "max_num_seqs":
|
||||
lo.maxNumSeqs = parseInt32(v, lo.maxNumSeqs)
|
||||
case "max_num_batched_tokens":
|
||||
lo.maxNumBatchedTokens = parseInt32(v, lo.maxNumBatchedTokens)
|
||||
case "max_model_len":
|
||||
lo.maxModelLen = parseInt32(v, lo.maxModelLen)
|
||||
case "scheduling_policy", "schedule_policy":
|
||||
lo.schedulingPolicy = strings.TrimSpace(v)
|
||||
case "tool_parser", "tool_call_parser":
|
||||
case "tool_parser":
|
||||
lo.toolParser = strings.TrimSpace(v)
|
||||
case "reasoning_parser":
|
||||
lo.reasoningParser = strings.TrimSpace(v)
|
||||
case "speculative_config":
|
||||
lo.speculativeConfig = strings.TrimSpace(v)
|
||||
case "kv_transfer_config":
|
||||
lo.kvTransferConfig = strings.TrimSpace(v)
|
||||
case "tokenizer_config", "tokenizer_config_path":
|
||||
lo.tokenizerConfigPath = strings.TrimSpace(v)
|
||||
case "enable_prefix_caching", "enable_radix_attention":
|
||||
if b, err := strconv.ParseBool(strings.TrimSpace(v)); err == nil {
|
||||
lo.enablePrefixCaching = boolTriState(b)
|
||||
}
|
||||
case "enable_jump_forward":
|
||||
if b, err := strconv.ParseBool(strings.TrimSpace(v)); err == nil {
|
||||
lo.enableJumpForward = boolTriState(b)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// applyEngineArgs overlays the `engine_args:` JSON object. A document that does
|
||||
// not parse is logged and skipped: engine_args is shared with the other engines
|
||||
// (the vLLM and SGLang backends read the same field), so a stray key must not
|
||||
// take the model down.
|
||||
func applyEngineArgs(lo *loadOptions, engineArgs string) {
|
||||
if strings.TrimSpace(engineArgs) == "" {
|
||||
return
|
||||
}
|
||||
var args map[string]any
|
||||
if err := json.Unmarshal([]byte(engineArgs), &args); err != nil {
|
||||
xlog.Warn("[vllm-cpp] ignoring unparseable engine_args", "error", err)
|
||||
return
|
||||
}
|
||||
for k, v := range args {
|
||||
switch k {
|
||||
case "block_size":
|
||||
lo.blockSize = jsonInt32(v, lo.blockSize)
|
||||
case "num_blocks":
|
||||
lo.numBlocks = jsonInt32(v, lo.numBlocks)
|
||||
case "max_num_seqs":
|
||||
lo.maxNumSeqs = jsonInt32(v, lo.maxNumSeqs)
|
||||
case "max_num_batched_tokens":
|
||||
lo.maxNumBatchedTokens = jsonInt32(v, lo.maxNumBatchedTokens)
|
||||
case "max_model_len":
|
||||
lo.maxModelLen = jsonInt32(v, lo.maxModelLen)
|
||||
case "scheduling_policy", "schedule_policy":
|
||||
lo.schedulingPolicy = jsonString(v, lo.schedulingPolicy)
|
||||
case "tool_parser", "tool_call_parser":
|
||||
lo.toolParser = jsonString(v, lo.toolParser)
|
||||
case "reasoning_parser":
|
||||
lo.reasoningParser = jsonString(v, lo.reasoningParser)
|
||||
case "tokenizer_config", "tokenizer_config_path":
|
||||
lo.tokenizerConfigPath = jsonString(v, lo.tokenizerConfigPath)
|
||||
case "speculative_config":
|
||||
lo.speculativeConfig = jsonDocument(v, lo.speculativeConfig, k)
|
||||
case "kv_transfer_config":
|
||||
lo.kvTransferConfig = jsonDocument(v, lo.kvTransferConfig, k)
|
||||
case "enable_prefix_caching", "enable_radix_attention":
|
||||
if b, ok := v.(bool); ok {
|
||||
lo.enablePrefixCaching = boolTriState(b)
|
||||
}
|
||||
case "enable_jump_forward":
|
||||
if b, ok := v.(bool); ok {
|
||||
lo.enableJumpForward = boolTriState(b)
|
||||
}
|
||||
default:
|
||||
xlog.Debug("[vllm-cpp] ignoring unknown engine_args key", "key", k)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// boolTriState maps a YAML/JSON boolean onto the ABI's tri-state encoding. An
|
||||
// explicit `false` must reach the engine as force-OFF (2), NOT as the 0 that
|
||||
// means "defer". The difference is real in both directions: prefix caching
|
||||
// defaults ON for dense archs and OFF for hybrid ones, and jump forward defers
|
||||
// to VT_ENABLE_JUMP_FORWARD.
|
||||
func boolTriState(on bool) int32 {
|
||||
if on {
|
||||
return triStateOn
|
||||
}
|
||||
return triStateOff
|
||||
}
|
||||
|
||||
// jsonDocument normalises an object-valued engine_args entry to a JSON string
|
||||
// for the C ABI. YAML nesting arrives as a map (the natural spelling); a
|
||||
// pre-encoded JSON string is accepted too, since a config round-tripped through
|
||||
// a flat store may carry it that way.
|
||||
func jsonDocument(v any, fallback string, key string) string {
|
||||
switch t := v.(type) {
|
||||
case string:
|
||||
if strings.TrimSpace(t) == "" {
|
||||
return fallback
|
||||
}
|
||||
return t
|
||||
default:
|
||||
buf, err := json.Marshal(t)
|
||||
if err != nil {
|
||||
xlog.Warn("[vllm-cpp] ignoring unencodable engine_args value", "key", key, "error", err)
|
||||
return fallback
|
||||
}
|
||||
return string(buf)
|
||||
}
|
||||
}
|
||||
|
||||
func jsonString(v any, fallback string) string {
|
||||
s, ok := v.(string)
|
||||
if !ok {
|
||||
return fallback
|
||||
}
|
||||
return strings.TrimSpace(s)
|
||||
}
|
||||
|
||||
// jsonInt32 accepts the float64 a JSON number decodes to, plus the string
|
||||
// spelling a YAML config may produce. Non-positive values keep the fallback:
|
||||
// every knob this covers uses "<= 0 means the engine default".
|
||||
func jsonInt32(v any, fallback int32) int32 {
|
||||
switch t := v.(type) {
|
||||
case float64:
|
||||
if t <= 0 || t > 1<<31-1 {
|
||||
return fallback
|
||||
}
|
||||
return int32(t)
|
||||
case string:
|
||||
return parseInt32(t, fallback)
|
||||
default:
|
||||
return fallback
|
||||
}
|
||||
}
|
||||
|
||||
// resolveDraftModelPath rewrites a DFlash draft reference into an absolute path
|
||||
// the engine can actually open.
|
||||
//
|
||||
// The engine resolves `speculative_config.model` against a directory containing
|
||||
// config.json, or against ~/.cache/huggingface/hub/models--<org>--<repo>/
|
||||
// snapshots/* - and it NEVER downloads. LocalAI keeps models in its own
|
||||
// directory, so a bare HF repo id (the spelling the vLLM docs teach) misses the
|
||||
// HF cache and dies deep in the load with "draft checkpoint not found", which
|
||||
// reads like a broken checkpoint rather than a missing download.
|
||||
//
|
||||
// So: try the reference as given, then the last path segment under the models
|
||||
// dir (`z-lab/Qwen3.6-27B-DFlash` -> `<models>/Qwen3.6-27B-DFlash`, which is
|
||||
// what LocalAI's own downloader produces), then the whole reference under the
|
||||
// models dir. If none exist, fail HERE with a message naming both what was
|
||||
// asked for and where we looked.
|
||||
//
|
||||
// mtp and ngram carry no separate draft checkpoint, so they pass through. A
|
||||
// document that does not parse also passes through: the engine owns config
|
||||
// validation and produces the better error.
|
||||
func resolveDraftModelPath(speculativeConfig, modelsDir string) (string, error) {
|
||||
if strings.TrimSpace(speculativeConfig) == "" {
|
||||
return speculativeConfig, nil
|
||||
}
|
||||
var spec map[string]any
|
||||
if err := json.Unmarshal([]byte(speculativeConfig), &spec); err != nil {
|
||||
return speculativeConfig, nil
|
||||
}
|
||||
if method, _ := spec["method"].(string); !strings.EqualFold(method, "dflash") {
|
||||
return speculativeConfig, nil
|
||||
}
|
||||
|
||||
ref, _ := spec["model"].(string)
|
||||
ref = strings.TrimSpace(ref)
|
||||
if ref == "" {
|
||||
return "", fmt.Errorf(
|
||||
"vllm-cpp: speculative_config method %q requires a \"model\" key naming the draft checkpoint", "dflash")
|
||||
}
|
||||
|
||||
candidates := []string{ref}
|
||||
if modelsDir != "" {
|
||||
if base := path.Base(filepath.ToSlash(ref)); base != "" && base != "." && base != "/" {
|
||||
candidates = append(candidates, filepath.Join(modelsDir, base))
|
||||
}
|
||||
candidates = append(candidates, filepath.Join(modelsDir, filepath.FromSlash(ref)))
|
||||
}
|
||||
|
||||
for _, c := range candidates {
|
||||
if _, err := os.Stat(filepath.Join(c, "config.json")); err != nil {
|
||||
continue
|
||||
}
|
||||
abs, err := filepath.Abs(c)
|
||||
if err != nil {
|
||||
abs = c
|
||||
}
|
||||
spec["model"] = abs
|
||||
out, err := json.Marshal(spec)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("vllm-cpp: re-encoding speculative_config: %w", err)
|
||||
}
|
||||
xlog.Info("[vllm-cpp] resolved DFlash draft checkpoint", "reference", ref, "path", abs)
|
||||
return string(out), nil
|
||||
}
|
||||
|
||||
return "", fmt.Errorf(
|
||||
"vllm-cpp: DFlash draft checkpoint %q not found (looked in: %s). "+
|
||||
"The engine does not download drafts - install the draft model into LocalAI first, "+
|
||||
"or set speculative_config.model to an absolute path to a directory containing config.json",
|
||||
ref, strings.Join(candidates, ", "))
|
||||
return lo
|
||||
}
|
||||
|
||||
func parseInt32(s string, fallback int32) int32 {
|
||||
|
||||
@@ -43,50 +43,6 @@ elif [ -f "/lib/ld-linux-aarch64.so.1" ]; then
|
||||
cp -arfLv /lib/aarch64-linux-gnu/libpthread.so.0 $CURDIR/package/lib/libpthread.so.0
|
||||
elif [ $(uname -s) = "Darwin" ]; then
|
||||
echo "Detected Darwin"
|
||||
# Vendor the optional MLX GEMM provider, when libvllm was built against it.
|
||||
# Three facts drive every line below, each verified on an Apple M4 before it
|
||||
# was written:
|
||||
# 1. libvllm.dylib carries an LC_LOAD_DYLIB on @rpath/libmlx.dylib, and its
|
||||
# build-time LC_RPATH points inside the build venv. That path does not
|
||||
# exist on a user's machine, so it must become @loader_path/lib.
|
||||
# 2. MLX finds its ~100 MB mlx.metallib beside its OWN dylib, so the two
|
||||
# files have to land in the same directory or every Metal op dies with
|
||||
# "Failed to load the default metallib".
|
||||
# 3. install_name_tool invalidates the code signature, and macOS refuses to
|
||||
# load an arm64 image whose signature does not match, so the patched
|
||||
# library must be re-signed ad-hoc afterwards.
|
||||
if otool -L "$CURDIR/package/libvllm.dylib" 2>/dev/null | grep -q "libmlx.dylib"; then
|
||||
MLX_LIB_DIR="${MLX_ROOT}/lib"
|
||||
if [ ! -f "$MLX_LIB_DIR/libmlx.dylib" ] || [ ! -f "$MLX_LIB_DIR/mlx.metallib" ]; then
|
||||
echo "Error: libvllm.dylib links libmlx.dylib but $MLX_LIB_DIR is missing libmlx.dylib/mlx.metallib" >&2
|
||||
exit 1
|
||||
fi
|
||||
echo "Vendoring the MLX GEMM provider from $MLX_LIB_DIR"
|
||||
cp -fLv "$MLX_LIB_DIR/libmlx.dylib" "$CURDIR/package/lib/"
|
||||
cp -fLv "$MLX_LIB_DIR/mlx.metallib" "$CURDIR/package/lib/"
|
||||
# MLX is MIT and we redistribute its binaries, so its license ships with
|
||||
# them. mlx-metal is the wheel carrying the dylib and the metallib.
|
||||
MLX_LICENSE=$(ls "${MLX_ROOT}"/../mlx_metal-*.dist-info/licenses/LICENSE 2>/dev/null | head -1)
|
||||
if [ -z "$MLX_LICENSE" ]; then
|
||||
MLX_LICENSE=$(ls "${MLX_ROOT}"/../mlx-*.dist-info/licenses/LICENSE 2>/dev/null | head -1)
|
||||
fi
|
||||
if [ -z "$MLX_LICENSE" ]; then
|
||||
echo "Error: could not find the MLX LICENSE to redistribute alongside libmlx.dylib" >&2
|
||||
exit 1
|
||||
fi
|
||||
cp -fLv "$MLX_LICENSE" "$CURDIR/package/lib/LICENSE.mlx"
|
||||
# Drop every build-tree rpath, then point at the packaged copy.
|
||||
otool -l "$CURDIR/package/libvllm.dylib" | awk '/LC_RPATH/{f=1;next} f&&/ path /{print $2;f=0}' | while read -r rp; do
|
||||
install_name_tool -delete_rpath "$rp" "$CURDIR/package/libvllm.dylib" 2>/dev/null || true
|
||||
done
|
||||
install_name_tool -add_rpath "@loader_path/lib" "$CURDIR/package/libvllm.dylib"
|
||||
codesign -f -s - "$CURDIR/package/libvllm.dylib"
|
||||
# A broken rpath must fail the BUILD, not the user's first inference.
|
||||
if ! otool -l "$CURDIR/package/libvllm.dylib" | grep -q "@loader_path/lib"; then
|
||||
echo "Error: libvllm.dylib did not get the @loader_path/lib rpath" >&2
|
||||
exit 1
|
||||
fi
|
||||
fi
|
||||
else
|
||||
echo "Error: Could not detect architecture"
|
||||
exit 1
|
||||
|
||||
@@ -16,17 +16,10 @@ func TestVllmCpp(t *testing.T) {
|
||||
RunSpecs(t, "vllm-cpp suite")
|
||||
}
|
||||
|
||||
// The Go POD mirrors must match the C struct layout of vllm.h (ABI v10)
|
||||
// The Go POD mirrors must match the C struct layout of vllm.h (ABI v2)
|
||||
// byte-for-byte: these offsets are the C offsets on LP64 (linux/darwin
|
||||
// amd64+arm64). A failure here means govllmcpp.go drifted from vllm.h.
|
||||
var _ = Describe("C ABI struct mirrors", func() {
|
||||
It("declares the ABI version the pinned engine reports", func() {
|
||||
// VLLM_ABI_VERSION in the vllm.h of VLLM_CPP_VERSION (Makefile).
|
||||
// Moving the pin past this without growing the mirrors below ships a
|
||||
// backend that refuses every load at startup (issue #11379).
|
||||
Expect(abiVersion).To(Equal(10))
|
||||
})
|
||||
|
||||
It("cModelParams matches vllm_model_params", func() {
|
||||
var p cModelParams
|
||||
Expect(unsafe.Offsetof(p.ModelPath)).To(Equal(uintptr(0)))
|
||||
@@ -37,18 +30,10 @@ var _ = Describe("C ABI struct mirrors", func() {
|
||||
Expect(unsafe.Offsetof(p.MaxNumSeqs)).To(Equal(uintptr(28)))
|
||||
Expect(unsafe.Offsetof(p.ToolParser)).To(Equal(uintptr(32)))
|
||||
Expect(unsafe.Offsetof(p.ReasoningParser)).To(Equal(uintptr(40)))
|
||||
Expect(unsafe.Offsetof(p.SpeculativeConfig)).To(Equal(uintptr(48)))
|
||||
Expect(unsafe.Offsetof(p.EnablePrefixCaching)).To(Equal(uintptr(56)))
|
||||
Expect(unsafe.Offsetof(p.MaxNumBatchedTokens)).To(Equal(uintptr(60)))
|
||||
Expect(unsafe.Offsetof(p.SchedulingPolicy)).To(Equal(uintptr(64)))
|
||||
Expect(unsafe.Offsetof(p.KVTransferConfig)).To(Equal(uintptr(72)))
|
||||
Expect(unsafe.Offsetof(p.EnableJumpForward)).To(Equal(uintptr(80)))
|
||||
// 88, not 84: the struct is 8-aligned (it holds pointers), so the
|
||||
// trailing int32 is padded out. Go pads identically.
|
||||
Expect(unsafe.Sizeof(p)).To(Equal(uintptr(88)))
|
||||
Expect(unsafe.Sizeof(p)).To(Equal(uintptr(48)))
|
||||
})
|
||||
|
||||
It("cSamplingParams matches vllm_sampling_params (ABI v8)", func() {
|
||||
It("cSamplingParams matches vllm_sampling_params (ABI v2)", func() {
|
||||
var p cSamplingParams
|
||||
Expect(unsafe.Offsetof(p.Temperature)).To(Equal(uintptr(0)))
|
||||
Expect(unsafe.Offsetof(p.TopP)).To(Equal(uintptr(4)))
|
||||
@@ -70,9 +55,7 @@ var _ = Describe("C ABI struct mirrors", func() {
|
||||
Expect(unsafe.Offsetof(p.NStructuredChoice)).To(Equal(uintptr(96)))
|
||||
Expect(unsafe.Offsetof(p.StructuredGrammar)).To(Equal(uintptr(104)))
|
||||
Expect(unsafe.Offsetof(p.StructuredJSONObject)).To(Equal(uintptr(112)))
|
||||
Expect(unsafe.Offsetof(p.LogitsProcessor)).To(Equal(uintptr(120)))
|
||||
Expect(unsafe.Offsetof(p.LogitsProcessorUserData)).To(Equal(uintptr(128)))
|
||||
Expect(unsafe.Sizeof(p)).To(Equal(uintptr(136)))
|
||||
Expect(unsafe.Sizeof(p)).To(Equal(uintptr(120)))
|
||||
})
|
||||
|
||||
It("cCompletion matches vllm_completion", func() {
|
||||
@@ -85,23 +68,6 @@ var _ = Describe("C ABI struct mirrors", func() {
|
||||
})
|
||||
})
|
||||
|
||||
// Pin/mirror skew is the failure mode this backend is most exposed to: the Go
|
||||
// PODs above are hand-written against one VLLM_ABI_VERSION, and the Makefile
|
||||
// pins the vllm.cpp commit that produces it. This spec catches drift without
|
||||
// needing model weights - set VLLM_CPP_LIBRARY to a built libvllm and it binds
|
||||
// every symbol and compares the library's reported ABI against the mirrors'.
|
||||
var _ = Describe("real library ABI handshake", func() {
|
||||
It("binds every symbol and reports the ABI the mirrors were written against", func() {
|
||||
lib := os.Getenv("VLLM_CPP_LIBRARY")
|
||||
if lib == "" {
|
||||
Skip("VLLM_CPP_LIBRARY not set; skipping the real-library handshake")
|
||||
}
|
||||
Expect(registerLib(lib)).To(Succeed())
|
||||
Expect(vllmABIVersion()).To(Equal(int32(abiVersion)))
|
||||
Expect(vllmVersion()).NotTo(BeEmpty())
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("parseOptions", func() {
|
||||
It("extracts the engine sizing knobs", func() {
|
||||
lo := parseOptions(&pb.ModelOptions{Options: []string{
|
||||
@@ -117,129 +83,6 @@ var _ = Describe("parseOptions", func() {
|
||||
}})
|
||||
Expect(lo).To(Equal(loadOptions{}))
|
||||
})
|
||||
|
||||
It("carries a speculative_config JSON value through the legacy options list", func() {
|
||||
// strings.Cut splits on the FIRST colon only, so a JSON object value
|
||||
// survives the "key:value" spelling intact.
|
||||
lo := parseOptions(&pb.ModelOptions{Options: []string{
|
||||
`speculative_config:{"method":"mtp","num_speculative_tokens":1}`,
|
||||
}})
|
||||
Expect(lo.speculativeConfig).To(Equal(`{"method":"mtp","num_speculative_tokens":1}`))
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("engine_args", func() {
|
||||
It("maps every load knob onto the C model params", func() {
|
||||
lo := parseOptions(&pb.ModelOptions{EngineArgs: `{
|
||||
"block_size": 64,
|
||||
"num_blocks": 1024,
|
||||
"max_model_len": 16384,
|
||||
"max_num_seqs": 32,
|
||||
"max_num_batched_tokens": 8192,
|
||||
"enable_prefix_caching": true,
|
||||
"scheduling_policy": "lpm",
|
||||
"tool_parser": "qwen3",
|
||||
"reasoning_parser": "deepseek_r1",
|
||||
"tokenizer_config": "/models/tok/tokenizer_config.json"
|
||||
}`})
|
||||
Expect(lo.blockSize).To(Equal(int32(64)))
|
||||
Expect(lo.numBlocks).To(Equal(int32(1024)))
|
||||
Expect(lo.maxModelLen).To(Equal(int32(16384)))
|
||||
Expect(lo.maxNumSeqs).To(Equal(int32(32)))
|
||||
Expect(lo.maxNumBatchedTokens).To(Equal(int32(8192)))
|
||||
Expect(lo.enablePrefixCaching).To(Equal(int32(1)))
|
||||
Expect(lo.schedulingPolicy).To(Equal("lpm"))
|
||||
Expect(lo.toolParser).To(Equal("qwen3"))
|
||||
Expect(lo.reasoningParser).To(Equal("deepseek_r1"))
|
||||
Expect(lo.tokenizerConfigPath).To(Equal("/models/tok/tokenizer_config.json"))
|
||||
})
|
||||
|
||||
It("re-marshals a nested speculative_config object to JSON for the engine", func() {
|
||||
lo := parseOptions(&pb.ModelOptions{EngineArgs: `{
|
||||
"speculative_config": {"method": "mtp", "num_speculative_tokens": 1}
|
||||
}`})
|
||||
Expect(lo.speculativeConfig).To(MatchJSON(`{"method":"mtp","num_speculative_tokens":1}`))
|
||||
})
|
||||
|
||||
It("re-marshals a nested kv_transfer_config object (LMCache) to JSON", func() {
|
||||
lo := parseOptions(&pb.ModelOptions{EngineArgs: `{
|
||||
"kv_transfer_config": {
|
||||
"kv_connector": "LMCacheConnector",
|
||||
"kv_role": "kv_both",
|
||||
"kv_connector_extra_config": {"host": "127.0.0.1", "port": 65432}
|
||||
}
|
||||
}`})
|
||||
Expect(lo.kvTransferConfig).To(MatchJSON(`{
|
||||
"kv_connector":"LMCacheConnector",
|
||||
"kv_role":"kv_both",
|
||||
"kv_connector_extra_config":{"host":"127.0.0.1","port":65432}
|
||||
}`))
|
||||
})
|
||||
|
||||
It("accepts a pre-encoded JSON string for the object-valued knobs", func() {
|
||||
// A config written by hand (or round-tripped through a flat store) may
|
||||
// carry the object as a string; both spellings reach the engine the same.
|
||||
lo := parseOptions(&pb.ModelOptions{EngineArgs: `{
|
||||
"speculative_config": "{\"method\":\"ngram\",\"num_speculative_tokens\":4}"
|
||||
}`})
|
||||
Expect(lo.speculativeConfig).To(MatchJSON(`{"method":"ngram","num_speculative_tokens":4}`))
|
||||
})
|
||||
|
||||
It("maps enable_prefix_caching false onto the force-OFF tri-state", func() {
|
||||
// The C ABI tri-state is 0=model default, 1=on, 2=off, so an explicit
|
||||
// `false` must NOT collapse to the 0 that means "let the model decide".
|
||||
lo := parseOptions(&pb.ModelOptions{EngineArgs: `{"enable_prefix_caching": false}`})
|
||||
Expect(lo.enablePrefixCaching).To(Equal(int32(2)))
|
||||
})
|
||||
|
||||
It("leaves the prefix-caching tri-state at the model default when unset", func() {
|
||||
lo := parseOptions(&pb.ModelOptions{EngineArgs: `{"max_num_seqs": 4}`})
|
||||
Expect(lo.enablePrefixCaching).To(Equal(int32(0)))
|
||||
})
|
||||
|
||||
It("accepts the radix-attention alias upstream documents for prefix caching", func() {
|
||||
lo := parseOptions(&pb.ModelOptions{EngineArgs: `{"enable_radix_attention": true}`})
|
||||
Expect(lo.enablePrefixCaching).To(Equal(int32(1)))
|
||||
})
|
||||
|
||||
It("maps enable_jump_forward onto its own tri-state", func() {
|
||||
// ABI v10. Same tri-state shape as prefix caching, and the same trap:
|
||||
// an explicit false must be force-OFF (2), not the 0 that defers to the
|
||||
// environment.
|
||||
on := parseOptions(&pb.ModelOptions{EngineArgs: `{"enable_jump_forward": true}`})
|
||||
Expect(on.enableJumpForward).To(Equal(int32(1)))
|
||||
off := parseOptions(&pb.ModelOptions{EngineArgs: `{"enable_jump_forward": false}`})
|
||||
Expect(off.enableJumpForward).To(Equal(int32(2)))
|
||||
unset := parseOptions(&pb.ModelOptions{EngineArgs: `{"max_num_seqs": 4}`})
|
||||
Expect(unset.enableJumpForward).To(Equal(int32(0)))
|
||||
})
|
||||
|
||||
It("reads enable_jump_forward from the legacy options list too", func() {
|
||||
lo := parseOptions(&pb.ModelOptions{Options: []string{"enable_jump_forward:true"}})
|
||||
Expect(lo.enableJumpForward).To(Equal(int32(1)))
|
||||
})
|
||||
|
||||
It("lets engine_args override the legacy options list", func() {
|
||||
lo := parseOptions(&pb.ModelOptions{
|
||||
Options: []string{"max_num_seqs:8", "block_size:16"},
|
||||
EngineArgs: `{"max_num_seqs": 64}`,
|
||||
})
|
||||
Expect(lo.maxNumSeqs).To(Equal(int32(64))) // engine_args wins
|
||||
Expect(lo.blockSize).To(Equal(int32(16))) // untouched keys survive
|
||||
})
|
||||
|
||||
It("ignores malformed engine_args rather than failing the load", func() {
|
||||
lo := parseOptions(&pb.ModelOptions{
|
||||
Options: []string{"max_num_seqs:8"},
|
||||
EngineArgs: `{not json`,
|
||||
})
|
||||
Expect(lo.maxNumSeqs).To(Equal(int32(8)))
|
||||
})
|
||||
|
||||
It("ignores unknown keys", func() {
|
||||
lo := parseOptions(&pb.ModelOptions{EngineArgs: `{"gpu_memory_utilization": 0.9}`})
|
||||
Expect(lo).To(Equal(loadOptions{}))
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("samplingFromPredict", func() {
|
||||
@@ -292,91 +135,6 @@ var _ = Describe("samplingFromPredict", func() {
|
||||
})
|
||||
})
|
||||
|
||||
// The engine resolves speculative_config.model against a local directory or
|
||||
// ~/.cache/huggingface/hub ONLY - it never downloads. LocalAI keeps models in
|
||||
// its own directory, so a bare repo id would miss the HF cache and fail deep in
|
||||
// the load with a confusing "draft checkpoint not found". Resolve it here.
|
||||
var _ = Describe("resolveDraftModelPath", func() {
|
||||
var modelsDir string
|
||||
|
||||
BeforeEach(func() {
|
||||
modelsDir = GinkgoT().TempDir()
|
||||
})
|
||||
|
||||
// draftDir creates a plausible draft checkpoint under models/.
|
||||
draftDir := func(name string) string {
|
||||
d := filepath.Join(modelsDir, name)
|
||||
Expect(os.MkdirAll(d, 0o750)).To(Succeed())
|
||||
Expect(os.WriteFile(filepath.Join(d, "config.json"), []byte("{}"), 0o600)).To(Succeed())
|
||||
return d
|
||||
}
|
||||
|
||||
It("rewrites a repo id to the matching directory in the models dir", func() {
|
||||
want := draftDir("Qwen3.6-27B-DFlash")
|
||||
spec := `{"method":"dflash","model":"z-lab/Qwen3.6-27B-DFlash"}`
|
||||
out, err := resolveDraftModelPath(spec, modelsDir)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(out).To(MatchJSON(`{"method":"dflash","model":"` + want + `"}`))
|
||||
})
|
||||
|
||||
It("rewrites a models-dir-relative path", func() {
|
||||
want := draftDir("drafts__dflash")
|
||||
spec := `{"method":"dflash","model":"drafts__dflash"}`
|
||||
out, err := resolveDraftModelPath(spec, modelsDir)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(out).To(ContainSubstring(want))
|
||||
})
|
||||
|
||||
It("leaves an absolute path that already resolves alone", func() {
|
||||
abs := draftDir("elsewhere")
|
||||
spec := `{"method":"dflash","model":"` + abs + `"}`
|
||||
out, err := resolveDraftModelPath(spec, modelsDir)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(out).To(MatchJSON(spec))
|
||||
})
|
||||
|
||||
It("fails with an actionable error when the draft is nowhere on disk", func() {
|
||||
// Silently passing the repo id through would surface as an HF-cache
|
||||
// miss inside the engine, which reads as "your model is broken".
|
||||
spec := `{"method":"dflash","model":"z-lab/Not-Downloaded"}`
|
||||
_, err := resolveDraftModelPath(spec, modelsDir)
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("z-lab/Not-Downloaded"))
|
||||
Expect(err.Error()).To(ContainSubstring(modelsDir))
|
||||
})
|
||||
|
||||
It("requires a model key for dflash", func() {
|
||||
_, err := resolveDraftModelPath(`{"method":"dflash"}`, modelsDir)
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("model"))
|
||||
})
|
||||
|
||||
It("leaves mtp and ngram configs untouched", func() {
|
||||
// Neither has a separate draft checkpoint to resolve.
|
||||
for _, spec := range []string{
|
||||
`{"method":"mtp"}`,
|
||||
`{"method":"ngram","num_speculative_tokens":4}`,
|
||||
} {
|
||||
out, err := resolveDraftModelPath(spec, modelsDir)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(out).To(MatchJSON(spec))
|
||||
}
|
||||
})
|
||||
|
||||
It("passes a malformed document through for the engine to reject", func() {
|
||||
// The engine owns config validation and produces the better message.
|
||||
out, err := resolveDraftModelPath(`{not json`, modelsDir)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(out).To(Equal(`{not json`))
|
||||
})
|
||||
|
||||
It("is a no-op on an empty config", func() {
|
||||
out, err := resolveDraftModelPath("", modelsDir)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(out).To(BeEmpty())
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("validModelPath", func() {
|
||||
It("accepts a .gguf file", func() {
|
||||
dir := GinkgoT().TempDir()
|
||||
|
||||
@@ -8,7 +8,7 @@ JOBS?=$(shell nproc --ignore=1)
|
||||
|
||||
# whisper.cpp version
|
||||
WHISPER_REPO?=https://github.com/ggml-org/whisper.cpp
|
||||
WHISPER_CPP_VERSION?=306c88f4d1286aec1bf96e544632897886af5501
|
||||
WHISPER_CPP_VERSION?=2ca53bb45e38748d07b310eeb36245a7157ac882
|
||||
SO_TARGET?=libgowhisper.so
|
||||
|
||||
CMAKE_ARGS+=-DBUILD_SHARED_LIBS=OFF
|
||||
|
||||
@@ -193,22 +193,12 @@
|
||||
alias: "vllm-cpp"
|
||||
license: apache-2.0
|
||||
description: |
|
||||
ALPHA development builds. Try it, but llama-cpp stays the recommendation for
|
||||
production use.
|
||||
|
||||
vllm.cpp is an Apache-2.0 C++20 inference engine maintained by the LocalAI team,
|
||||
developed in its own repository and usable without LocalAI. It began as a port of
|
||||
vLLM and keeps vLLM as its reference implementation, checking output against it and
|
||||
benchmarking against it, while growing a featureset of its own. It implements vLLM's
|
||||
V1 architecture (paged KV cache, continuous batching, prefix caching, scheduler,
|
||||
sampler) on a portable tensor runtime with no Python, PyTorch or ggml at inference
|
||||
time. It loads GGUF as well as Hugging Face safetensors, supports structured output
|
||||
(JSON schema / regex / choice / GBNF grammar) enforced in-engine, ships speculative
|
||||
decoding and KV offload, and runs on CPU, NVIDIA CUDA (Blackwell-family), Apple
|
||||
Metal and Vulkan.
|
||||
|
||||
The project is expected to be renamed as it diverges further from vLLM; the new
|
||||
name is still to be decided.
|
||||
vllm.cpp is a from-scratch C++20 port of vLLM created and maintained by the LocalAI team.
|
||||
It mirrors vLLM's V1 architecture (paged KV cache, continuous batching, prefix caching,
|
||||
scheduler, sampler) on a portable tensor runtime with no Python, PyTorch or ggml at
|
||||
inference time. It loads Hugging Face safetensors and GGUF checkpoints, supports
|
||||
structured output (JSON schema / regex / choice / GBNF grammar) enforced in-engine,
|
||||
and runs on CPU, NVIDIA CUDA (Blackwell-family), Apple Metal and Vulkan.
|
||||
urls:
|
||||
- https://github.com/mudler/vllm.cpp
|
||||
tags:
|
||||
@@ -292,55 +282,6 @@
|
||||
nvidia-cuda-12: "cuda12-parakeet-cpp"
|
||||
nvidia-l4t-cuda-12: "nvidia-l4t-arm64-parakeet-cpp"
|
||||
nvidia-l4t-cuda-13: "cuda13-nvidia-l4t-arm64-parakeet-cpp"
|
||||
- &nemospeechcpp
|
||||
name: "nemo-speech-cpp"
|
||||
alias: "nemo-speech-cpp"
|
||||
license: apache-2.0
|
||||
icon: https://avatars.githubusercontent.com/u/1728152?s=200&v=4
|
||||
description: |
|
||||
NVIDIA NeMo-Speech.cpp, a C++/ggml runtime for NVIDIA Nemotron Speech models.
|
||||
One backend serves four model families, selected automatically from the GGUF
|
||||
general.architecture key: automatic speech recognition (offline, cache-aware
|
||||
streaming and live transcription, with optional Silero VAD, punctuation,
|
||||
inverse text normalization and Sortformer speaker diarization attached),
|
||||
standalone Sortformer diarization, MagpieTTS text-to-speech over NanoCodec,
|
||||
and Riva-Translate text translation. Runs on CPU, NVIDIA CUDA, Vulkan,
|
||||
NVIDIA Jetson (L4T) and Apple Metal.
|
||||
urls:
|
||||
- https://github.com/NVIDIA/NeMo-Speech.cpp
|
||||
tags:
|
||||
- audio-transcription
|
||||
- text-to-speech
|
||||
- diarization
|
||||
- text-to-text
|
||||
- CPU
|
||||
- GPU
|
||||
- CUDA
|
||||
- Metal
|
||||
# No amd and no intel key on purpose: upstream NeMo-Speech.cpp has no ROCm/HIP
|
||||
# and no SYCL backend, so there is nothing to point those at. A host reporting
|
||||
# either capability falls through to "default" (SystemState.Capability) and
|
||||
# gets the CPU build, which is the honest answer rather than a broken tag.
|
||||
#
|
||||
# Listing only nvidia-l4t would be a silent downgrade: a Jetson that reports a
|
||||
# CUDA-refined capability would miss the map and fall back to the CPU build.
|
||||
#
|
||||
# The two nvidia-l4t-cuda-* keys point at DIFFERENT images on purpose. The
|
||||
# JetPack r36.4.0 base links ggml against CUDA 12, so serving it to a host that
|
||||
# reports nvidia-l4t-cuda-13 would fail at dlopen on a missing libcudart.so.12.
|
||||
# That is worse than no key at all, since a missing key falls back to a working
|
||||
# CPU build. Hence the separate cuda13 L4T image, as parakeet-cpp and
|
||||
# moss-transcribe-cpp both do.
|
||||
capabilities:
|
||||
default: "cpu-nemo-speech-cpp"
|
||||
nvidia: "cuda12-nemo-speech-cpp"
|
||||
metal: "metal-nemo-speech-cpp"
|
||||
vulkan: "vulkan-nemo-speech-cpp"
|
||||
nvidia-l4t: "nvidia-l4t-arm64-nemo-speech-cpp"
|
||||
nvidia-cuda-13: "cuda13-nemo-speech-cpp"
|
||||
nvidia-cuda-12: "cuda12-nemo-speech-cpp"
|
||||
nvidia-l4t-cuda-12: "nvidia-l4t-arm64-nemo-speech-cpp"
|
||||
nvidia-l4t-cuda-13: "cuda13-nvidia-l4t-arm64-nemo-speech-cpp"
|
||||
- &mosstranscribecpp
|
||||
name: "moss-transcribe-cpp"
|
||||
alias: "moss-transcribe-cpp"
|
||||
@@ -3323,89 +3264,6 @@
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-gpu-nvidia-cuda-13-parakeet-cpp"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-gpu-nvidia-cuda-13-parakeet-cpp
|
||||
## nemo-speech-cpp
|
||||
- !!merge <<: *nemospeechcpp
|
||||
name: "nemo-speech-cpp-development"
|
||||
capabilities:
|
||||
default: "cpu-nemo-speech-cpp-development"
|
||||
nvidia: "cuda12-nemo-speech-cpp-development"
|
||||
metal: "metal-nemo-speech-cpp-development"
|
||||
vulkan: "vulkan-nemo-speech-cpp-development"
|
||||
nvidia-l4t: "nvidia-l4t-arm64-nemo-speech-cpp-development"
|
||||
nvidia-cuda-13: "cuda13-nemo-speech-cpp-development"
|
||||
nvidia-cuda-12: "cuda12-nemo-speech-cpp-development"
|
||||
nvidia-l4t-cuda-12: "nvidia-l4t-arm64-nemo-speech-cpp-development"
|
||||
nvidia-l4t-cuda-13: "cuda13-nvidia-l4t-arm64-nemo-speech-cpp-development"
|
||||
- !!merge <<: *nemospeechcpp
|
||||
name: "cpu-nemo-speech-cpp"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-cpu-nemo-speech-cpp"
|
||||
mirrors:
|
||||
- localai/localai-backends:latest-cpu-nemo-speech-cpp
|
||||
- !!merge <<: *nemospeechcpp
|
||||
name: "cpu-nemo-speech-cpp-development"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-cpu-nemo-speech-cpp"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-cpu-nemo-speech-cpp
|
||||
- !!merge <<: *nemospeechcpp
|
||||
name: "cuda12-nemo-speech-cpp"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-gpu-nvidia-cuda-12-nemo-speech-cpp"
|
||||
mirrors:
|
||||
- localai/localai-backends:latest-gpu-nvidia-cuda-12-nemo-speech-cpp
|
||||
- !!merge <<: *nemospeechcpp
|
||||
name: "cuda12-nemo-speech-cpp-development"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-gpu-nvidia-cuda-12-nemo-speech-cpp"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-gpu-nvidia-cuda-12-nemo-speech-cpp
|
||||
- !!merge <<: *nemospeechcpp
|
||||
name: "cuda13-nemo-speech-cpp"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-gpu-nvidia-cuda-13-nemo-speech-cpp"
|
||||
mirrors:
|
||||
- localai/localai-backends:latest-gpu-nvidia-cuda-13-nemo-speech-cpp
|
||||
- !!merge <<: *nemospeechcpp
|
||||
name: "cuda13-nemo-speech-cpp-development"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-gpu-nvidia-cuda-13-nemo-speech-cpp"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-gpu-nvidia-cuda-13-nemo-speech-cpp
|
||||
- !!merge <<: *nemospeechcpp
|
||||
name: "vulkan-nemo-speech-cpp"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-gpu-vulkan-nemo-speech-cpp"
|
||||
mirrors:
|
||||
- localai/localai-backends:latest-gpu-vulkan-nemo-speech-cpp
|
||||
- !!merge <<: *nemospeechcpp
|
||||
name: "vulkan-nemo-speech-cpp-development"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-gpu-vulkan-nemo-speech-cpp"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-gpu-vulkan-nemo-speech-cpp
|
||||
- !!merge <<: *nemospeechcpp
|
||||
name: "nvidia-l4t-arm64-nemo-speech-cpp"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-nvidia-l4t-arm64-nemo-speech-cpp"
|
||||
mirrors:
|
||||
- localai/localai-backends:latest-nvidia-l4t-arm64-nemo-speech-cpp
|
||||
- !!merge <<: *nemospeechcpp
|
||||
name: "nvidia-l4t-arm64-nemo-speech-cpp-development"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-nvidia-l4t-arm64-nemo-speech-cpp"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-nvidia-l4t-arm64-nemo-speech-cpp
|
||||
- !!merge <<: *nemospeechcpp
|
||||
name: "cuda13-nvidia-l4t-arm64-nemo-speech-cpp"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-nvidia-l4t-cuda-13-arm64-nemo-speech-cpp"
|
||||
mirrors:
|
||||
- localai/localai-backends:latest-nvidia-l4t-cuda-13-arm64-nemo-speech-cpp
|
||||
- !!merge <<: *nemospeechcpp
|
||||
name: "cuda13-nvidia-l4t-arm64-nemo-speech-cpp-development"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-nvidia-l4t-cuda-13-arm64-nemo-speech-cpp"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-nvidia-l4t-cuda-13-arm64-nemo-speech-cpp
|
||||
- !!merge <<: *nemospeechcpp
|
||||
name: "metal-nemo-speech-cpp"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-metal-darwin-arm64-nemo-speech-cpp"
|
||||
mirrors:
|
||||
- localai/localai-backends:latest-metal-darwin-arm64-nemo-speech-cpp
|
||||
- !!merge <<: *nemospeechcpp
|
||||
name: "metal-nemo-speech-cpp-development"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-metal-darwin-arm64-nemo-speech-cpp"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-metal-darwin-arm64-nemo-speech-cpp
|
||||
## moss-transcribe-cpp
|
||||
- !!merge <<: *mosstranscribecpp
|
||||
name: "moss-transcribe-cpp-development"
|
||||
|
||||
@@ -883,34 +883,6 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
|
||||
return backend_pb2.Result(message="Media generated", success=True)
|
||||
|
||||
def UpscaleImage(self, request, context):
|
||||
try:
|
||||
if not request.src:
|
||||
return backend_pb2.Result(success=False, message="No source image provided")
|
||||
if not request.dst:
|
||||
return backend_pb2.Result(success=False, message="No destination path provided")
|
||||
|
||||
scale = request.scale if request.scale > 0 else 2
|
||||
image = Image.open(request.src).convert("RGB")
|
||||
|
||||
# If the loaded pipeline supports upscaling (e.g. StableDiffusionUpscalePipeline),
|
||||
# use it; otherwise fall back to high-quality Lanczos resize.
|
||||
if self.pipe is not None and self.PipelineType in ("StableDiffusionUpscalePipeline", "StableDiffusionLatentUpscalePipeline"):
|
||||
print(f"UpscaleImage: using diffusers upscale pipeline ({self.PipelineType})", file=sys.stderr)
|
||||
upscaled = self.pipe(prompt="", image=image).images[0]
|
||||
else:
|
||||
# Fallback: high-quality Lanczos resize
|
||||
print(f"UpscaleImage: no upscale pipeline loaded, using Lanczos resize (scale={scale})", file=sys.stderr)
|
||||
new_w = image.width * scale
|
||||
new_h = image.height * scale
|
||||
upscaled = image.resize((new_w, new_h), Image.LANCZOS)
|
||||
|
||||
upscaled.save(request.dst)
|
||||
return backend_pb2.Result(message="Image upscaled", success=True)
|
||||
except Exception as e:
|
||||
print(f"UpscaleImage error: {e}", file=sys.stderr)
|
||||
return backend_pb2.Result(success=False, message=str(e))
|
||||
|
||||
def GenerateVideo(self, request, context):
|
||||
try:
|
||||
prompt = request.prompt
|
||||
|
||||
@@ -15,12 +15,3 @@ sglang[all]>=0.5.11
|
||||
# load-bearing for flash-attn-4, and this is the narrower change. Raise the
|
||||
# bound once 0.46.0 final ships.
|
||||
nvidia-modelopt<0.46
|
||||
|
||||
# Same failure mode as the nvidia-modelopt bound above, via a different
|
||||
# package. sglang -> flashinfer-python -> cuda-tile, unbounded, and the
|
||||
# global --prerelease=allow resolves it to 1.6.0rc3, whose build backend
|
||||
# imports wheel_stub without declaring it in build-system.requires. With
|
||||
# --no-build-isolation nothing installs it and the build dies with
|
||||
# "No module named 'wheel_stub'". 1.5.0 is the newest stable release.
|
||||
# Raise the bound once 1.6.0 final ships.
|
||||
cuda-tile<1.6
|
||||
|
||||
@@ -15,12 +15,3 @@ sglang[all]>=0.5.11
|
||||
# load-bearing for flash-attn-4, and this is the narrower change. Raise the
|
||||
# bound once 0.46.0 final ships.
|
||||
nvidia-modelopt<0.46
|
||||
|
||||
# Same failure mode as the nvidia-modelopt bound above, via a different
|
||||
# package. sglang -> flashinfer-python -> cuda-tile, unbounded, and the
|
||||
# global --prerelease=allow resolves it to 1.6.0rc3, whose build backend
|
||||
# imports wheel_stub without declaring it in build-system.requires. With
|
||||
# --no-build-isolation nothing installs it and the build dies with
|
||||
# "No module named 'wheel_stub'". 1.5.0 is the newest stable release.
|
||||
# Raise the bound once 1.6.0 final ships.
|
||||
cuda-tile<1.6
|
||||
|
||||
@@ -13,12 +13,3 @@
|
||||
# FunctionCallParser, ReasoningParser); the [all] extras are optional
|
||||
# accelerators not required at import time.
|
||||
sglang>=0.5.11
|
||||
|
||||
# Same failure mode the cublas profiles carry an nvidia-modelopt bound for,
|
||||
# reached through a different package. sglang -> flashinfer-python ->
|
||||
# cuda-tile, unbounded, and the global --prerelease=allow resolves it to
|
||||
# 1.6.0rc3, whose build backend imports wheel_stub without declaring it in
|
||||
# build-system.requires. With --no-build-isolation nothing installs it and
|
||||
# the build dies with "No module named 'wheel_stub'". 1.5.0 is the newest
|
||||
# stable release. Raise the bound once 1.6.0 final ships.
|
||||
cuda-tile<1.6
|
||||
|
||||
@@ -1,7 +1,6 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"os"
|
||||
"path/filepath"
|
||||
|
||||
@@ -108,13 +107,6 @@ For documentation and support:
|
||||
// Run the thing!
|
||||
err = ctx.Run(&cli.CLI.Context)
|
||||
if err != nil {
|
||||
// A command that has already told the user what went wrong returns
|
||||
// only a status. Logging it as well would print a bare "exit status 1"
|
||||
// underneath the explanation they just read.
|
||||
var reported cli.ExitCodeError
|
||||
if errors.As(err, &reported) {
|
||||
os.Exit(reported.Code)
|
||||
}
|
||||
xlog.Fatal("Error running the application", "error", err)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -553,17 +553,12 @@ func (a *Application) start() error {
|
||||
// once at startup and reused across chat sessions that opt in via metadata.
|
||||
if !a.applicationConfig.DisableLocalAIAssistant {
|
||||
holder := mcpTools.NewLocalAIAssistantHolder()
|
||||
var nodeRegistry *nodes.NodeRegistry
|
||||
if a.distributed != nil {
|
||||
nodeRegistry = a.distributed.Registry
|
||||
}
|
||||
assistantClient := localaiInproc.New(
|
||||
a.applicationConfig,
|
||||
a.applicationConfig.SystemState,
|
||||
a.backendLoader,
|
||||
a.modelLoader,
|
||||
a.galleryService,
|
||||
nodeRegistry,
|
||||
)
|
||||
// Wire usage tracking so the assistant's get_usage_stats tool
|
||||
// returns real data; nil values keep the tool returning a clear
|
||||
|
||||
@@ -444,13 +444,6 @@ func New(opts ...config.AppOption) (*Application, error) {
|
||||
// when gallery data refreshes instead of using a fixed TTL.
|
||||
vram.SetGalleryGenerationFunc(gallery.GalleryGeneration)
|
||||
|
||||
// Fill those caches ahead of the first visitor. An estimate for an entry
|
||||
// nobody has asked about yet costs a remote probe of its weight files, and
|
||||
// the model gallery asks for one per row, so without this the first page
|
||||
// spends seconds filling in its own sizes while somebody watches it.
|
||||
// Non-blocking, and bounded: see DefaultEstimateWarmConfig.
|
||||
gallery.WarmEstimateCache(options.Context, options.Galleries, options.SystemState, gallery.EstimateWarmConfigFromEnv())
|
||||
|
||||
if options.ConfigFile != "" {
|
||||
if err := application.ModelConfigLoader().LoadMultipleModelConfigsSingleFile(options.ConfigFile, configLoaderOpts...); err != nil {
|
||||
xlog.Error("error loading config file", "error", err)
|
||||
|
||||
@@ -1,37 +0,0 @@
|
||||
package backend
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"github.com/mudler/LocalAI/core/config"
|
||||
"github.com/mudler/LocalAI/pkg/grpc/proto"
|
||||
model "github.com/mudler/LocalAI/pkg/model"
|
||||
)
|
||||
|
||||
// ImageUpscale loads the model specified in modelConfig and calls UpscaleImage
|
||||
// on the backend, writing the result to dst.
|
||||
func ImageUpscale(ctx context.Context, src, dst string, scale int, loader *model.ModelLoader, modelConfig config.ModelConfig, appConfig *config.ApplicationConfig) (func() error, error) {
|
||||
opts := ModelOptions(modelConfig, appConfig, model.WithContext(ctx))
|
||||
inferenceModel, err := loader.Load(opts...)
|
||||
if err != nil {
|
||||
recordModelLoadFailure(appConfig, modelConfig.Name, modelConfig.Backend, err, nil)
|
||||
return nil, err
|
||||
}
|
||||
|
||||
fn := func() error {
|
||||
_, err := inferenceModel.UpscaleImage(
|
||||
ctx,
|
||||
&proto.UpscaleImageRequest{
|
||||
Src: src,
|
||||
Dst: dst,
|
||||
Scale: int32(scale),
|
||||
},
|
||||
)
|
||||
return err
|
||||
}
|
||||
|
||||
return fn, nil
|
||||
}
|
||||
|
||||
// ImageUpscaleFunc is a test-friendly indirection.
|
||||
var ImageUpscaleFunc = ImageUpscale
|
||||
30
core/cli/chat/chat.go
Normal file
30
core/cli/chat/chat.go
Normal file
@@ -0,0 +1,30 @@
|
||||
package chat
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"strings"
|
||||
)
|
||||
|
||||
type Options struct {
|
||||
Model string
|
||||
BaseURL string
|
||||
APIKey string
|
||||
In io.Reader
|
||||
Out io.Writer
|
||||
}
|
||||
|
||||
func Run(ctx context.Context, opts Options) error {
|
||||
if opts.In == nil {
|
||||
opts.In = strings.NewReader("")
|
||||
}
|
||||
if opts.Out == nil {
|
||||
opts.Out = io.Discard
|
||||
}
|
||||
|
||||
session, err := newChatSession(ctx, newLocalAIChatClient(opts.BaseURL, opts.APIKey), opts.Model)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return runTerminalChat(ctx, session, opts.In, opts.Out)
|
||||
}
|
||||
172
core/cli/chat/chat_test.go
Normal file
172
core/cli/chat/chat_test.go
Normal file
@@ -0,0 +1,172 @@
|
||||
package chat
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
)
|
||||
|
||||
var _ = Describe("Run chat", func() {
|
||||
It("streams a single chat response", func() {
|
||||
var capturedModel string
|
||||
var capturedAuth string
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path == "/v1/models" {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
writeResponse(w, `{"object":"list","data":[{"id":"test-model","object":"model"}]}`)
|
||||
return
|
||||
}
|
||||
|
||||
Expect(r.URL.Path).To(Equal("/v1/chat/completions"))
|
||||
capturedAuth = r.Header.Get("Authorization")
|
||||
|
||||
var body struct {
|
||||
Model string `json:"model"`
|
||||
Messages []struct {
|
||||
Role string `json:"role"`
|
||||
Content string `json:"content"`
|
||||
} `json:"messages"`
|
||||
}
|
||||
Expect(json.NewDecoder(r.Body).Decode(&body)).To(Succeed())
|
||||
capturedModel = body.Model
|
||||
Expect(body.Messages).To(HaveLen(1))
|
||||
Expect(body.Messages[0].Role).To(Equal("user"))
|
||||
Expect(body.Messages[0].Content).To(Equal("hello"))
|
||||
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
writeResponse(w, "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"hi\"}}]}\n\n")
|
||||
writeResponse(w, "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"!\"}}]}\n\n")
|
||||
writeResponse(w, "data: [DONE]\n\n")
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
var out bytes.Buffer
|
||||
err := Run(GinkgoT().Context(), Options{
|
||||
Model: "test-model",
|
||||
BaseURL: server.URL + "/v1",
|
||||
APIKey: "secret",
|
||||
In: strings.NewReader("hello\n/exit\n"),
|
||||
Out: &out,
|
||||
})
|
||||
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(capturedModel).To(Equal("test-model"))
|
||||
Expect(capturedAuth).To(Equal("Bearer secret"))
|
||||
Expect(out.String()).To(ContainSubstring("assistant: hi!"))
|
||||
Expect(out.String()).To(ContainSubstring("bye"))
|
||||
})
|
||||
|
||||
It("auto-selects the only available model", func() {
|
||||
server := chatTestServer([]string{"solo"}, nil)
|
||||
defer server.Close()
|
||||
|
||||
var out bytes.Buffer
|
||||
err := Run(GinkgoT().Context(), Options{
|
||||
BaseURL: server.URL + "/v1",
|
||||
In: strings.NewReader("/exit\n"),
|
||||
Out: &out,
|
||||
})
|
||||
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(out.String()).To(ContainSubstring("LocalAI chat (solo)"))
|
||||
})
|
||||
|
||||
It("returns an actionable error when no models are installed", func() {
|
||||
server := chatTestServer(nil, nil)
|
||||
defer server.Close()
|
||||
|
||||
err := Run(GinkgoT().Context(), Options{
|
||||
BaseURL: server.URL + "/v1",
|
||||
In: strings.NewReader(""),
|
||||
})
|
||||
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("no chat models are installed"))
|
||||
Expect(err.Error()).To(ContainSubstring("local-ai models install <model>"))
|
||||
})
|
||||
|
||||
It("returns an actionable error when multiple models are available without a selection", func() {
|
||||
server := chatTestServer([]string{"alpha", "beta"}, nil)
|
||||
defer server.Close()
|
||||
|
||||
err := Run(GinkgoT().Context(), Options{
|
||||
BaseURL: server.URL + "/v1",
|
||||
In: strings.NewReader(""),
|
||||
})
|
||||
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("multiple models are available"))
|
||||
Expect(err.Error()).To(ContainSubstring("--model"))
|
||||
Expect(err.Error()).To(ContainSubstring("alpha"))
|
||||
Expect(err.Error()).To(ContainSubstring("beta"))
|
||||
})
|
||||
|
||||
It("lists and switches models inside the chat", func() {
|
||||
requestedModels := []string{}
|
||||
server := chatTestServer([]string{"alpha", "beta"}, func(model string) {
|
||||
requestedModels = append(requestedModels, model)
|
||||
})
|
||||
defer server.Close()
|
||||
|
||||
var out bytes.Buffer
|
||||
err := Run(GinkgoT().Context(), Options{
|
||||
Model: "alpha",
|
||||
BaseURL: server.URL + "/v1",
|
||||
In: strings.NewReader("/models\n/model beta\nhello\n/exit\n"),
|
||||
Out: &out,
|
||||
})
|
||||
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(out.String()).To(ContainSubstring("* alpha"))
|
||||
Expect(out.String()).To(ContainSubstring(" beta"))
|
||||
Expect(out.String()).To(ContainSubstring("switched to beta; conversation cleared"))
|
||||
Expect(requestedModels).To(Equal([]string{"beta"}))
|
||||
})
|
||||
})
|
||||
|
||||
func chatTestServer(models []string, onChat func(model string)) *httptest.Server {
|
||||
return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
switch r.URL.Path {
|
||||
case "/v1/models":
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
writeResponse(w, `{"object":"list","data":[`)
|
||||
for i, model := range models {
|
||||
if i > 0 {
|
||||
writeResponse(w, ",")
|
||||
}
|
||||
writeResponsef(w, `{"id":%q,"object":"model"}`, model)
|
||||
}
|
||||
writeResponse(w, `]}`)
|
||||
case "/v1/chat/completions":
|
||||
var body struct {
|
||||
Model string `json:"model"`
|
||||
}
|
||||
Expect(json.NewDecoder(r.Body).Decode(&body)).To(Succeed())
|
||||
if onChat != nil {
|
||||
onChat(body.Model)
|
||||
}
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
writeResponse(w, "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"ok\"}}]}\n\n")
|
||||
writeResponse(w, "data: [DONE]\n\n")
|
||||
default:
|
||||
w.WriteHeader(http.StatusNotFound)
|
||||
}
|
||||
}))
|
||||
}
|
||||
|
||||
func writeResponse(w io.Writer, text string) {
|
||||
_, err := fmt.Fprint(w, text)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
}
|
||||
|
||||
func writeResponsef(w io.Writer, format string, args ...any) {
|
||||
_, err := fmt.Fprintf(w, format, args...)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
}
|
||||
114
core/cli/chat/client.go
Normal file
114
core/cli/chat/client.go
Normal file
@@ -0,0 +1,114 @@
|
||||
package chat
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"sort"
|
||||
"strings"
|
||||
|
||||
openai "github.com/sashabaranov/go-openai"
|
||||
)
|
||||
|
||||
type chatClient interface {
|
||||
ListModels(ctx context.Context) ([]string, error)
|
||||
StreamChat(ctx context.Context, model string, messages []chatMessage, out io.Writer) (string, error)
|
||||
}
|
||||
|
||||
type localAIChatClient struct {
|
||||
client *openai.Client
|
||||
}
|
||||
|
||||
func newLocalAIChatClient(baseURL string, apiKey string) *localAIChatClient {
|
||||
cfg := openai.DefaultConfig(apiKey)
|
||||
cfg.BaseURL = baseURL
|
||||
return &localAIChatClient{client: openai.NewClientWithConfig(cfg)}
|
||||
}
|
||||
|
||||
func (c *localAIChatClient) ListModels(ctx context.Context) ([]string, error) {
|
||||
resp, err := c.client.ListModels(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
models := make([]string, 0, len(resp.Models))
|
||||
for _, model := range resp.Models {
|
||||
if model.ID != "" {
|
||||
models = append(models, model.ID)
|
||||
}
|
||||
}
|
||||
sort.Strings(models)
|
||||
return models, nil
|
||||
}
|
||||
|
||||
func (c *localAIChatClient) StreamChat(ctx context.Context, model string, messages []chatMessage, out io.Writer) (string, error) {
|
||||
stream, err := c.client.CreateChatCompletionStream(ctx, openai.ChatCompletionRequest{
|
||||
Model: model,
|
||||
Messages: openAIChatMessages(messages),
|
||||
})
|
||||
if err != nil {
|
||||
return "", friendlyChatError(err, model)
|
||||
}
|
||||
defer func() {
|
||||
_ = stream.Close()
|
||||
}()
|
||||
|
||||
var answer strings.Builder
|
||||
for {
|
||||
resp, err := stream.Recv()
|
||||
if errors.Is(err, io.EOF) {
|
||||
break
|
||||
}
|
||||
if err != nil {
|
||||
return answer.String(), friendlyChatError(err, model)
|
||||
}
|
||||
if len(resp.Choices) == 0 {
|
||||
continue
|
||||
}
|
||||
|
||||
token := resp.Choices[0].Delta.Content
|
||||
if token == "" {
|
||||
continue
|
||||
}
|
||||
answer.WriteString(token)
|
||||
if _, err := fmt.Fprint(out, token); err != nil {
|
||||
return answer.String(), err
|
||||
}
|
||||
}
|
||||
|
||||
return answer.String(), nil
|
||||
}
|
||||
|
||||
func openAIChatMessages(messages []chatMessage) []openai.ChatCompletionMessage {
|
||||
converted := make([]openai.ChatCompletionMessage, len(messages))
|
||||
for i, message := range messages {
|
||||
converted[i] = openai.ChatCompletionMessage{
|
||||
Role: message.Role,
|
||||
Content: message.Content,
|
||||
}
|
||||
}
|
||||
return converted
|
||||
}
|
||||
|
||||
func friendlyChatError(err error, model string) error {
|
||||
var apiErr *openai.APIError
|
||||
if errors.As(err, &apiErr) {
|
||||
switch apiErr.HTTPStatusCode {
|
||||
case 404:
|
||||
return fmt.Errorf("model %q is not available. Run `local-ai models list`, install a model with `local-ai models install <model>`, or switch with `/model <name>`", model)
|
||||
case 403:
|
||||
return fmt.Errorf("model %q is disabled. Enable it from LocalAI settings or choose another model with `/model <name>`", model)
|
||||
}
|
||||
if apiErr.Message != "" {
|
||||
return errors.New(apiErr.Message)
|
||||
}
|
||||
}
|
||||
|
||||
msg := err.Error()
|
||||
if strings.Contains(msg, "model") && strings.Contains(msg, "not found") {
|
||||
return fmt.Errorf("model %q is not available. Run `local-ai models list`, install a model with `local-ai models install <model>`, or switch with `/model <name>`", model)
|
||||
}
|
||||
|
||||
return err
|
||||
}
|
||||
17
core/cli/chat/models.go
Normal file
17
core/cli/chat/models.go
Normal file
@@ -0,0 +1,17 @@
|
||||
package chat
|
||||
|
||||
import "strings"
|
||||
|
||||
func formatChatModelList(models []string, current string) string {
|
||||
var b strings.Builder
|
||||
for _, model := range models {
|
||||
prefix := " "
|
||||
if model == current {
|
||||
prefix = "* "
|
||||
}
|
||||
b.WriteString(prefix)
|
||||
b.WriteString(model)
|
||||
b.WriteByte('\n')
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
@@ -1,153 +0,0 @@
|
||||
package chat
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
|
||||
"gopkg.in/yaml.v3"
|
||||
)
|
||||
|
||||
// stateDirMode matches the mode nib uses for the same directory. The directory
|
||||
// holds an API key, so it stays owner-only.
|
||||
const stateDirMode = 0o700
|
||||
|
||||
// configFileMode keeps the config owner-only: nib stores the user's API key in
|
||||
// it alongside the keys written here.
|
||||
const configFileMode = 0o600
|
||||
|
||||
// StateDir resolves where the chat agent keeps its config, plugins, and
|
||||
// skills. This is user-scoped rather than server-scoped: chat is a client that
|
||||
// may target a remote LocalAI, so it does not belong under LOCALAI_CONFIG_DIR.
|
||||
func StateDir(override string) (string, error) {
|
||||
if override != "" {
|
||||
return override, nil
|
||||
}
|
||||
if xdg := os.Getenv("XDG_CONFIG_HOME"); xdg != "" {
|
||||
return filepath.Join(xdg, "localai", "chat"), nil
|
||||
}
|
||||
home, err := os.UserHomeDir()
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("resolving home directory for the agent state dir: %w", err)
|
||||
}
|
||||
return filepath.Join(home, ".config", "localai", "chat"), nil
|
||||
}
|
||||
|
||||
// ConfigPath is the agent's config file inside dir.
|
||||
func ConfigPath(dir string) string { return filepath.Join(dir, "config.yaml") }
|
||||
|
||||
// EnsureStateDir creates dir and, on first run only, seeds a config file
|
||||
// pointing at baseURL. It deliberately does not seed a model: a baked-in model
|
||||
// name goes stale as soon as the user installs a different one.
|
||||
//
|
||||
// The config file is machine-managed from here on: nib rewrites it whenever it
|
||||
// self-configures, so hand-written comments in it do not survive.
|
||||
func EnsureStateDir(dir, baseURL string) error {
|
||||
if err := os.MkdirAll(dir, stateDirMode); err != nil {
|
||||
return fmt.Errorf("creating agent state dir %s: %w", dir, err)
|
||||
}
|
||||
path := ConfigPath(dir)
|
||||
if _, err := os.Stat(path); err == nil {
|
||||
return nil // already configured; never overwrite the user's file
|
||||
} else if !os.IsNotExist(err) {
|
||||
return fmt.Errorf("checking agent config %s: %w", path, err)
|
||||
}
|
||||
|
||||
seed := map[string]string{"base_url": baseURL}
|
||||
data, err := yaml.Marshal(seed)
|
||||
if err != nil {
|
||||
return fmt.Errorf("encoding seed agent config: %w", err)
|
||||
}
|
||||
if err := writeConfigFile(path, data); err != nil {
|
||||
return fmt.Errorf("writing seed agent config: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// PersistModel records the chosen model in the agent config, preserving every
|
||||
// other key the user may have set, including the api_key nib writes there.
|
||||
//
|
||||
// The file is machine-managed: this overlays the model onto the parsed keys and
|
||||
// re-marshals, which drops comments. That is deliberate rather than an
|
||||
// oversight, because nib's own save path does the same thing and would erase
|
||||
// them on its next write regardless.
|
||||
func PersistModel(dir, model string) error {
|
||||
// PersistModel is callable before EnsureStateDir, so it cannot assume the
|
||||
// directory exists.
|
||||
if err := os.MkdirAll(dir, stateDirMode); err != nil {
|
||||
return fmt.Errorf("creating agent state dir %s: %w", dir, err)
|
||||
}
|
||||
path := ConfigPath(dir)
|
||||
|
||||
values := map[string]any{}
|
||||
// #nosec G304 -- path is the fixed config.yaml name under the user-selected
|
||||
// chat state directory; selecting that directory is the documented override.
|
||||
data, err := os.ReadFile(path)
|
||||
if err != nil && !os.IsNotExist(err) {
|
||||
return fmt.Errorf("reading agent config %s: %w", path, err)
|
||||
}
|
||||
if err == nil {
|
||||
if err := yaml.Unmarshal(data, &values); err != nil {
|
||||
return fmt.Errorf("parsing agent config %s: %w", path, err)
|
||||
}
|
||||
}
|
||||
values["model"] = model
|
||||
|
||||
out, err := yaml.Marshal(values)
|
||||
if err != nil {
|
||||
return fmt.Errorf("encoding agent config: %w", err)
|
||||
}
|
||||
if err := writeConfigFile(path, out); err != nil {
|
||||
return fmt.Errorf("writing agent config: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// writeConfigFile replaces path with data atomically: it writes a temporary
|
||||
// file next to the target and renames it over the target. Writing the target in
|
||||
// place would truncate it first, so an interrupted or out-of-disk write would
|
||||
// leave a half-written config and destroy the api_key nib keeps in the same
|
||||
// file. The temporary file must share the directory because rename is only
|
||||
// atomic within one filesystem.
|
||||
func writeConfigFile(path string, data []byte) error {
|
||||
dir := filepath.Dir(path)
|
||||
|
||||
// A randomized name rather than a fixed config.yaml.tmp, so two concurrent
|
||||
// writers cannot corrupt each other's temporary file.
|
||||
tmp, err := os.CreateTemp(dir, "config.yaml.*.tmp")
|
||||
if err != nil {
|
||||
return fmt.Errorf("creating temp file in %s: %w", dir, err)
|
||||
}
|
||||
tmpPath := tmp.Name()
|
||||
renamed := false
|
||||
defer func() {
|
||||
if !renamed {
|
||||
// Leave no litter behind on any failure path.
|
||||
_ = os.Remove(tmpPath)
|
||||
}
|
||||
}()
|
||||
|
||||
if _, err := tmp.Write(data); err != nil {
|
||||
_ = tmp.Close()
|
||||
return fmt.Errorf("writing %s: %w", tmpPath, err)
|
||||
}
|
||||
// Flush before the rename: renaming a file whose contents are still only in
|
||||
// the page cache can still lose them across a crash.
|
||||
if err := tmp.Sync(); err != nil {
|
||||
_ = tmp.Close()
|
||||
return fmt.Errorf("syncing %s: %w", tmpPath, err)
|
||||
}
|
||||
if err := tmp.Close(); err != nil {
|
||||
return fmt.Errorf("closing %s: %w", tmpPath, err)
|
||||
}
|
||||
// CreateTemp already asks for 0600, but the umask can only ever clear bits,
|
||||
// so set the mode explicitly rather than inheriting whatever survived.
|
||||
if err := os.Chmod(tmpPath, configFileMode); err != nil {
|
||||
return fmt.Errorf("setting mode on %s: %w", tmpPath, err)
|
||||
}
|
||||
if err := os.Rename(tmpPath, path); err != nil {
|
||||
return fmt.Errorf("replacing %s: %w", path, err)
|
||||
}
|
||||
renamed = true
|
||||
return nil
|
||||
}
|
||||
@@ -1,186 +0,0 @@
|
||||
package chat
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
"gopkg.in/yaml.v3"
|
||||
)
|
||||
|
||||
// richConfig stands in for a config nib has already taken ownership of: a
|
||||
// comment, a secret, and a nested block. A flat scalar alone would not catch a
|
||||
// writer that mangles structure or drops a key it does not know about.
|
||||
const richConfig = `# hand written note
|
||||
base_url: http://x.invalid/v1
|
||||
api_key: secret-token
|
||||
mcp_servers:
|
||||
files:
|
||||
command: mcp-files
|
||||
args:
|
||||
- --root
|
||||
- /tmp
|
||||
`
|
||||
|
||||
var _ = Describe("Agent state directory", func() {
|
||||
Describe("StateDir", func() {
|
||||
It("prefers an explicit override", func() {
|
||||
Expect(StateDir("/custom/dir")).To(Equal("/custom/dir"))
|
||||
})
|
||||
|
||||
It("uses XDG_CONFIG_HOME when set", func() {
|
||||
tmp := GinkgoT().TempDir()
|
||||
GinkgoT().Setenv("XDG_CONFIG_HOME", tmp)
|
||||
Expect(StateDir("")).To(Equal(filepath.Join(tmp, "localai", "chat")))
|
||||
})
|
||||
|
||||
It("falls back to ~/.config/localai/chat", func() {
|
||||
tmp := GinkgoT().TempDir()
|
||||
GinkgoT().Setenv("XDG_CONFIG_HOME", "")
|
||||
GinkgoT().Setenv("HOME", tmp)
|
||||
Expect(StateDir("")).To(Equal(filepath.Join(tmp, ".config", "localai", "chat")))
|
||||
})
|
||||
|
||||
It("fails when neither XDG_CONFIG_HOME nor a home directory is resolvable", func() {
|
||||
GinkgoT().Setenv("XDG_CONFIG_HOME", "")
|
||||
GinkgoT().Setenv("HOME", "")
|
||||
|
||||
dir, err := StateDir("")
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("agent state dir"))
|
||||
// No silent fallback to a relative path: writing an API key into the
|
||||
// working directory would be worse than refusing.
|
||||
Expect(dir).To(BeEmpty())
|
||||
})
|
||||
})
|
||||
|
||||
Describe("EnsureStateDir", func() {
|
||||
It("creates the directory and seeds base_url on first run", func() {
|
||||
dir := filepath.Join(GinkgoT().TempDir(), "chat")
|
||||
Expect(EnsureStateDir(dir, "http://127.0.0.1:8080/v1")).To(Succeed())
|
||||
|
||||
data, err := os.ReadFile(ConfigPath(dir))
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(string(data)).To(ContainSubstring("base_url: http://127.0.0.1:8080/v1"))
|
||||
// A model must NOT be seeded: it goes stale as soon as the user
|
||||
// installs a different one.
|
||||
Expect(string(data)).ToNot(ContainSubstring("model:"))
|
||||
})
|
||||
|
||||
It("keeps the seeded config and its directory owner-only", func() {
|
||||
dir := filepath.Join(GinkgoT().TempDir(), "chat")
|
||||
Expect(EnsureStateDir(dir, "http://127.0.0.1:8080/v1")).To(Succeed())
|
||||
|
||||
// nib writes the user's api_key into this same file, so the modes are
|
||||
// load-bearing, not cosmetic.
|
||||
config, err := os.Stat(ConfigPath(dir))
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(config.Mode().Perm()).To(Equal(os.FileMode(0o600)))
|
||||
|
||||
state, err := os.Stat(dir)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(state.Mode().Perm()).To(Equal(os.FileMode(0o700)))
|
||||
})
|
||||
|
||||
It("leaves an existing config byte-for-byte untouched", func() {
|
||||
dir := GinkgoT().TempDir()
|
||||
Expect(os.WriteFile(ConfigPath(dir), []byte(richConfig), 0o600)).To(Succeed())
|
||||
|
||||
Expect(EnsureStateDir(dir, "http://127.0.0.1:8080/v1")).To(Succeed())
|
||||
|
||||
data, err := os.ReadFile(ConfigPath(dir))
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
// Byte-exact against a fixture carrying a comment and a nested block:
|
||||
// an implementation that "preserves" by re-marshaling through a map
|
||||
// fails here rather than passing on a flat scalar.
|
||||
Expect(string(data)).To(Equal(richConfig))
|
||||
})
|
||||
})
|
||||
|
||||
Describe("PersistModel", func() {
|
||||
It("adds a model to an existing config, preserving other keys", func() {
|
||||
dir := GinkgoT().TempDir()
|
||||
Expect(os.WriteFile(ConfigPath(dir), []byte("base_url: http://x.invalid/v1\n"), 0o600)).To(Succeed())
|
||||
|
||||
Expect(PersistModel(dir, "chosen-model")).To(Succeed())
|
||||
|
||||
data, err := os.ReadFile(ConfigPath(dir))
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(string(data)).To(ContainSubstring("base_url: http://x.invalid/v1"))
|
||||
Expect(string(data)).To(ContainSubstring("model: chosen-model"))
|
||||
})
|
||||
|
||||
It("replaces an existing model rather than duplicating the key", func() {
|
||||
dir := GinkgoT().TempDir()
|
||||
Expect(os.WriteFile(ConfigPath(dir), []byte("model: old\nbase_url: http://x.invalid/v1\n"), 0o600)).To(Succeed())
|
||||
|
||||
Expect(PersistModel(dir, "new")).To(Succeed())
|
||||
|
||||
data, err := os.ReadFile(ConfigPath(dir))
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(string(data)).To(ContainSubstring("model: new"))
|
||||
Expect(string(data)).ToNot(ContainSubstring("model: old"))
|
||||
})
|
||||
|
||||
It("preserves secrets and nested blocks it does not understand", func() {
|
||||
dir := GinkgoT().TempDir()
|
||||
Expect(os.WriteFile(ConfigPath(dir), []byte(richConfig), 0o600)).To(Succeed())
|
||||
|
||||
Expect(PersistModel(dir, "chosen-model")).To(Succeed())
|
||||
|
||||
data, err := os.ReadFile(ConfigPath(dir))
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
var got map[string]any
|
||||
Expect(yaml.Unmarshal(data, &got)).To(Succeed())
|
||||
Expect(got).To(HaveKeyWithValue("model", "chosen-model"))
|
||||
Expect(got).To(HaveKeyWithValue("base_url", "http://x.invalid/v1"))
|
||||
// Losing this key logs the user out of their own server.
|
||||
Expect(got).To(HaveKeyWithValue("api_key", "secret-token"))
|
||||
Expect(got).To(HaveKeyWithValue("mcp_servers",
|
||||
HaveKeyWithValue("files", And(
|
||||
HaveKeyWithValue("command", "mcp-files"),
|
||||
HaveKeyWithValue("args", ConsistOf("--root", "/tmp")),
|
||||
)),
|
||||
))
|
||||
|
||||
// Documented, accepted behavior rather than an aspiration: the overlay
|
||||
// re-marshals, so comments do not survive. nib's own save path erases
|
||||
// them too, so preserving them here would buy nothing.
|
||||
Expect(string(data)).ToNot(ContainSubstring("# hand written note"))
|
||||
})
|
||||
|
||||
It("keeps the rewritten config owner-only and leaves no temp file behind", func() {
|
||||
dir := GinkgoT().TempDir()
|
||||
Expect(os.WriteFile(ConfigPath(dir), []byte(richConfig), 0o600)).To(Succeed())
|
||||
|
||||
Expect(PersistModel(dir, "chosen-model")).To(Succeed())
|
||||
|
||||
info, err := os.Stat(ConfigPath(dir))
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(info.Mode().Perm()).To(Equal(os.FileMode(0o600)))
|
||||
|
||||
// The atomic write stages through a sibling temp file; it must not
|
||||
// survive a successful write.
|
||||
entries, err := os.ReadDir(dir)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
names := []string{}
|
||||
for _, entry := range entries {
|
||||
names = append(names, entry.Name())
|
||||
}
|
||||
Expect(names).To(ConsistOf("config.yaml"))
|
||||
})
|
||||
|
||||
It("creates the state directory when it does not exist yet", func() {
|
||||
// Task 4 may persist a picked model before anything else has run.
|
||||
dir := filepath.Join(GinkgoT().TempDir(), "chat")
|
||||
|
||||
Expect(PersistModel(dir, "chosen-model")).To(Succeed())
|
||||
|
||||
data, err := os.ReadFile(ConfigPath(dir))
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(string(data)).To(ContainSubstring("model: chosen-model"))
|
||||
})
|
||||
})
|
||||
})
|
||||
@@ -1,86 +0,0 @@
|
||||
package chat
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/url"
|
||||
|
||||
openai "github.com/sashabaranov/go-openai"
|
||||
)
|
||||
|
||||
var (
|
||||
// ErrUnreachable means nothing answered at the endpoint. Callers use this
|
||||
// to decide whether offering to start a server makes sense.
|
||||
ErrUnreachable = errors.New("no LocalAI server reachable")
|
||||
// ErrUnauthorized means the server answered but rejected the credentials.
|
||||
ErrUnauthorized = errors.New("LocalAI server rejected the API key")
|
||||
)
|
||||
|
||||
// Probe lists the models the endpoint advertises. It classifies the two
|
||||
// failures that need different advice: nothing listening, and bad credentials.
|
||||
//
|
||||
// The returned list is what the server advertises, verbatim and in server
|
||||
// order. LocalAI happily lists non-model entries it finds in the models
|
||||
// directory (stray archives, dotfiles), and guessing which advertised IDs are
|
||||
// real belongs to whoever presents them, not here.
|
||||
func Probe(ctx context.Context, baseURL, apiKey string) ([]string, error) {
|
||||
cfg := openai.DefaultConfig(apiKey)
|
||||
cfg.BaseURL = baseURL
|
||||
|
||||
resp, err := openai.NewClientWithConfig(cfg).ListModels(ctx)
|
||||
if err != nil {
|
||||
if status, answered := responseStatus(err); answered {
|
||||
if status == http.StatusUnauthorized || status == http.StatusForbidden {
|
||||
return nil, fmt.Errorf("%w: %w", ErrUnauthorized, err)
|
||||
}
|
||||
// The server answered, so it is up; surface its error as-is.
|
||||
return nil, fmt.Errorf("listing models at %s: %w", baseURL, err)
|
||||
}
|
||||
// A caller who cancelled the probe learned nothing about the endpoint,
|
||||
// so claiming it is unreachable would send them to fix a server that
|
||||
// may be fine. A deadline is left alone: an endpoint that cannot answer
|
||||
// within the probe's budget is unreachable for our purposes.
|
||||
var urlErr *url.Error
|
||||
if errors.As(err, &urlErr) && !errors.Is(err, context.Canceled) {
|
||||
// Only a failure to complete the round trip means nothing is
|
||||
// listening. A reply we could not parse is a different problem,
|
||||
// so it falls through to the generic error below.
|
||||
return nil, fmt.Errorf("%w at %s: %w", ErrUnreachable, baseURL, err)
|
||||
}
|
||||
return nil, fmt.Errorf("listing models at %s: %w", baseURL, err)
|
||||
}
|
||||
|
||||
models := make([]string, 0, len(resp.Models))
|
||||
for _, m := range resp.Models {
|
||||
if m.ID != "" {
|
||||
models = append(models, m.ID)
|
||||
}
|
||||
}
|
||||
return models, nil
|
||||
}
|
||||
|
||||
// responseStatus reports the HTTP status a failed call came back with, and
|
||||
// whether there was one at all.
|
||||
//
|
||||
// go-openai splits this across two types depending on the error body, and both
|
||||
// occur against a real LocalAI: it returns *openai.APIError when the body
|
||||
// parses as an OpenAI error envelope, which is what LocalAI's normal error
|
||||
// handler sends, and *openai.RequestError when it does not, which is what
|
||||
// LocalAI sends when started with opaque errors, since that handler replies
|
||||
// with a bare status and no body.
|
||||
func responseStatus(err error) (int, bool) {
|
||||
// *RequestError is checked first because it is the outer type when
|
||||
// go-openai nests one error inside the other; the inner value in that case
|
||||
// carries no status.
|
||||
var reqErr *openai.RequestError
|
||||
if errors.As(err, &reqErr) {
|
||||
return reqErr.HTTPStatusCode, true
|
||||
}
|
||||
var apiErr *openai.APIError
|
||||
if errors.As(err, &apiErr) {
|
||||
return apiErr.HTTPStatusCode, true
|
||||
}
|
||||
return 0, false
|
||||
}
|
||||
@@ -1,169 +0,0 @@
|
||||
package chat
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"time"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
)
|
||||
|
||||
var _ = Describe("Probe", func() {
|
||||
It("returns the advertised models", func() {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
Expect(json.NewEncoder(w).Encode(map[string]any{
|
||||
"object": "list",
|
||||
"data": []map[string]string{
|
||||
{"id": "model-a", "object": "model"},
|
||||
{"id": "model-b", "object": "model"},
|
||||
},
|
||||
})).To(Succeed())
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
models, err := Probe(context.Background(), srv.URL+"/v1", "")
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(models).To(Equal([]string{"model-a", "model-b"}))
|
||||
})
|
||||
|
||||
It("reports an unreachable server distinguishably", func() {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))
|
||||
url := srv.URL
|
||||
srv.Close() // nothing is listening now
|
||||
|
||||
_, err := Probe(context.Background(), url+"/v1", "")
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(errors.Is(err, ErrUnreachable)).To(BeTrue(), "want ErrUnreachable, got %v", err)
|
||||
})
|
||||
|
||||
It("reports an auth failure distinguishably", func() {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(http.StatusUnauthorized)
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
_, err := Probe(context.Background(), srv.URL+"/v1", "bad-key")
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(errors.Is(err, ErrUnauthorized)).To(BeTrue(), "want ErrUnauthorized, got %v", err)
|
||||
})
|
||||
|
||||
// LocalAI's normal error handler replies with an OpenAI error envelope, and
|
||||
// its opaque-errors handler replies with a bare status and no body. Those
|
||||
// reach the client as two different go-openai types, so both have to be
|
||||
// classified the same way.
|
||||
It("reports an auth failure carrying an error envelope distinguishably", func() {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusUnauthorized)
|
||||
Expect(json.NewEncoder(w).Encode(map[string]any{
|
||||
"error": map[string]any{"message": "invalid api key", "code": http.StatusUnauthorized},
|
||||
})).To(Succeed())
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
_, err := Probe(context.Background(), srv.URL+"/v1", "bad-key")
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(errors.Is(err, ErrUnauthorized)).To(BeTrue(), "want ErrUnauthorized, got %v", err)
|
||||
})
|
||||
|
||||
It("does not call a server that answered with an error unreachable", func() {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(http.StatusInternalServerError)
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
_, err := Probe(context.Background(), srv.URL+"/v1", "")
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(errors.Is(err, ErrUnreachable)).To(BeFalse(), "a server that replied is not unreachable, got %v", err)
|
||||
Expect(errors.Is(err, ErrUnauthorized)).To(BeFalse(), "500 is not an auth failure, got %v", err)
|
||||
})
|
||||
|
||||
// Pointing chat at some other service that happens to be listening is a
|
||||
// different problem from nothing listening, and needs different advice.
|
||||
It("does not call a reply it could not parse unreachable", func() {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "text/html")
|
||||
_, err := w.Write([]byte("<html><body>not LocalAI</body></html>"))
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
_, err := Probe(context.Background(), srv.URL+"/v1", "")
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(errors.Is(err, ErrUnreachable)).To(BeFalse(), "something answered, got %v", err)
|
||||
})
|
||||
|
||||
It("returns every advertised id, including ones that are not models", func() {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
Expect(json.NewEncoder(w).Encode(map[string]any{
|
||||
"object": "list",
|
||||
"data": []map[string]string{
|
||||
{"id": "zeta", "object": "model"},
|
||||
{"id": ".gitignore", "object": "model"},
|
||||
{"id": "alpha", "object": "model"},
|
||||
{"id": "voice.tar.bz2", "object": "model"},
|
||||
},
|
||||
})).To(Succeed())
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
// Verbatim and in server order: deciding which of these are real, and
|
||||
// what order to show them in, belongs to the caller.
|
||||
models, err := Probe(context.Background(), srv.URL+"/v1", "")
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(models).To(Equal([]string{"zeta", ".gitignore", "alpha", "voice.tar.bz2"}))
|
||||
})
|
||||
|
||||
It("stops early when the context is already cancelled", func() {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
Expect(json.NewEncoder(w).Encode(map[string]any{"object": "list", "data": []any{}})).To(Succeed())
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
|
||||
_, err := Probe(ctx, srv.URL+"/v1", "")
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(errors.Is(err, context.Canceled)).To(BeTrue(), "want the cancellation preserved, got %v", err)
|
||||
// A cancelled probe learned nothing about the endpoint, so it must not
|
||||
// send the caller off to start a server that may already be running.
|
||||
Expect(errors.Is(err, ErrUnreachable)).To(BeFalse(), "cancelling is not a verdict on the server, got %v", err)
|
||||
})
|
||||
|
||||
It("reports a server that never answers as unreachable", func() {
|
||||
release := make(chan struct{})
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
<-release
|
||||
}))
|
||||
defer srv.Close()
|
||||
defer close(release)
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
|
||||
defer cancel()
|
||||
|
||||
_, err := Probe(ctx, srv.URL+"/v1", "")
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(errors.Is(err, ErrUnreachable)).To(BeTrue(), "want ErrUnreachable, got %v", err)
|
||||
Expect(errors.Is(err, context.DeadlineExceeded)).To(BeTrue(), "want the deadline preserved, got %v", err)
|
||||
})
|
||||
|
||||
It("returns an empty list when the server has no models", func() {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
Expect(json.NewEncoder(w).Encode(map[string]any{"object": "list", "data": []any{}})).To(Succeed())
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
models, err := Probe(context.Background(), srv.URL+"/v1", "")
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(models).To(BeEmpty())
|
||||
})
|
||||
})
|
||||
@@ -1,95 +0,0 @@
|
||||
package chat
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"slices"
|
||||
"sort"
|
||||
"strings"
|
||||
|
||||
"github.com/mudler/xlog"
|
||||
)
|
||||
|
||||
// ModelChooser asks the user to pick one of models. It is nil when the session
|
||||
// is not interactive.
|
||||
type ModelChooser func(models []string) (string, error)
|
||||
|
||||
// ModelRequest is everything model resolution needs.
|
||||
type ModelRequest struct {
|
||||
Flag string // --model
|
||||
Configured string // model recorded in the agent config
|
||||
Available []string // models the server advertises
|
||||
StateDir string // where an interactive choice is persisted
|
||||
Choose ModelChooser // nil means non-interactive
|
||||
// Notify reports a problem that is worth telling the user about but not
|
||||
// worth failing over. Nil discards it. It exists because the one such
|
||||
// problem here, a choice that could not be saved, changes what the user
|
||||
// should expect next: they will be asked again. A log line does not reach
|
||||
// them, since the agent runs at log level error by default.
|
||||
Notify func(message string)
|
||||
}
|
||||
|
||||
// ResolveModel picks the model for this invocation. A flag or a configured
|
||||
// value wins outright and is not persisted; only an interactive choice is
|
||||
// written back, so the prompt appears at most once.
|
||||
//
|
||||
// Available is used exactly as the server gave it. LocalAI advertises stray
|
||||
// files it finds in the models directory alongside real models, but real model
|
||||
// IDs contain dots too (lfm2.5-8b-a1b), so any client-side "looks like a
|
||||
// filename" heuristic would eventually hide a model the user has. Deciding
|
||||
// which advertised IDs are real belongs to the endpoint, not to a guess here.
|
||||
func ResolveModel(req ModelRequest) (string, error) {
|
||||
if req.Flag != "" {
|
||||
return req.Flag, nil
|
||||
}
|
||||
if req.Configured != "" {
|
||||
return req.Configured, nil
|
||||
}
|
||||
|
||||
// The server's /v1/models ordering is not stable between calls, so sort
|
||||
// before showing or listing: the same number must mean the same model on
|
||||
// the next run. Sort a copy; the caller's slice is not ours to reorder.
|
||||
available := append([]string(nil), req.Available...)
|
||||
sort.Strings(available)
|
||||
|
||||
switch len(available) {
|
||||
case 0:
|
||||
return "", errors.New("the LocalAI server has no models installed. Install one with 'local-ai models install <name>', then run 'local-ai chat' again")
|
||||
case 1:
|
||||
return available[0], nil
|
||||
}
|
||||
|
||||
if req.Choose == nil {
|
||||
return "", fmt.Errorf(
|
||||
"several models are available; pick one with --model. Available: %s",
|
||||
strings.Join(available, ", "),
|
||||
)
|
||||
}
|
||||
|
||||
chosen, err := req.Choose(available)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
// Choose is an interface, so its answer is checked rather than trusted.
|
||||
// What comes back is persisted and every later run starts against it, so a
|
||||
// chooser that returns an empty string or a name of its own would record a
|
||||
// model the server never offered and there would be nothing left to catch
|
||||
// it.
|
||||
if !slices.Contains(available, chosen) {
|
||||
return "", fmt.Errorf(
|
||||
"the model chooser answered %q, which is not one of the available models: %s",
|
||||
chosen, strings.Join(available, ", "),
|
||||
)
|
||||
}
|
||||
if req.StateDir != "" {
|
||||
if err := PersistModel(req.StateDir, chosen); err != nil {
|
||||
// A failure to remember the choice must not block the session: the
|
||||
// user picked a model, so honour it and say what will happen.
|
||||
xlog.Warn("could not save the model choice", "error", err, "model", chosen)
|
||||
if req.Notify != nil {
|
||||
req.Notify(fmt.Sprintf("Your choice of %s could not be saved, so this question comes back next time: %v", chosen, err))
|
||||
}
|
||||
}
|
||||
}
|
||||
return chosen, nil
|
||||
}
|
||||
@@ -1,156 +0,0 @@
|
||||
package chat
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"os"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
)
|
||||
|
||||
var _ = Describe("ResolveModel", func() {
|
||||
It("prefers the flag over everything", func() {
|
||||
got, err := ResolveModel(ModelRequest{
|
||||
Flag: "from-flag",
|
||||
Configured: "from-config",
|
||||
Available: []string{"a", "b"},
|
||||
})
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(got).To(Equal("from-flag"))
|
||||
})
|
||||
|
||||
It("uses the configured model when no flag is given", func() {
|
||||
got, err := ResolveModel(ModelRequest{
|
||||
Configured: "from-config",
|
||||
Available: []string{"a", "b"},
|
||||
})
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(got).To(Equal("from-config"))
|
||||
})
|
||||
|
||||
It("auto-selects when the server offers exactly one model", func() {
|
||||
got, err := ResolveModel(ModelRequest{Available: []string{"only-one"}})
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(got).To(Equal("only-one"))
|
||||
})
|
||||
|
||||
It("errors and lists the options when several models exist and there is no chooser", func() {
|
||||
_, err := ResolveModel(ModelRequest{Available: []string{"a", "b"}})
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("a"))
|
||||
Expect(err.Error()).To(ContainSubstring("b"))
|
||||
Expect(err.Error()).To(ContainSubstring("--model"))
|
||||
})
|
||||
|
||||
It("sorts before offering, so the same number means the same model next run", func() {
|
||||
var offered []string
|
||||
available := []string{"zeta", "alpha", "mid"}
|
||||
_, err := ResolveModel(ModelRequest{
|
||||
Available: available,
|
||||
StateDir: GinkgoT().TempDir(),
|
||||
Choose: func(models []string) (string, error) {
|
||||
offered = models
|
||||
return models[0], nil
|
||||
},
|
||||
})
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
// The server's /v1/models ordering is unstable between calls.
|
||||
Expect(offered).To(Equal([]string{"alpha", "mid", "zeta"}))
|
||||
// Sorting must happen on a copy: the caller still owns this slice, and
|
||||
// reordering it under them would move whatever they index into it.
|
||||
Expect(available).To(Equal([]string{"zeta", "alpha", "mid"}))
|
||||
})
|
||||
|
||||
It("lists models in sorted order in the several-models error", func() {
|
||||
_, err := ResolveModel(ModelRequest{Available: []string{"zeta", "alpha"}})
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("alpha, zeta"))
|
||||
})
|
||||
|
||||
It("asks the chooser when several models exist, and persists the answer", func() {
|
||||
dir := GinkgoT().TempDir()
|
||||
got, err := ResolveModel(ModelRequest{
|
||||
Available: []string{"a", "b"},
|
||||
StateDir: dir,
|
||||
Choose: func(models []string) (string, error) { return models[1], nil },
|
||||
})
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(got).To(Equal("b"))
|
||||
|
||||
data, err := os.ReadFile(ConfigPath(dir))
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(string(data)).To(ContainSubstring("model: b"))
|
||||
})
|
||||
|
||||
// The answer is persisted and every later run starts against it, and
|
||||
// ModelChooser is exported, so the invariant has to hold for choosers this
|
||||
// package did not write.
|
||||
DescribeTable("refuses an answer the chooser was not offered",
|
||||
func(answer string) {
|
||||
dir := GinkgoT().TempDir()
|
||||
got, err := ResolveModel(ModelRequest{
|
||||
Available: []string{"alpha", "zeta"},
|
||||
StateDir: dir,
|
||||
Choose: func([]string) (string, error) { return answer, nil },
|
||||
})
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(got).To(BeEmpty())
|
||||
Expect(err.Error()).To(ContainSubstring("alpha, zeta"))
|
||||
|
||||
_, statErr := os.Stat(ConfigPath(dir))
|
||||
Expect(os.IsNotExist(statErr)).To(BeTrue(), "nothing may be recorded for an answer that was refused")
|
||||
},
|
||||
Entry("nothing at all", ""),
|
||||
Entry("a model the server never offered", "gamma"),
|
||||
Entry("an offered model with stray whitespace", " alpha"),
|
||||
Entry("an offered model in the wrong case", "Alpha"),
|
||||
)
|
||||
|
||||
It("notifies, and still honours the choice, when it cannot be persisted", func() {
|
||||
dir := GinkgoT().TempDir()
|
||||
// A directory where the config file belongs: the write fails for any
|
||||
// user, including root.
|
||||
Expect(os.MkdirAll(ConfigPath(dir), 0o700)).To(Succeed())
|
||||
|
||||
var notices []string
|
||||
got, err := ResolveModel(ModelRequest{
|
||||
Available: []string{"a", "b"},
|
||||
StateDir: dir,
|
||||
Choose: func(models []string) (string, error) { return models[0], nil },
|
||||
Notify: func(message string) { notices = append(notices, message) },
|
||||
})
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(got).To(Equal("a"))
|
||||
Expect(notices).To(HaveLen(1))
|
||||
Expect(notices[0]).To(ContainSubstring("a"))
|
||||
Expect(notices[0]).To(ContainSubstring("could not be saved"))
|
||||
})
|
||||
|
||||
It("says nothing when the choice was saved", func() {
|
||||
var notices []string
|
||||
_, err := ResolveModel(ModelRequest{
|
||||
Available: []string{"a", "b"},
|
||||
StateDir: GinkgoT().TempDir(),
|
||||
Choose: func(models []string) (string, error) { return models[0], nil },
|
||||
Notify: func(message string) { notices = append(notices, message) },
|
||||
})
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(notices).To(BeEmpty())
|
||||
})
|
||||
|
||||
It("propagates a chooser cancellation", func() {
|
||||
cancelled := errors.New("cancelled")
|
||||
_, err := ResolveModel(ModelRequest{
|
||||
Available: []string{"a", "b"},
|
||||
StateDir: GinkgoT().TempDir(),
|
||||
Choose: func([]string) (string, error) { return "", cancelled },
|
||||
})
|
||||
Expect(errors.Is(err, cancelled)).To(BeTrue())
|
||||
})
|
||||
|
||||
It("errors with an install hint when the server has no models", func() {
|
||||
_, err := ResolveModel(ModelRequest{Available: nil})
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("local-ai models install"))
|
||||
})
|
||||
})
|
||||
@@ -1,475 +0,0 @@
|
||||
package chat
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"os/signal"
|
||||
"strconv"
|
||||
"strings"
|
||||
"syscall"
|
||||
"time"
|
||||
|
||||
"github.com/mudler/nib/app"
|
||||
nibcmd "github.com/mudler/nib/cmd"
|
||||
nibconfig "github.com/mudler/nib/config"
|
||||
nibtypes "github.com/mudler/nib/types"
|
||||
"golang.org/x/term"
|
||||
)
|
||||
|
||||
// Options is everything the chat command passes down from its flags.
|
||||
type Options struct {
|
||||
Args []string // forwarded to the agent verbatim
|
||||
Endpoint string // the server root, e.g. http://127.0.0.1:8080
|
||||
BaseURL string // the API base, e.g. http://127.0.0.1:8080/v1
|
||||
APIKey string
|
||||
Model string
|
||||
StateDir string
|
||||
TraceDir string
|
||||
Yolo bool
|
||||
// ProbeTimeout bounds each check of the server. Zero means
|
||||
// defaultProbeTimeout.
|
||||
ProbeTimeout time.Duration
|
||||
|
||||
In io.Reader
|
||||
Out io.Writer
|
||||
ErrOut io.Writer
|
||||
}
|
||||
|
||||
// ExitStatus reports the status the process should exit with for an agent run
|
||||
// that failed, and whether err is such a failure.
|
||||
//
|
||||
// nib writes what went wrong to the error stream itself and hands back nothing
|
||||
// but a code, so an error that satisfies this has already been explained to the
|
||||
// user and must not be reported a second time. The refusal to open a
|
||||
// full-screen session on a stdin that cannot be read arrives this way, and it
|
||||
// is the one a user is most likely to meet: 'echo q | local-ai chat' names
|
||||
// --cli, and burying that under a second message would hide the fix.
|
||||
func ExitStatus(err error) (int, bool) {
|
||||
var exit app.ExitError
|
||||
if errors.As(err, &exit) {
|
||||
return exit.Code, true
|
||||
}
|
||||
return 0, false
|
||||
}
|
||||
|
||||
// shutdownSignals end the session. SIGHUP is one of them because this is a
|
||||
// terminal program: once the terminal is gone there is nobody left to talk to,
|
||||
// and a server started for the session has to go with it.
|
||||
var shutdownSignals = []os.Signal{os.Interrupt, syscall.SIGTERM, syscall.SIGHUP}
|
||||
|
||||
// shutdownContext derives a context that is cancelled when the process is
|
||||
// asked to stop.
|
||||
//
|
||||
// Without it a signal kills this process where it stands, skipping every
|
||||
// deferred call, and a 'local-ai run' started for the session is reparented to
|
||||
// init with nothing left that knows to shut it down. An interactive Ctrl+C is
|
||||
// safe on its own, because the child shares this process' foreground process
|
||||
// group and the terminal signals all of it, but a SIGTERM from a supervisor or
|
||||
// a script reaches only this process.
|
||||
//
|
||||
// Since nib v0.5.1 cancelling this context does end the session: RunTUI passes
|
||||
// it to bubbletea, which unwinds the program and reports the context's own
|
||||
// error. The server is still stopped on cancellation rather than on the way
|
||||
// out (see runSession), because registering here removes SIGHUP's default
|
||||
// terminate disposition, and a guarantee about a server this process owns is
|
||||
// not worth resting on how promptly a third party unwinds its interface.
|
||||
//
|
||||
// A handler rather than SysProcAttr.Pdeathsig on the child: Pdeathsig is
|
||||
// Linux-only, and in Go it is delivered when the OS thread that forked exits
|
||||
// rather than when the process does, so it can fire on a perfectly healthy
|
||||
// parent. Setpgid is not an alternative either, since taking the child out of
|
||||
// the foreground process group is what would break the Ctrl+C that works
|
||||
// today. SIGKILL stays uncovered, as it must: nothing in the process can
|
||||
// observe it.
|
||||
func shutdownContext(parent context.Context) (context.Context, context.CancelFunc) {
|
||||
return signal.NotifyContext(parent, shutdownSignals...)
|
||||
}
|
||||
|
||||
// Run starts the agent: resolve where state lives, make sure a server is
|
||||
// reachable, pick a model, then hand off to nib.
|
||||
func Run(ctx context.Context, opts Options) error {
|
||||
ctx, stop := shutdownContext(ctx)
|
||||
defer stop()
|
||||
|
||||
p, err := prepare(ctx, opts, isTerminal(opts.In))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
// A server this process started belongs to this session, and Stop is
|
||||
// nil-safe and idempotent, so one defer covers both cases and costs nothing
|
||||
// when runSession has already stopped it.
|
||||
defer p.server.Stop()
|
||||
|
||||
return runSession(ctx, p.server, func(ctx context.Context) error {
|
||||
return runAgent(ctx, p.dir, p.model, opts)
|
||||
})
|
||||
}
|
||||
|
||||
// runSession hands the terminal to agent, and stops a server started for this
|
||||
// session as soon as the context is cancelled rather than when agent returns.
|
||||
//
|
||||
// The difference matters because the deferred Stop in Run is only reached once
|
||||
// agent returns, and how long that takes is nib's business rather than ours.
|
||||
// nib v0.5.1 does unwind the TUI on a cancelled context, so it does return; a
|
||||
// SIGHUP no longer leaves the interface on screen with the server behind it,
|
||||
// which it did before, when bubbletea's own SIGINT and SIGTERM handler was the
|
||||
// only thing that ever quit the program and registering for SIGHUP had removed
|
||||
// the default disposition that used to end the process. Watching the context
|
||||
// keeps the guarantee independent of what the agent does with it.
|
||||
func runSession(ctx context.Context, server *StartedServer, agent func(context.Context) error) error {
|
||||
returned := make(chan struct{})
|
||||
defer close(returned)
|
||||
|
||||
go func() {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
server.Stop()
|
||||
case <-returned:
|
||||
}
|
||||
}()
|
||||
|
||||
return agent(ctx)
|
||||
}
|
||||
|
||||
// preparation is what the agent needs once the environment is ready: where its
|
||||
// state lives, which model to talk to, and the server this process started on
|
||||
// the user's behalf, if any.
|
||||
type preparation struct {
|
||||
dir string
|
||||
model string
|
||||
server *StartedServer
|
||||
}
|
||||
|
||||
// prepare does everything that has to happen before the agent takes over the
|
||||
// terminal. It is split out of Run because all of it is testable and none of
|
||||
// what follows is: once app.Run has the terminal there is no seam left.
|
||||
//
|
||||
// interactive says whether there is a user to prompt. It is a parameter rather
|
||||
// than a second read of opts.In so the prompts can be driven over a pipe.
|
||||
func prepare(ctx context.Context, opts Options, interactive bool) (_ *preparation, err error) {
|
||||
dir, dirErr := StateDir(opts.StateDir)
|
||||
if dirErr != nil {
|
||||
return nil, dirErr
|
||||
}
|
||||
if err := EnsureStateDir(dir, opts.BaseURL); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if isLocalOnlyArgs(opts.Args) {
|
||||
return &preparation{dir: dir}, nil
|
||||
}
|
||||
|
||||
// One prompter for every question this run asks; see its doc comment for
|
||||
// why the reader cannot be rebuilt per question.
|
||||
var prompts *prompter
|
||||
if interactive {
|
||||
prompts = newPrompter(opts.In, opts.ErrOut)
|
||||
}
|
||||
|
||||
var started *StartedServer
|
||||
defer func() {
|
||||
// Nothing after the spawn may leave a server behind: the caller only
|
||||
// learns about it through a successful return.
|
||||
if err != nil {
|
||||
started.Stop()
|
||||
}
|
||||
}()
|
||||
|
||||
models, err := probeModels(ctx, opts)
|
||||
if err != nil {
|
||||
if errors.Is(err, ErrUnauthorized) {
|
||||
return nil, fmt.Errorf("the LocalAI server at %s rejected the API key. Pass --api-key or set LOCALAI_API_KEY", opts.Endpoint)
|
||||
}
|
||||
if !errors.Is(err, ErrUnreachable) {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var confirm Confirmer
|
||||
if interactive {
|
||||
confirm = prompts.yesNo
|
||||
}
|
||||
var startErr error
|
||||
started, startErr = OfferToStart(ctx, StartOptions{
|
||||
Endpoint: opts.Endpoint,
|
||||
Confirm: confirm,
|
||||
Stderr: opts.ErrOut,
|
||||
})
|
||||
if startErr != nil {
|
||||
err = startErr
|
||||
if errors.Is(startErr, ErrDeclined) {
|
||||
err = fmt.Errorf("no LocalAI server at %s. Start one with 'local-ai run', or point elsewhere with --endpoint", opts.Endpoint)
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
say(opts.ErrOut, "Started a temporary LocalAI server; it stops when you exit. Use 'local-ai run' for a persistent one.\n")
|
||||
|
||||
if models, err = probeModels(ctx, opts); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
|
||||
var chooser ModelChooser
|
||||
if interactive {
|
||||
chooser = prompts.choose
|
||||
}
|
||||
model, err := ResolveModel(ModelRequest{
|
||||
Flag: opts.Model,
|
||||
Configured: configuredModel(dir),
|
||||
Available: models,
|
||||
StateDir: dir,
|
||||
Choose: chooser,
|
||||
Notify: func(message string) { say(opts.ErrOut, "%s\n", message) },
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return &preparation{dir: dir, model: model, server: started}, nil
|
||||
}
|
||||
|
||||
func runAgent(ctx context.Context, dir, model string, opts Options) error {
|
||||
return app.Run(ctx, agentOptions(dir, model, opts))
|
||||
}
|
||||
|
||||
// agentOptions builds the request handed to nib. It is split out of runAgent
|
||||
// because app.Run takes the terminal and cannot be called from a test, while
|
||||
// what is asked of it is exactly the part worth pinning.
|
||||
//
|
||||
// The stream fields are the interesting ones, and they are not symmetric.
|
||||
//
|
||||
// nib reads a non-nil stream as "the embedder wants this used", and refuses
|
||||
// every mode but --cli when such a stream is not a terminal, because the
|
||||
// full-screen interface renders on /dev/tty and would otherwise ignore it in
|
||||
// silence. Nil means "not injected": nib falls back to the process stream and
|
||||
// behaves as standalone nib does.
|
||||
//
|
||||
// Stdin is passed through as it comes. A piped or redirected stdin really is
|
||||
// ignored by the interface, so the refusal is the honest answer there, and it
|
||||
// is the one users meet: 'echo q | local-ai chat' says to re-run with --cli
|
||||
// rather than opening a full-screen session that will never read the question.
|
||||
//
|
||||
// Stdout is different, and the process stream is deliberately sent as nil. The
|
||||
// interface does write to stdout even when it is a pipe: that is the whole of
|
||||
// nib's shell-capture idiom, out=$(local-ai chat --height 50%), which is what
|
||||
// the Ctrl+Space widget emitted by --init is built on. Injecting os.Stdout
|
||||
// there would refuse the widget for a stream nib was going to use anyway.
|
||||
//
|
||||
// The test is identity with os.Stdout rather than whether it happens to be a
|
||||
// terminal, which means a shell redirect goes the same way as the widget:
|
||||
// 'local-ai chat > out.txt' no longer refuses either, and renders on /dev/tty
|
||||
// with the capture line landing in the file. That is not a second decision, it
|
||||
// is the same one. Both are the process stdout as the shell handed it over,
|
||||
// differing only in being a pipe rather than a regular file, which nib's gate
|
||||
// does not look at and should not. Refusing one would refuse the other.
|
||||
//
|
||||
// What stays injected, and so stays subject to the refusal, is a writer some
|
||||
// in-process caller chose for itself rather than inherited: a bytes.Buffer, or
|
||||
// an *os.File it opened. The specs rely on that.
|
||||
//
|
||||
// Stderr is never gated by nib, so it is passed through unchanged.
|
||||
//
|
||||
// The config values go through Overrides rather than Defaults, and that is not
|
||||
// a detail. Defaults are seeds: they sit BENEATH the config file, so the file
|
||||
// silently undoes them. Everything here is a decision this invocation already
|
||||
// made on the user's behalf, and a flag that the file can undo is not a flag.
|
||||
// It was not a rare case either, since EnsureStateDir writes base_url on the
|
||||
// first run and an interactive choice writes model, so from the second run on
|
||||
// the file carried a value for both and --endpoint and --model did nothing.
|
||||
//
|
||||
// The one asymmetry to plan around is that nib cannot tell "set to the zero
|
||||
// value" from "not set", so an override only ever raises a field. --yolo can
|
||||
// turn approval off, but nothing on the command line can turn it back on over
|
||||
// an approval_mode: auto in the file; that needs a config edit. Same shape for
|
||||
// the strings, which is what makes an unset --api-key or --trace-dir leave the
|
||||
// file's value standing, as it should.
|
||||
//
|
||||
// nib's own --trace-dir and --yolo, and their NIB_TRACE_DIR and NIB_YOLO twins,
|
||||
// are resolved after the config load and so still outrank these. That is
|
||||
// deliberate upstream: they are instructions to nib rather than ambient
|
||||
// environment.
|
||||
func agentOptions(dir, model string, opts Options) app.Options {
|
||||
// Model is the model this run resolved, which already prefers --model and
|
||||
// falls back to the file's own model, so the override restates the file's
|
||||
// value rather than fighting it whenever no flag was given.
|
||||
//
|
||||
// BaseURL is the endpoint this run probed, offered to start a server for,
|
||||
// and seeded the config with. Handing nib a different one is precisely the
|
||||
// split that made --endpoint a no-op, so the agent talks to the server
|
||||
// LocalAI checked. Pointing somewhere else for good is LOCALAI_CHAT_ENDPOINT
|
||||
// or --endpoint, not a hand-edited base_url the probe never reads.
|
||||
//
|
||||
// APIKey and TraceDir are the flags as given, empty when they were not, and
|
||||
// an empty override leaves the file alone. TraceDir is runtime-only in nib
|
||||
// (yaml:"-"), so no file value exists for it to beat today; it belongs here
|
||||
// with the other flags rather than one rung down for a reason that could
|
||||
// quietly stop being true.
|
||||
overrides := nibtypes.Config{
|
||||
Model: model,
|
||||
APIKey: opts.APIKey,
|
||||
BaseURL: opts.BaseURL,
|
||||
TraceDir: opts.TraceDir,
|
||||
}
|
||||
if opts.Yolo {
|
||||
overrides.ApprovalMode = "auto"
|
||||
}
|
||||
|
||||
return app.Options{
|
||||
Args: opts.Args,
|
||||
ProgramName: "local-ai chat",
|
||||
BaseDir: dir,
|
||||
Overrides: overrides,
|
||||
SkipSetup: true,
|
||||
SkipBareEnv: true,
|
||||
Stdin: opts.In,
|
||||
Stdout: ownStdout(opts.Out),
|
||||
Stderr: opts.ErrOut,
|
||||
}
|
||||
}
|
||||
|
||||
// ownStdout reports the writer as nib's own rather than as an injected one when
|
||||
// it is the process stdout, by answering nil for it. See agentOptions for why
|
||||
// that distinction is the difference between a working Ctrl+Space widget and a
|
||||
// refused one.
|
||||
func ownStdout(w io.Writer) io.Writer {
|
||||
if f, ok := w.(*os.File); ok && f == os.Stdout {
|
||||
return nil
|
||||
}
|
||||
return w
|
||||
}
|
||||
|
||||
// defaultProbeTimeout bounds a check of the server. Listing models is cheap,
|
||||
// so this is long enough that a loaded server is never given up on and short
|
||||
// enough that a hung one does not leave the user staring at nothing.
|
||||
const defaultProbeTimeout = 30 * time.Second
|
||||
|
||||
// probeModels lists what the endpoint offers, under a budget.
|
||||
func probeModels(ctx context.Context, opts Options) ([]string, error) {
|
||||
timeout := opts.ProbeTimeout
|
||||
if timeout <= 0 {
|
||||
timeout = defaultProbeTimeout
|
||||
}
|
||||
// A real deadline rather than a cancel plus a timer. Probe reads
|
||||
// context.Canceled as "the caller gave up", which is a statement about the
|
||||
// caller and not about the endpoint, and only a deadline as "nothing
|
||||
// answered in time". Expiring the budget as a cancellation would stop
|
||||
// ErrUnreachable firing for precisely the hung servers that the offer to
|
||||
// start one exists for.
|
||||
probeCtx, cancel := context.WithTimeout(ctx, timeout)
|
||||
defer cancel()
|
||||
return Probe(probeCtx, opts.BaseURL, opts.APIKey)
|
||||
}
|
||||
|
||||
// isLocalOnlyArgs reports whether the forwarded arguments do their work
|
||||
// without ever reaching a model, in which case demanding a running server (and
|
||||
// offering to start one) would be an obstacle rather than a service.
|
||||
//
|
||||
// Two groups qualify. The management subcommands edit nib's own state: plugin,
|
||||
// skill, and the mcp verbs that add or remove configured servers, which is
|
||||
// asked of nib rather than restated, because bare 'mcp' and its transport
|
||||
// flags do serve the agent and do need a model. The other group is the flags
|
||||
// that only print something, above all --init: its shell snippet goes into an
|
||||
// rc file, typically long before any server exists.
|
||||
func isLocalOnlyArgs(args []string) bool {
|
||||
if len(args) == 0 {
|
||||
return false
|
||||
}
|
||||
// A scan rather than a look at args[0]: the mode flags this command
|
||||
// translates are prepended, so --init is not necessarily first. Positional
|
||||
// text cannot be mistaken for a flag here, since nib ignores what is left
|
||||
// after flag parsing.
|
||||
for _, a := range args {
|
||||
switch {
|
||||
case a == "--init", a == "-init", strings.HasPrefix(a, "--init="), strings.HasPrefix(a, "-init="):
|
||||
return true
|
||||
case a == "--version", a == "-version":
|
||||
return true
|
||||
}
|
||||
}
|
||||
switch args[0] {
|
||||
case "plugin", "skill":
|
||||
return true
|
||||
case "mcp":
|
||||
return len(args) >= 2 && nibcmd.IsMCPManageSubcommand(args[1])
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// configuredModel reads the model already recorded in the agent config, if any.
|
||||
func configuredModel(dir string) string {
|
||||
cfg := nibconfig.LoadWith(nibconfig.LoadOptions{BaseDir: dir, SkipBareEnv: true})
|
||||
return cfg.Model
|
||||
}
|
||||
|
||||
func isTerminal(in io.Reader) bool {
|
||||
f, ok := in.(*os.File)
|
||||
return ok && term.IsTerminal(int(f.Fd()))
|
||||
}
|
||||
|
||||
// say writes a line of interactive chatter: a question, or a notice about
|
||||
// something that did not stop the session. A write that fails is not worth
|
||||
// failing over, and when the terminal really is gone the read that follows the
|
||||
// question says so.
|
||||
func say(w io.Writer, format string, args ...any) {
|
||||
_, _ = fmt.Fprintf(w, format, args...)
|
||||
}
|
||||
|
||||
// prompter asks this run's questions on the user's terminal.
|
||||
//
|
||||
// It owns the buffered reader rather than wrapping opts.In per question,
|
||||
// because bufio reads ahead: a throwaway reader for the "start a server?"
|
||||
// question swallows the model choice that was typed behind it, and the next
|
||||
// question then sees EOF. A real run asks both, one after the other.
|
||||
type prompter struct {
|
||||
in *bufio.Reader
|
||||
out io.Writer
|
||||
}
|
||||
|
||||
func newPrompter(in io.Reader, out io.Writer) *prompter {
|
||||
return &prompter{in: bufio.NewReader(in), out: out}
|
||||
}
|
||||
|
||||
// yesNo satisfies Confirmer. Anything that is not an explicit yes is a no, so
|
||||
// a closed stream declines rather than proceeding on the user's behalf.
|
||||
func (p *prompter) yesNo(question string) (bool, error) {
|
||||
say(p.out, "%s [y/N]: ", question)
|
||||
line, err := p.in.ReadString('\n')
|
||||
if err != nil && !errors.Is(err, io.EOF) {
|
||||
return false, fmt.Errorf("reading the answer: %w", err)
|
||||
}
|
||||
switch strings.ToLower(strings.TrimSpace(line)) {
|
||||
case "y", "yes":
|
||||
return true, nil
|
||||
}
|
||||
return false, nil
|
||||
}
|
||||
|
||||
// choose satisfies ModelChooser. It answers with a list index rather than with
|
||||
// what the user typed, so the result can only ever be one of the models it was
|
||||
// offered: a model name is not something to accept unvalidated here, since
|
||||
// ResolveModel persists whatever comes back and every later run then starts
|
||||
// against it.
|
||||
func (p *prompter) choose(models []string) (string, error) {
|
||||
if len(models) == 0 {
|
||||
return "", errors.New("there is nothing to choose from")
|
||||
}
|
||||
say(p.out, "Several models are available:\n")
|
||||
for i, m := range models {
|
||||
say(p.out, " %d) %s\n", i+1, m)
|
||||
}
|
||||
say(p.out, "Pick one [1-%d]: ", len(models))
|
||||
|
||||
line, err := p.in.ReadString('\n')
|
||||
if err != nil && !errors.Is(err, io.EOF) {
|
||||
return "", fmt.Errorf("reading the choice: %w", err)
|
||||
}
|
||||
answer := strings.TrimSpace(line)
|
||||
n, err := strconv.Atoi(answer)
|
||||
if err != nil || n < 1 || n > len(models) {
|
||||
return "", fmt.Errorf("not a valid choice: %q. Pick a number between 1 and %d, or pass --model", answer, len(models))
|
||||
}
|
||||
return models[n-1], nil
|
||||
}
|
||||
@@ -1,629 +0,0 @@
|
||||
package chat
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"syscall"
|
||||
"time"
|
||||
|
||||
"github.com/mudler/nib/app"
|
||||
nibconfig "github.com/mudler/nib/config"
|
||||
nibtypes "github.com/mudler/nib/types"
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
)
|
||||
|
||||
// modelServer answers /v1/models with the given ids, as LocalAI does.
|
||||
func modelServer(ids ...string) *httptest.Server {
|
||||
return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
data := make([]map[string]string, 0, len(ids))
|
||||
for _, id := range ids {
|
||||
data = append(data, map[string]string{"id": id, "object": "model"})
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
Expect(json.NewEncoder(w).Encode(map[string]any{"object": "list", "data": data})).To(Succeed())
|
||||
}))
|
||||
}
|
||||
|
||||
var _ = Describe("prepare", func() {
|
||||
var (
|
||||
dir string
|
||||
errOut *bytes.Buffer
|
||||
)
|
||||
|
||||
BeforeEach(func() {
|
||||
dir = GinkgoT().TempDir()
|
||||
errOut = &bytes.Buffer{}
|
||||
})
|
||||
|
||||
// optionsFor points a run at srv, with no input to read: the default is a
|
||||
// session nobody can be asked anything in.
|
||||
optionsFor := func(srv *httptest.Server) Options {
|
||||
endpoint := "http://127.0.0.1:0"
|
||||
base := endpoint + "/v1"
|
||||
if srv != nil {
|
||||
endpoint, base = srv.URL, srv.URL+"/v1"
|
||||
}
|
||||
return Options{
|
||||
Endpoint: endpoint,
|
||||
BaseURL: base,
|
||||
StateDir: dir,
|
||||
In: strings.NewReader(""),
|
||||
Out: &bytes.Buffer{},
|
||||
ErrOut: errOut,
|
||||
}
|
||||
}
|
||||
|
||||
It("uses the only model the server offers", func() {
|
||||
srv := modelServer("the-only-model")
|
||||
defer srv.Close()
|
||||
|
||||
p, err := prepare(context.Background(), optionsFor(srv), false)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(p.model).To(Equal("the-only-model"))
|
||||
Expect(p.dir).To(Equal(dir))
|
||||
Expect(p.server).To(BeNil(), "nothing was started, so nothing is owned")
|
||||
})
|
||||
|
||||
It("seeds the agent config with the endpoint on first run", func() {
|
||||
srv := modelServer("m")
|
||||
defer srv.Close()
|
||||
|
||||
_, err := prepare(context.Background(), optionsFor(srv), false)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
data, err := os.ReadFile(ConfigPath(dir))
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(string(data)).To(ContainSubstring(srv.URL + "/v1"))
|
||||
})
|
||||
|
||||
It("lets --model win over what the server offers", func() {
|
||||
srv := modelServer("a", "b")
|
||||
defer srv.Close()
|
||||
|
||||
opts := optionsFor(srv)
|
||||
opts.Model = "not-listed-yet"
|
||||
p, err := prepare(context.Background(), opts, false)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(p.model).To(Equal("not-listed-yet"))
|
||||
})
|
||||
|
||||
It("advises about the API key when the server rejects it", func() {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(http.StatusUnauthorized)
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
_, err := prepare(context.Background(), optionsFor(srv), false)
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("--api-key"))
|
||||
Expect(err.Error()).To(ContainSubstring(srv.URL))
|
||||
})
|
||||
|
||||
// Not interactive means nobody can answer the offer, so the advice has to
|
||||
// stand on its own.
|
||||
It("advises how to start a server when none is reachable", func() {
|
||||
srv := modelServer()
|
||||
url := srv.URL
|
||||
srv.Close() // nothing is listening now
|
||||
|
||||
opts := optionsFor(nil)
|
||||
opts.Endpoint, opts.BaseURL = url, url+"/v1"
|
||||
_, err := prepare(context.Background(), opts, false)
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("local-ai run"))
|
||||
Expect(err.Error()).To(ContainSubstring(url))
|
||||
})
|
||||
|
||||
// A server that accepts the connection and then never replies is the case
|
||||
// the offer to start one exists for, so the budget has to expire as a
|
||||
// deadline: Probe reads a cancellation as "the caller gave up" and refuses
|
||||
// to call the endpoint unreachable on the strength of it.
|
||||
It("treats a server that never answers as one that is not there", func(ctx SpecContext) {
|
||||
release := make(chan struct{})
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
select {
|
||||
case <-release:
|
||||
case <-r.Context().Done():
|
||||
}
|
||||
}))
|
||||
defer srv.Close()
|
||||
defer close(release)
|
||||
|
||||
opts := optionsFor(srv)
|
||||
opts.ProbeTimeout = 100 * time.Millisecond
|
||||
_, err := prepare(context.Background(), opts, false)
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("local-ai run"), "want the offer-a-server advice, got %v", err)
|
||||
}, SpecTimeout(30*time.Second))
|
||||
|
||||
It("asks which model to use and remembers the answer", func() {
|
||||
srv := modelServer("zeta", "alpha")
|
||||
defer srv.Close()
|
||||
|
||||
opts := optionsFor(srv)
|
||||
opts.In = strings.NewReader("2\n")
|
||||
p, err := prepare(context.Background(), opts, true)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
// The list is sorted before it is shown, so 2 is zeta, not the second
|
||||
// thing the server happened to name.
|
||||
Expect(p.model).To(Equal("zeta"))
|
||||
Expect(errOut.String()).To(ContainSubstring("1) alpha"))
|
||||
Expect(errOut.String()).To(ContainSubstring("2) zeta"))
|
||||
|
||||
data, err := os.ReadFile(ConfigPath(dir))
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(string(data)).To(ContainSubstring("zeta"))
|
||||
})
|
||||
|
||||
// The choice is prompted for once and remembered. When remembering it fails
|
||||
// the user is about to be asked again on every future run, so they have to
|
||||
// be told here: a log line is invisible at the default log level.
|
||||
It("says so on the prompt when the choice cannot be remembered", func() {
|
||||
srv := modelServer("zeta", "alpha")
|
||||
defer srv.Close()
|
||||
|
||||
// A directory where the config file belongs: writable state dir,
|
||||
// unwritable config, on any platform and as any user.
|
||||
Expect(os.MkdirAll(ConfigPath(dir), 0o700)).To(Succeed())
|
||||
|
||||
opts := optionsFor(srv)
|
||||
opts.In = strings.NewReader("1\n")
|
||||
p, err := prepare(context.Background(), opts, true)
|
||||
|
||||
// Failing to remember the choice must not cost the user their session.
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(p.model).To(Equal("alpha"))
|
||||
Expect(errOut.String()).To(ContainSubstring("could not be saved"), "the user has to learn they will be asked again")
|
||||
})
|
||||
|
||||
It("does not ask again once a model is recorded", func() {
|
||||
srv := modelServer("zeta", "alpha")
|
||||
defer srv.Close()
|
||||
|
||||
Expect(PersistModel(dir, "alpha")).To(Succeed())
|
||||
|
||||
opts := optionsFor(srv)
|
||||
opts.In = strings.NewReader("") // an answer would have nothing to read
|
||||
p, err := prepare(context.Background(), opts, true)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(p.model).To(Equal("alpha"))
|
||||
Expect(errOut.String()).To(BeEmpty())
|
||||
})
|
||||
|
||||
It("says what to install when the server has no models", func() {
|
||||
srv := modelServer()
|
||||
defer srv.Close()
|
||||
|
||||
_, err := prepare(context.Background(), optionsFor(srv), false)
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("models install"))
|
||||
})
|
||||
|
||||
Describe("arguments that only touch local state", func() {
|
||||
unreachable := func(args ...string) Options {
|
||||
opts := optionsFor(nil) // port 0: nothing can ever answer here
|
||||
opts.Args = args
|
||||
return opts
|
||||
}
|
||||
|
||||
DescribeTable("skips the server entirely",
|
||||
func(args ...string) {
|
||||
p, err := prepare(context.Background(), unreachable(args...), false)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(p.model).To(BeEmpty())
|
||||
Expect(p.server).To(BeNil())
|
||||
},
|
||||
Entry("plugin", "plugin", "list"),
|
||||
Entry("skill", "skill", "list"),
|
||||
Entry("mcp add", "mcp", "add", "srv"),
|
||||
Entry("mcp list", "mcp", "list"),
|
||||
// The shell snippet is what a user puts in their rc file, long
|
||||
// before any server exists.
|
||||
Entry("the shell integration script", "--init", "zsh"),
|
||||
Entry("the version", "--version"),
|
||||
)
|
||||
|
||||
// Bare 'mcp' and its transport flags serve the agent over MCP, so they
|
||||
// need a model like any other session. Only the verbs that edit the
|
||||
// configured servers are local.
|
||||
DescribeTable("still needs a server",
|
||||
func(args ...string) {
|
||||
_, err := prepare(context.Background(), unreachable(args...), false)
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("local-ai run"))
|
||||
},
|
||||
Entry("mcp over stdio", "mcp", "--stdio"),
|
||||
Entry("bare mcp", "mcp"),
|
||||
)
|
||||
})
|
||||
|
||||
// A reader per question would read ahead into a buffer it then discards, so
|
||||
// the second question would see EOF whenever both answers were typed ahead.
|
||||
// That is the shape of a real run: the offer to start a server is followed
|
||||
// by the model prompt.
|
||||
It("keeps reading answers from the same stream across questions", func() {
|
||||
out := &bytes.Buffer{}
|
||||
p := newPrompter(strings.NewReader("y\n2\n"), out)
|
||||
|
||||
yes, err := p.yesNo("Start one now?")
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(yes).To(BeTrue())
|
||||
|
||||
chosen, err := p.choose([]string{"alpha", "zeta"})
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(chosen).To(Equal("zeta"))
|
||||
})
|
||||
|
||||
// Whatever the chooser returns is persisted and used for every later run,
|
||||
// so an answer that is not one of the offered models must never come back
|
||||
// as one.
|
||||
Describe("the model prompt", func() {
|
||||
offered := []string{"alpha", "zeta"}
|
||||
|
||||
DescribeTable("refuses an answer that is not one of the numbers shown",
|
||||
func(answer string) {
|
||||
chosen, err := newPrompter(strings.NewReader(answer), &bytes.Buffer{}).choose(offered)
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(chosen).To(BeEmpty())
|
||||
},
|
||||
Entry("nothing at all", ""),
|
||||
Entry("a blank line", "\n"),
|
||||
Entry("only spaces", " \n"),
|
||||
Entry("zero", "0\n"),
|
||||
Entry("past the end", "3\n"),
|
||||
Entry("negative", "-1\n"),
|
||||
Entry("a model name", "zeta\n"),
|
||||
Entry("a number with a suffix", "1x\n"),
|
||||
)
|
||||
|
||||
It("says how to answer when the answer was not a number", func() {
|
||||
_, err := newPrompter(strings.NewReader("banana\n"), &bytes.Buffer{}).choose(offered)
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("between 1 and 2"))
|
||||
Expect(err.Error()).To(ContainSubstring("--model"))
|
||||
})
|
||||
|
||||
It("returns the model shown against the number", func() {
|
||||
chosen, err := newPrompter(strings.NewReader("1\n"), &bytes.Buffer{}).choose(offered)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(chosen).To(Equal("alpha"))
|
||||
})
|
||||
|
||||
It("refuses to ask when there is nothing to offer", func() {
|
||||
chosen, err := newPrompter(strings.NewReader("1\n"), &bytes.Buffer{}).choose(nil)
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(chosen).To(BeEmpty())
|
||||
})
|
||||
})
|
||||
|
||||
// A server started for this session is stopped by a deferred call, which a
|
||||
// signal skips: the process dies where it stands and leaves 'local-ai run'
|
||||
// reparented to init.
|
||||
Describe("shutdown signals", func() {
|
||||
It("ends the session when the terminal goes away", func() {
|
||||
ctx, stop := shutdownContext(context.Background())
|
||||
defer stop()
|
||||
|
||||
self, err := os.FindProcess(os.Getpid())
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(self.Signal(syscall.SIGHUP)).To(Succeed())
|
||||
|
||||
Eventually(ctx.Done()).WithTimeout(5 * time.Second).Should(BeClosed())
|
||||
Expect(ctx.Err()).To(MatchError(context.Canceled))
|
||||
})
|
||||
|
||||
// SIGINT and SIGTERM cannot be delivered here to prove the same thing:
|
||||
// Ginkgo registers for both to abort the suite, and a signal goes to
|
||||
// every registered listener.
|
||||
It("also listens for an interrupt and a terminate", func() {
|
||||
Expect(shutdownSignals).To(ContainElements(os.Signal(os.Interrupt), os.Signal(syscall.SIGTERM)))
|
||||
})
|
||||
})
|
||||
|
||||
// Cancelling the context does unwind nib's TUI since v0.5.1, but how long
|
||||
// that takes is nib's business, and the deferred Stop in Run is only reached
|
||||
// once the agent returns. A server this process started is ours to end, so
|
||||
// the guarantee is made here instead, where it does not depend on the agent
|
||||
// at all. Before v0.5.1 there was no guarantee to be had on the SIGHUP path:
|
||||
// bubbletea's own SIGINT and SIGTERM handler was the only thing that ever
|
||||
// quit the program, and registering for SIGHUP took away the default
|
||||
// disposition that used to end the process.
|
||||
Describe("runSession", func() {
|
||||
It("stops the session's server on cancellation, without waiting for the agent", func() {
|
||||
server, proc := stoppableServer()
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
|
||||
err := runSession(ctx, server, func(ctx context.Context) error {
|
||||
cancel()
|
||||
Eventually(func() int32 { return proc.interrupts.Load() }).
|
||||
WithTimeout(5 * time.Second).
|
||||
Should(BeNumerically(">", 0), "the server has to be stopped while the agent is still running")
|
||||
return nil
|
||||
})
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(proc.lastSignal.Load()).To(Equal(os.Interrupt))
|
||||
})
|
||||
|
||||
It("leaves the server alone for as long as the session lasts", func() {
|
||||
server, proc := stoppableServer()
|
||||
|
||||
Expect(runSession(context.Background(), server, func(context.Context) error {
|
||||
return nil
|
||||
})).To(Succeed())
|
||||
Expect(proc.interrupts.Load()).To(BeZero())
|
||||
Expect(proc.kills.Load()).To(BeZero())
|
||||
})
|
||||
|
||||
It("returns what the agent returned", func() {
|
||||
failed := errors.New("the agent gave up")
|
||||
server, _ := stoppableServer()
|
||||
|
||||
Expect(runSession(context.Background(), server, func(context.Context) error {
|
||||
return failed
|
||||
})).To(MatchError(failed))
|
||||
})
|
||||
|
||||
// Most sessions run against a server the user already had, and there is
|
||||
// nothing to stop then.
|
||||
It("copes with a session that started no server", func() {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
|
||||
Expect(runSession(ctx, nil, func(context.Context) error {
|
||||
return nil
|
||||
})).To(Succeed())
|
||||
})
|
||||
})
|
||||
|
||||
// Which streams reach nib decides two user-visible behaviours at once, and
|
||||
// they pull in opposite directions, so both are pinned here rather than left
|
||||
// to whoever next edits the literal.
|
||||
//
|
||||
// nib refuses every mode but --cli when a stream it was handed is not a
|
||||
// terminal. That refusal is wanted for stdin, where it is what tells someone
|
||||
// piping a question to re-run with --cli. It is not wanted for the process
|
||||
// stdout, where it would refuse the Ctrl+Space widget that --init emits:
|
||||
// out=$(local-ai chat --height 50%) puts a pipe on stdout by construction,
|
||||
// and writing the chosen command into that pipe is the entire point.
|
||||
Describe("agentOptions", func() {
|
||||
// optionsWithStreams is a request that differs from the next only in
|
||||
// what it was told to read and write.
|
||||
optionsWithStreams := func(in io.Reader, out, errOut io.Writer) Options {
|
||||
return Options{
|
||||
BaseURL: "http://127.0.0.1:8080/v1",
|
||||
In: in,
|
||||
Out: out,
|
||||
ErrOut: errOut,
|
||||
}
|
||||
}
|
||||
|
||||
Describe("stdout", func() {
|
||||
// The regression this exists to catch: reinstating
|
||||
// 'Stdout: opts.Out' breaks Ctrl+Space and nothing else notices.
|
||||
It("hands nib nothing for the process stdout, so the capture widget is not refused", func() {
|
||||
o := agentOptions(dir, "a-model", optionsWithStreams(os.Stdin, os.Stdout, os.Stderr))
|
||||
Expect(o.Stdout).To(BeNil(), "injecting os.Stdout is what refuses out=$(local-ai chat)")
|
||||
})
|
||||
|
||||
It("keeps a stdout the caller chose, which the refusal still guards", func() {
|
||||
out := &bytes.Buffer{}
|
||||
o := agentOptions(dir, "a-model", optionsWithStreams(os.Stdin, out, os.Stderr))
|
||||
Expect(o.Stdout).To(BeIdenticalTo(out))
|
||||
})
|
||||
|
||||
// Being an *os.File is not what makes a stream nib's own; being the
|
||||
// process stdout is. This is a file an in-process caller opened for
|
||||
// itself, not one a shell redirect handed over as stdout, which
|
||||
// still arrives as os.Stdout and is still nil-ed. It was never going
|
||||
// to receive the interface, so it stays injected and stays refused.
|
||||
It("keeps a file that is not the process stdout", func() {
|
||||
f, err := os.CreateTemp(GinkgoT().TempDir(), "captured")
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
DeferCleanup(f.Close)
|
||||
|
||||
o := agentOptions(dir, "a-model", optionsWithStreams(os.Stdin, f, os.Stderr))
|
||||
Expect(o.Stdout).To(BeIdenticalTo(f))
|
||||
})
|
||||
})
|
||||
|
||||
Describe("stdin", func() {
|
||||
// The opposite regression: nilling stdin the way stdout is nilled
|
||||
// would silently drop the refusal that names --cli.
|
||||
It("hands the process stdin over, so a piped session is still refused", func() {
|
||||
o := agentOptions(dir, "a-model", optionsWithStreams(os.Stdin, os.Stdout, os.Stderr))
|
||||
Expect(o.Stdin).To(BeIdenticalTo(os.Stdin))
|
||||
})
|
||||
|
||||
It("hands over a stdin the caller chose", func() {
|
||||
in := strings.NewReader("a question")
|
||||
o := agentOptions(dir, "a-model", optionsWithStreams(in, os.Stdout, os.Stderr))
|
||||
Expect(o.Stdin).To(BeIdenticalTo(in))
|
||||
})
|
||||
})
|
||||
|
||||
// nib gates stdin and stdout and nothing else, so there is no reason to
|
||||
// hide the error stream from it.
|
||||
It("hands the error stream over whatever it is", func() {
|
||||
errOut := &bytes.Buffer{}
|
||||
o := agentOptions(dir, "a-model", optionsWithStreams(os.Stdin, os.Stdout, errOut))
|
||||
Expect(o.Stderr).To(BeIdenticalTo(errOut))
|
||||
|
||||
o = agentOptions(dir, "a-model", optionsWithStreams(os.Stdin, os.Stdout, os.Stderr))
|
||||
Expect(o.Stderr).To(BeIdenticalTo(os.Stderr))
|
||||
})
|
||||
|
||||
It("names the command a user would type, not the binary nib ships as", func() {
|
||||
o := agentOptions(dir, "a-model", optionsWithStreams(os.Stdin, os.Stdout, os.Stderr))
|
||||
Expect(o.ProgramName).To(Equal("local-ai chat"),
|
||||
"the --init widget invokes this name, so a user has to be able to run it")
|
||||
})
|
||||
|
||||
It("carries the resolved session through to nib", func() {
|
||||
opts := optionsWithStreams(os.Stdin, os.Stdout, os.Stderr)
|
||||
opts.Args = []string{"--cli"}
|
||||
opts.APIKey = "a-key"
|
||||
opts.TraceDir = "/traces"
|
||||
|
||||
o := agentOptions(dir, "the-model", opts)
|
||||
Expect(o.Args).To(Equal([]string{"--cli"}))
|
||||
Expect(o.BaseDir).To(Equal(dir))
|
||||
Expect(o.Overrides.Model).To(Equal("the-model"))
|
||||
Expect(o.Overrides.APIKey).To(Equal("a-key"))
|
||||
Expect(o.Overrides.BaseURL).To(Equal("http://127.0.0.1:8080/v1"))
|
||||
Expect(o.Overrides.TraceDir).To(Equal("/traces"))
|
||||
// The model and the server are settled before nib starts, and the
|
||||
// bare MODEL and API_KEY variables belong to some other tool.
|
||||
Expect(o.SkipSetup).To(BeTrue())
|
||||
Expect(o.SkipBareEnv).To(BeTrue())
|
||||
})
|
||||
|
||||
// Defaults sit beneath the config file. Anything routed through them is
|
||||
// accepted from the command line and then thrown away the moment the
|
||||
// file carries the same key, which is the normal state rather than an
|
||||
// edge case. Nothing this command resolves belongs there, so the channel
|
||||
// stays empty and this says so: it is what fails if the block is moved
|
||||
// back a rung.
|
||||
It("seeds nothing, because a seed is not a flag", func() {
|
||||
opts := optionsWithStreams(os.Stdin, os.Stdout, os.Stderr)
|
||||
opts.APIKey = "a-key"
|
||||
opts.TraceDir = "/traces"
|
||||
opts.Yolo = true
|
||||
|
||||
Expect(agentOptions(dir, "the-model", opts).Defaults).To(Equal(nibtypes.Config{}),
|
||||
"Defaults lose to the config file, so a value placed there is a flag that does nothing")
|
||||
})
|
||||
|
||||
It("asks for automatic approval only when --yolo was given", func() {
|
||||
opts := optionsWithStreams(os.Stdin, os.Stdout, os.Stderr)
|
||||
Expect(agentOptions(dir, "a-model", opts).Overrides.ApprovalMode).To(BeEmpty())
|
||||
|
||||
opts.Yolo = true
|
||||
Expect(agentOptions(dir, "a-model", opts).Overrides.ApprovalMode).To(Equal("auto"))
|
||||
})
|
||||
|
||||
// The specs above pin what is handed over. These pin what nib does with
|
||||
// it, which is the part that was wrong: every value below reached
|
||||
// app.Options intact and was then discarded by the config load, so a
|
||||
// spec that stops at the struct cannot see the bug. Resolving the config
|
||||
// the way app.Run resolves it can.
|
||||
Describe("the config nib actually resolves", func() {
|
||||
// writeConfig puts a config file where nib will read it, with values
|
||||
// that disagree with every flag under test.
|
||||
writeConfig := func(body string) {
|
||||
Expect(os.WriteFile(ConfigPath(dir), []byte(body), 0o600)).To(Succeed())
|
||||
}
|
||||
|
||||
// resolve loads the config exactly as app.Run does, so the precedence
|
||||
// under test is nib's own rather than a restatement of it here.
|
||||
resolve := func(o app.Options) nibtypes.Config {
|
||||
return nibconfig.LoadWith(nibconfig.LoadOptions{
|
||||
BaseDir: o.BaseDir,
|
||||
Defaults: o.Defaults,
|
||||
Overrides: o.Overrides,
|
||||
SkipBareEnv: o.SkipBareEnv,
|
||||
})
|
||||
}
|
||||
|
||||
It("sends the requests to the endpoint the flag named, not the one on disk", func() {
|
||||
writeConfig("base_url: http://127.0.0.1:9999/v1\n")
|
||||
|
||||
opts := optionsWithStreams(os.Stdin, os.Stdout, os.Stderr)
|
||||
opts.BaseURL = "http://127.0.0.1:8080/v1"
|
||||
|
||||
cfg := resolve(agentOptions(dir, "a-model", opts))
|
||||
Expect(cfg.BaseURL).To(Equal("http://127.0.0.1:8080/v1"),
|
||||
"--endpoint probed 8080; every turn has to go there too")
|
||||
})
|
||||
|
||||
It("uses the model the flag named, not the one the picker recorded", func() {
|
||||
writeConfig("model: recorded-model\n")
|
||||
|
||||
cfg := resolve(agentOptions(dir, "flag-model", optionsWithStreams(os.Stdin, os.Stdout, os.Stderr)))
|
||||
Expect(cfg.Model).To(Equal("flag-model"))
|
||||
})
|
||||
|
||||
It("uses the key the flag named, not the one nib saved", func() {
|
||||
writeConfig("api_key: saved-key\n")
|
||||
|
||||
opts := optionsWithStreams(os.Stdin, os.Stdout, os.Stderr)
|
||||
opts.APIKey = "flag-key"
|
||||
|
||||
cfg := resolve(agentOptions(dir, "a-model", opts))
|
||||
Expect(cfg.APIKey).To(Equal("flag-key"))
|
||||
})
|
||||
|
||||
It("turns approval off for --yolo even when the file demands it", func() {
|
||||
writeConfig("approval_mode: prompt\n")
|
||||
|
||||
opts := optionsWithStreams(os.Stdin, os.Stdout, os.Stderr)
|
||||
opts.Yolo = true
|
||||
|
||||
cfg := resolve(agentOptions(dir, "a-model", opts))
|
||||
Expect(cfg.ApprovalMode).To(Equal("auto"))
|
||||
})
|
||||
|
||||
// The other half of the same rule, and the reason an unset flag is
|
||||
// not a demand for the empty string: an override only ever raises a
|
||||
// field, so what the user configured survives a run that said
|
||||
// nothing about it.
|
||||
It("leaves what the file configured alone when no flag was given", func() {
|
||||
writeConfig("api_key: saved-key\napproval_mode: prompt\n")
|
||||
|
||||
cfg := resolve(agentOptions(dir, "a-model", optionsWithStreams(os.Stdin, os.Stdout, os.Stderr)))
|
||||
Expect(cfg.APIKey).To(Equal("saved-key"))
|
||||
Expect(cfg.ApprovalMode).To(Equal("prompt"))
|
||||
})
|
||||
})
|
||||
})
|
||||
|
||||
// nib reports its own failures on the error stream and returns nothing but
|
||||
// a status, so anything that reaches here as one has already been explained
|
||||
// once. The refusal to open a full-screen session on a stdin that cannot be
|
||||
// read is the one users meet: 'echo q | local-ai chat' names --cli, and a
|
||||
// second message on top would bury the fix.
|
||||
Describe("ExitStatus", func() {
|
||||
It("recognises a status the agent already explained", func() {
|
||||
code, reported := ExitStatus(app.ExitError{Code: 2})
|
||||
Expect(reported).To(BeTrue())
|
||||
Expect(code).To(Equal(2))
|
||||
})
|
||||
|
||||
It("finds one that has been wrapped", func() {
|
||||
code, reported := ExitStatus(fmt.Errorf("running the agent: %w", app.ExitError{Code: 1}))
|
||||
Expect(reported).To(BeTrue())
|
||||
Expect(code).To(Equal(1))
|
||||
})
|
||||
|
||||
It("leaves an ordinary failure to be reported", func() {
|
||||
_, reported := ExitStatus(errors.New("no LocalAI server at http://127.0.0.1:8080"))
|
||||
Expect(reported).To(BeFalse())
|
||||
})
|
||||
|
||||
It("says nothing about a run that succeeded", func() {
|
||||
_, reported := ExitStatus(nil)
|
||||
Expect(reported).To(BeFalse())
|
||||
})
|
||||
})
|
||||
|
||||
It("reports a state dir it cannot create", func() {
|
||||
blocked := filepath.Join(dir, "a-file")
|
||||
Expect(os.WriteFile(blocked, []byte("not a dir"), 0o600)).To(Succeed())
|
||||
|
||||
opts := optionsFor(nil)
|
||||
opts.StateDir = filepath.Join(blocked, "chat")
|
||||
_, err := prepare(context.Background(), opts, false)
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("agent state dir"))
|
||||
})
|
||||
})
|
||||
@@ -1,276 +0,0 @@
|
||||
package chat
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"os"
|
||||
"os/exec"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/mudler/LocalAI/pkg/httpclient"
|
||||
)
|
||||
|
||||
// ErrDeclined means no server was started, either because the session is not
|
||||
// interactive or because the user said no.
|
||||
var ErrDeclined = errors.New("no server started")
|
||||
|
||||
// errServerExited means the process we spawned died before it ever reported
|
||||
// ready, so there is no point in polling out the rest of the budget.
|
||||
var errServerExited = errors.New("the LocalAI server exited before it became ready")
|
||||
|
||||
const (
|
||||
// defaultReadyTimeout bounds the wait for a freshly spawned server. A cold
|
||||
// start probes hardware and may pull a backend, so the budget is generous.
|
||||
defaultReadyTimeout = 2 * time.Minute
|
||||
// readyPollInterval is how long to wait between readiness polls.
|
||||
readyPollInterval = 500 * time.Millisecond
|
||||
// readyProbeTimeout bounds a single readiness request, so one connection
|
||||
// that hangs cannot swallow the whole budget.
|
||||
readyProbeTimeout = 5 * time.Second
|
||||
// shutdownGrace is how long a server we started gets to unload models and
|
||||
// stop its backends after SIGINT before it is killed outright.
|
||||
shutdownGrace = 10 * time.Second
|
||||
// childOutputDrainDelay bounds how long cmd.Wait keeps copying the child's
|
||||
// output after the child itself has exited.
|
||||
//
|
||||
// This is not a theoretical guard for LocalAI. 'local-ai run' spawns backend
|
||||
// subprocesses, and they inherit the write end of the pipe exec created for
|
||||
// the child's stderr. A backend that outlives its parent holds that pipe
|
||||
// open, so an unbounded cmd.Wait would block on the copy goroutine long
|
||||
// after the server itself is gone: exited would never close, Stop would burn
|
||||
// its whole grace period even on a clean shutdown, and the waiter goroutine
|
||||
// would leak.
|
||||
//
|
||||
// The value is long enough that a legitimate final burst of logs is never
|
||||
// truncated even on a loaded machine, where the copy itself takes
|
||||
// microseconds. It must stay strictly below shutdownGrace: at or above it,
|
||||
// every wedged-pipe shutdown would exhaust the grace period and then SIGKILL
|
||||
// a process that had already exited cleanly.
|
||||
childOutputDrainDelay = 5 * time.Second
|
||||
)
|
||||
|
||||
// Confirmer asks a yes/no question. Nil means the session is not interactive.
|
||||
type Confirmer func(question string) (bool, error)
|
||||
|
||||
// StartOptions configures OfferToStart.
|
||||
type StartOptions struct {
|
||||
// Endpoint is the address the user expected a server on, used in the
|
||||
// question and polled for readiness. This is the endpoint root, not the
|
||||
// /v1 API base URL: readiness is served at the root.
|
||||
Endpoint string
|
||||
// Confirm asks whether to start a server. Nil means never start.
|
||||
Confirm Confirmer
|
||||
// Stderr receives the child's output.
|
||||
Stderr io.Writer
|
||||
// Executable overrides the binary to run. Empty means os.Executable().
|
||||
Executable string
|
||||
// ReadyTimeout bounds the wait for readiness. Zero means defaultReadyTimeout.
|
||||
ReadyTimeout time.Duration
|
||||
}
|
||||
|
||||
// StartedServer is a server this process started and is responsible for.
|
||||
type StartedServer struct {
|
||||
// exited is closed once the child has been reaped. One background waiter
|
||||
// owns cmd.Wait: it may only be called once, and it is what closes the
|
||||
// pipes exec created for Stdout/Stderr and joins the goroutines copying
|
||||
// them, so calling os.Process.Wait directly instead would leak both.
|
||||
exited chan struct{}
|
||||
// waitErr is the child's exit status. It is written before exited is
|
||||
// closed and must only be read after that channel is observed closed.
|
||||
waitErr error
|
||||
|
||||
// proc is the child. It is an interface rather than *os.Process so that
|
||||
// Stop's contract, in particular that the child is asked to stop exactly
|
||||
// once however often Stop is called, can be pinned without a live process
|
||||
// to signal. Nil means nothing was ever started.
|
||||
proc processControl
|
||||
|
||||
stopOnce sync.Once
|
||||
}
|
||||
|
||||
// processControl is the part of *os.Process that Stop needs.
|
||||
//
|
||||
// One interface rather than a pair of independent function fields: two fields
|
||||
// can be wired to each other's operation, or one left nil, and no test can tell,
|
||||
// because a fake satisfies any combination. There is nothing to swap or forget
|
||||
// here, since the sole implementation is the real process and the method names
|
||||
// carry the meaning.
|
||||
type processControl interface {
|
||||
Signal(os.Signal) error
|
||||
Kill() error
|
||||
}
|
||||
|
||||
// *os.Process satisfies processControl unmodified, so production needs no
|
||||
// adapter and no nil branch: the wiring is a single assignment.
|
||||
var _ processControl = (*os.Process)(nil)
|
||||
|
||||
// newServerCommand builds the child process. Split out from OfferToStart so the
|
||||
// process' configuration can be asserted on without spawning anything.
|
||||
func newServerCommand(bin string, stderr io.Writer) *exec.Cmd {
|
||||
cmd := exec.Command(bin, "run")
|
||||
// Stdin is left nil, so the child gets /dev/null: it is a background
|
||||
// server, and sharing the terminal would have it stealing keystrokes from
|
||||
// the agent.
|
||||
cmd.Stdout = stderr // the child's logs are diagnostics, not chat output
|
||||
cmd.Stderr = stderr
|
||||
// Bound the wait for the child's output pipes; see childOutputDrainDelay.
|
||||
cmd.WaitDelay = childOutputDrainDelay
|
||||
return cmd
|
||||
}
|
||||
|
||||
// OfferToStart asks whether to start a LocalAI server and, if allowed, spawns
|
||||
// one and waits for it to report ready.
|
||||
//
|
||||
// A child process rather than an in-process boot: RunCMD.Run installs its own
|
||||
// signal handling and blocks until shutdown, so re-entering it from a chat
|
||||
// session would entangle two lifecycles in one process.
|
||||
func OfferToStart(ctx context.Context, opts StartOptions) (*StartedServer, error) {
|
||||
if opts.Confirm == nil {
|
||||
// Not interactive. Spawning a server nobody asked for is the one thing
|
||||
// this function must never do: in CI, in a pipeline, or under a
|
||||
// supervisor there is no one to see it or shut it down.
|
||||
return nil, ErrDeclined
|
||||
}
|
||||
ok, err := opts.Confirm(fmt.Sprintf("No LocalAI server at %s. Start one now?", opts.Endpoint))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("asking whether to start a server: %w", err)
|
||||
}
|
||||
if !ok {
|
||||
return nil, ErrDeclined
|
||||
}
|
||||
|
||||
bin := opts.Executable
|
||||
if bin == "" {
|
||||
if bin, err = os.Executable(); err != nil {
|
||||
return nil, fmt.Errorf("locating the local-ai binary: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
cmd := newServerCommand(bin, opts.Stderr)
|
||||
if err := cmd.Start(); err != nil {
|
||||
return nil, fmt.Errorf("starting a LocalAI server with %s: %w", bin, err)
|
||||
}
|
||||
|
||||
s := &StartedServer{exited: make(chan struct{}), proc: cmd.Process}
|
||||
go func() {
|
||||
s.waitErr = cmd.Wait()
|
||||
close(s.exited)
|
||||
}()
|
||||
|
||||
timeout := opts.ReadyTimeout
|
||||
if timeout <= 0 {
|
||||
timeout = defaultReadyTimeout
|
||||
}
|
||||
if err := waitReady(ctx, opts.Endpoint, timeout, s.exited); err != nil {
|
||||
if errors.Is(err, errServerExited) {
|
||||
// Safe to read: errServerExited is only returned once exited has
|
||||
// been observed closed, which happens after waitErr is written.
|
||||
err = describeExit(err, s.waitErr)
|
||||
}
|
||||
s.Stop()
|
||||
return nil, fmt.Errorf("%w. Run 'local-ai run' in another terminal to see why it did not come up", err)
|
||||
}
|
||||
return s, nil
|
||||
}
|
||||
|
||||
// describeExit adds what is known about how the child died to exitErr, without
|
||||
// putting os/exec's plumbing in front of the user.
|
||||
//
|
||||
// waitErr is exec.ErrWaitDelay when the child exited cleanly but something it
|
||||
// spawned still held its output pipe open past childOutputDrainDelay. The
|
||||
// sentinel's own text names the WaitDelay field, which is meaningless to a
|
||||
// user, so it is translated. Nothing is swallowed: os/exec only substitutes
|
||||
// ErrWaitDelay when the process itself exited without an error of its own (see
|
||||
// Cmd.Wait, "Report an error from the copying goroutines only if the program
|
||||
// otherwise exited normally"), so it can never stand in for an *ExitError.
|
||||
func describeExit(exitErr, waitErr error) error {
|
||||
switch {
|
||||
case waitErr == nil:
|
||||
return exitErr
|
||||
case errors.Is(waitErr, exec.ErrWaitDelay):
|
||||
return fmt.Errorf("%w, and left a subprocess of its own still running", exitErr)
|
||||
default:
|
||||
return fmt.Errorf("%w: %w", exitErr, waitErr)
|
||||
}
|
||||
}
|
||||
|
||||
// Stop terminates the server this process started, giving it a chance to shut
|
||||
// down cleanly first. It is safe to call on a nil or never-started server, and
|
||||
// safe to call more than once.
|
||||
func (s *StartedServer) Stop() {
|
||||
if s == nil || s.proc == nil {
|
||||
return
|
||||
}
|
||||
s.stopOnce.Do(func() {
|
||||
// SIGINT rather than SIGKILL: local-ai run installs its own handler and
|
||||
// needs it to unload models and stop backend subprocesses. Killing it
|
||||
// outright would strand those children.
|
||||
_ = s.proc.Signal(os.Interrupt)
|
||||
|
||||
select {
|
||||
case <-s.exited:
|
||||
case <-time.After(shutdownGrace):
|
||||
// It ignored the interrupt or wedged on the way down. The user is
|
||||
// waiting on their shell prompt, so stop being polite.
|
||||
_ = s.proc.Kill()
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// waitReady polls the endpoint's /readyz until the server reports ready, the
|
||||
// budget expires, the caller gives up, or exited signals that the process we
|
||||
// are waiting on is gone. A nil exited channel means there is no process to
|
||||
// watch.
|
||||
//
|
||||
// Readiness lives on the endpoint ROOT, not under the /v1 API base URL, and it
|
||||
// answers 503 for as long as startup is still in progress.
|
||||
func waitReady(ctx context.Context, endpoint string, timeout time.Duration, exited <-chan struct{}) error {
|
||||
url := strings.TrimSuffix(endpoint, "/") + "/readyz"
|
||||
|
||||
// A real deadline rather than context.WithCancel plus a timer: the latter
|
||||
// expires as context.Canceled, which every classifier here reads as "the
|
||||
// caller gave up" rather than "the endpoint never answered".
|
||||
waitCtx, cancel := context.WithTimeout(ctx, timeout)
|
||||
defer cancel()
|
||||
|
||||
client := httpclient.NewWithTimeout(readyProbeTimeout)
|
||||
ticker := time.NewTicker(readyPollInterval)
|
||||
defer ticker.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-exited:
|
||||
return errServerExited
|
||||
case <-waitCtx.Done():
|
||||
// Distinguish our budget from the caller's: only ours is advice
|
||||
// about the server.
|
||||
if err := ctx.Err(); err != nil {
|
||||
return err
|
||||
}
|
||||
return fmt.Errorf("the LocalAI server did not become ready within %s", timeout)
|
||||
case <-ticker.C:
|
||||
}
|
||||
|
||||
req, err := http.NewRequestWithContext(waitCtx, http.MethodGet, url, nil)
|
||||
if err != nil {
|
||||
return fmt.Errorf("building the readiness request for %s: %w", url, err)
|
||||
}
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
continue // nothing listening yet
|
||||
}
|
||||
// Drain before closing so the next poll can reuse the connection
|
||||
// instead of opening a socket every 500ms for two minutes.
|
||||
_, _ = io.Copy(io.Discard, resp.Body)
|
||||
_ = resp.Body.Close()
|
||||
if resp.StatusCode == http.StatusOK {
|
||||
return nil
|
||||
}
|
||||
// Anything else means startup is still in progress; keep polling.
|
||||
}
|
||||
}
|
||||
@@ -1,375 +0,0 @@
|
||||
package chat
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
)
|
||||
|
||||
// unusedPort is a loopback address nothing listens on, used wherever a spec
|
||||
// needs a readiness poll to keep failing. Port 1 is privileged, so no test
|
||||
// process could have bound it.
|
||||
const unusedPort = "http://127.0.0.1:1"
|
||||
|
||||
var _ = Describe("OfferToStart", func() {
|
||||
It("never spawns anything when there is no confirmer", func() {
|
||||
started, err := OfferToStart(context.Background(), StartOptions{
|
||||
Endpoint: "http://127.0.0.1:59999",
|
||||
Confirm: nil,
|
||||
Stderr: io.Discard,
|
||||
Executable: "/nonexistent/binary-that-must-not-run",
|
||||
})
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(errors.Is(err, ErrDeclined)).To(BeTrue(), "want ErrDeclined, got %v", err)
|
||||
Expect(started).To(BeNil())
|
||||
})
|
||||
|
||||
It("does not spawn when the user declines", func() {
|
||||
asked := false
|
||||
started, err := OfferToStart(context.Background(), StartOptions{
|
||||
Endpoint: "http://127.0.0.1:59999",
|
||||
Confirm: func(string) (bool, error) {
|
||||
asked = true
|
||||
return false, nil
|
||||
},
|
||||
Stderr: io.Discard,
|
||||
Executable: "/nonexistent/binary-that-must-not-run",
|
||||
})
|
||||
Expect(asked).To(BeTrue(), "the user should have been asked")
|
||||
Expect(errors.Is(err, ErrDeclined)).To(BeTrue())
|
||||
Expect(started).To(BeNil())
|
||||
})
|
||||
|
||||
It("names the endpoint in the question", func() {
|
||||
var question string
|
||||
_, _ = OfferToStart(context.Background(), StartOptions{
|
||||
Endpoint: "http://example.invalid:9090",
|
||||
Confirm: func(q string) (bool, error) {
|
||||
question = q
|
||||
return false, nil
|
||||
},
|
||||
Stderr: io.Discard,
|
||||
Executable: "/nonexistent/binary-that-must-not-run",
|
||||
})
|
||||
Expect(question).To(ContainSubstring("http://example.invalid:9090"))
|
||||
})
|
||||
|
||||
It("propagates a confirmer error", func() {
|
||||
boom := errors.New("boom")
|
||||
_, err := OfferToStart(context.Background(), StartOptions{
|
||||
Endpoint: "http://127.0.0.1:59999",
|
||||
Confirm: func(string) (bool, error) { return false, boom },
|
||||
Stderr: io.Discard,
|
||||
Executable: "/nonexistent/binary-that-must-not-run",
|
||||
})
|
||||
Expect(errors.Is(err, boom)).To(BeTrue())
|
||||
})
|
||||
|
||||
It("reports which binary it failed to launch", func() {
|
||||
started, err := OfferToStart(context.Background(), StartOptions{
|
||||
Endpoint: "http://127.0.0.1:59999",
|
||||
Confirm: func(string) (bool, error) { return true, nil },
|
||||
Stderr: io.Discard,
|
||||
Executable: "/nonexistent/binary-that-must-not-run",
|
||||
})
|
||||
Expect(started).To(BeNil())
|
||||
Expect(err).To(MatchError(ContainSubstring("starting a LocalAI server")))
|
||||
Expect(err).To(MatchError(ContainSubstring("/nonexistent/binary-that-must-not-run")))
|
||||
})
|
||||
|
||||
It("stops waiting as soon as the process it started exits", func() {
|
||||
// A harmless no-op binary rather than a real server: this exercises the
|
||||
// early-exit path without starting LocalAI, binding a port, or running
|
||||
// 'local-ai run'. Without early-exit detection the call would sit here
|
||||
// polling until ReadyTimeout.
|
||||
bin, lookErr := exec.LookPath("true")
|
||||
if lookErr != nil {
|
||||
Skip("no 'true' binary on PATH to stand in for a server that dies at once")
|
||||
}
|
||||
|
||||
start := time.Now()
|
||||
started, err := OfferToStart(context.Background(), StartOptions{
|
||||
Endpoint: unusedPort,
|
||||
Confirm: func(string) (bool, error) { return true, nil },
|
||||
Stderr: io.Discard,
|
||||
Executable: bin,
|
||||
ReadyTimeout: 30 * time.Second,
|
||||
})
|
||||
Expect(started).To(BeNil())
|
||||
Expect(err).To(MatchError(ContainSubstring("exited before it became ready")))
|
||||
Expect(time.Since(start)).To(BeNumerically("<", 10*time.Second),
|
||||
"the wait should end with the process, not with the readiness budget")
|
||||
})
|
||||
|
||||
It("gives up on a child whose grandchildren still hold its output pipe", func() {
|
||||
// The real LocalAI shape: 'local-ai run' exits but a backend
|
||||
// subprocess it spawned inherited the stderr pipe and keeps it open.
|
||||
// Without cmd.WaitDelay, cmd.Wait blocks on the copy goroutine, exited
|
||||
// never closes, and the readiness wait runs out the full budget instead
|
||||
// of reporting that the server died.
|
||||
sh, lookErr := exec.LookPath("sh")
|
||||
if lookErr != nil {
|
||||
Skip("no 'sh' binary on PATH to stand in for a server with a lingering child")
|
||||
}
|
||||
|
||||
dir := GinkgoT().TempDir()
|
||||
pidFile := filepath.Join(dir, "grandchild.pid")
|
||||
script := filepath.Join(dir, "server-with-lingering-child")
|
||||
// #nosec G306 -- this has to be executable to stand in for a binary.
|
||||
Expect(os.WriteFile(script,
|
||||
[]byte("#!"+sh+"\nsleep 30 &\necho $! > "+pidFile+"\nexit 0\n"),
|
||||
0o700)).To(Succeed())
|
||||
|
||||
// Reap the grandchild whatever happens: it outlives its own parent by
|
||||
// design, so nothing else will clean it up.
|
||||
DeferCleanup(func() {
|
||||
raw, err := os.ReadFile(pidFile)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
pid, err := strconv.Atoi(strings.TrimSpace(string(raw)))
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
proc, err := os.FindProcess(pid)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
_ = proc.Kill()
|
||||
_, _ = proc.Wait()
|
||||
})
|
||||
|
||||
start := time.Now()
|
||||
started, err := OfferToStart(context.Background(), StartOptions{
|
||||
Endpoint: unusedPort,
|
||||
Confirm: func(string) (bool, error) { return true, nil },
|
||||
Stderr: io.Discard,
|
||||
Executable: script,
|
||||
ReadyTimeout: 25 * time.Second,
|
||||
})
|
||||
elapsed := time.Since(start)
|
||||
|
||||
Expect(started).To(BeNil())
|
||||
Expect(err).To(MatchError(ContainSubstring("exited before it became ready")),
|
||||
"an unbounded cmd.Wait would report a readiness timeout instead")
|
||||
Expect(elapsed).To(BeNumerically("<", 20*time.Second),
|
||||
"the wait must be bounded by the output drain, not by the readiness budget")
|
||||
|
||||
// This is the case where cmd.Wait returns exec.ErrWaitDelay, whose own
|
||||
// text names a struct field of os/exec. Users get told what happened
|
||||
// instead.
|
||||
Expect(err).NotTo(MatchError(ContainSubstring("WaitDelay")),
|
||||
"os/exec plumbing must not reach the user")
|
||||
Expect(err).NotTo(MatchError(ContainSubstring("exec:")))
|
||||
Expect(err).To(MatchError(ContainSubstring("left a subprocess of its own still running")))
|
||||
})
|
||||
|
||||
It("reports the exit status of a server that failed outright", func() {
|
||||
// The counterpart to the case above: translating ErrWaitDelay must not
|
||||
// cost a real exit status, which is the one diagnostic worth having.
|
||||
bin, lookErr := exec.LookPath("false")
|
||||
if lookErr != nil {
|
||||
Skip("no 'false' binary on PATH to stand in for a server that fails")
|
||||
}
|
||||
|
||||
_, err := OfferToStart(context.Background(), StartOptions{
|
||||
Endpoint: unusedPort,
|
||||
Confirm: func(string) (bool, error) { return true, nil },
|
||||
Stderr: io.Discard,
|
||||
Executable: bin,
|
||||
ReadyTimeout: 30 * time.Second,
|
||||
})
|
||||
Expect(err).To(MatchError(ContainSubstring("exited before it became ready")))
|
||||
Expect(err).To(MatchError(ContainSubstring("exit status 1")))
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("StartedServer.Stop", func() {
|
||||
It("is a no-op on a server that was never started", func() {
|
||||
var nilServer *StartedServer
|
||||
Expect(nilServer.Stop).NotTo(Panic())
|
||||
Expect((&StartedServer{}).Stop).NotTo(Panic())
|
||||
})
|
||||
|
||||
It("interrupts the child exactly once however often it is called", func() {
|
||||
s, proc := stoppableServer()
|
||||
|
||||
s.Stop()
|
||||
s.Stop()
|
||||
s.Stop()
|
||||
|
||||
Expect(proc.interrupts.Load()).To(Equal(int32(1)),
|
||||
"a second Stop must not signal the child again")
|
||||
Expect(proc.kills.Load()).To(BeZero(), "a child that already exited must not be killed")
|
||||
})
|
||||
|
||||
It("interrupts the child exactly once when called concurrently", func() {
|
||||
// The realistic double-Stop: a deferred Stop on the way out racing the
|
||||
// signal handler that also owns shutting the server down.
|
||||
const callers = 8
|
||||
|
||||
s, proc := stoppableServer()
|
||||
|
||||
var wg sync.WaitGroup
|
||||
wg.Add(callers)
|
||||
for range callers {
|
||||
go func() {
|
||||
defer GinkgoRecover()
|
||||
defer wg.Done()
|
||||
s.Stop()
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
|
||||
Expect(proc.interrupts.Load()).To(Equal(int32(1)))
|
||||
Expect(proc.kills.Load()).To(BeZero())
|
||||
})
|
||||
|
||||
It("asks the child to interrupt rather than killing it outright", func() {
|
||||
// The escalation order is the whole point of the grace period: SIGKILL
|
||||
// first would strand the backend subprocesses local-ai run owns.
|
||||
s, proc := stoppableServer()
|
||||
|
||||
s.Stop()
|
||||
|
||||
Expect(proc.lastSignal.Load()).To(Equal(os.Interrupt))
|
||||
Expect(proc.kills.Load()).To(BeZero())
|
||||
})
|
||||
})
|
||||
|
||||
// countingProcess stands in for the *os.Process that Stop drives, recording
|
||||
// what it was asked to do.
|
||||
type countingProcess struct {
|
||||
interrupts atomic.Int32
|
||||
kills atomic.Int32
|
||||
lastSignal atomic.Value
|
||||
}
|
||||
|
||||
func (p *countingProcess) Signal(sig os.Signal) error {
|
||||
p.interrupts.Add(1)
|
||||
p.lastSignal.Store(sig)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (p *countingProcess) Kill() error {
|
||||
p.kills.Add(1)
|
||||
return nil
|
||||
}
|
||||
|
||||
// stoppableServer builds a StartedServer whose child has already exited, driven
|
||||
// by a countingProcess rather than a real one. Nothing is spawned.
|
||||
func stoppableServer() (*StartedServer, *countingProcess) {
|
||||
proc := &countingProcess{}
|
||||
exited := make(chan struct{})
|
||||
close(exited)
|
||||
return &StartedServer{exited: exited, proc: proc}, proc
|
||||
}
|
||||
|
||||
var _ = Describe("newServerCommand", func() {
|
||||
It("bounds how long it will wait for the child's output pipes", func() {
|
||||
cmd := newServerCommand("/nonexistent/binary-that-must-not-run", io.Discard)
|
||||
|
||||
// An unbounded wait is the failure mode: backend subprocesses inherit
|
||||
// the child's stderr pipe and can hold it open long after the server
|
||||
// itself is gone.
|
||||
Expect(cmd.WaitDelay).To(BeNumerically(">", 0), "cmd.Wait must not be unbounded")
|
||||
Expect(cmd.WaitDelay).To(BeNumerically("<", shutdownGrace),
|
||||
"a drain longer than the shutdown grace would kill a cleanly exited server")
|
||||
})
|
||||
|
||||
It("runs the server subcommand without giving it the terminal", func() {
|
||||
cmd := newServerCommand("/nonexistent/binary-that-must-not-run", io.Discard)
|
||||
|
||||
Expect(cmd.Args).To(Equal([]string{"/nonexistent/binary-that-must-not-run", "run"}))
|
||||
Expect(cmd.Stdin).To(BeNil(), "the child must not compete with the agent for stdin")
|
||||
Expect(cmd.Stdout).NotTo(BeNil())
|
||||
Expect(cmd.Stderr).NotTo(BeNil())
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("waitReady", func() {
|
||||
It("polls /readyz on the endpoint root and returns only once it answers 200", func() {
|
||||
// readyOnPoll is deliberately above 1. A handler that answers 200 to the
|
||||
// first poll cannot tell a correct implementation apart from one that
|
||||
// treats 503 as ready, because both return after a single request; the
|
||||
// poll count is what makes 503-as-ready observable.
|
||||
const readyOnPoll = 3
|
||||
|
||||
var polls atomic.Int32
|
||||
var paths atomic.Value
|
||||
paths.Store("")
|
||||
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
paths.Store(r.URL.Path)
|
||||
if polls.Add(1) < readyOnPoll {
|
||||
// What LocalAI answers while startup is still in progress.
|
||||
w.WriteHeader(http.StatusServiceUnavailable)
|
||||
return
|
||||
}
|
||||
w.WriteHeader(http.StatusOK)
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
Expect(waitReady(context.Background(), srv.URL, 20*time.Second, nil)).To(Succeed())
|
||||
Expect(paths.Load()).To(Equal("/readyz"), "readiness lives on the endpoint root, not under /v1")
|
||||
Expect(polls.Load()).To(BeNumerically(">=", readyOnPoll),
|
||||
"503 means startup is still in progress and must never be accepted as ready")
|
||||
})
|
||||
|
||||
It("tolerates a trailing slash on the endpoint", func() {
|
||||
var path atomic.Value
|
||||
path.Store("")
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
path.Store(r.URL.Path)
|
||||
w.WriteHeader(http.StatusOK)
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
Expect(waitReady(context.Background(), srv.URL+"/", 20*time.Second, nil)).To(Succeed())
|
||||
Expect(path.Load()).To(Equal("/readyz"))
|
||||
})
|
||||
|
||||
It("reports a timeout, not a cancellation, when the budget runs out", func() {
|
||||
err := waitReady(context.Background(), unusedPort, 1200*time.Millisecond, nil)
|
||||
Expect(err).To(HaveOccurred())
|
||||
// A budget built from context.WithCancel plus a timer would surface as
|
||||
// context.Canceled, which downstream code reads as "the caller gave up"
|
||||
// and would stop classifying a hung server as unreachable.
|
||||
Expect(errors.Is(err, context.Canceled)).To(BeFalse(), "got %v", err)
|
||||
Expect(err).To(MatchError(ContainSubstring("did not become ready")))
|
||||
})
|
||||
|
||||
It("returns the caller's cancellation when the caller gives up", func() {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
go func() {
|
||||
defer GinkgoRecover()
|
||||
time.Sleep(200 * time.Millisecond)
|
||||
cancel()
|
||||
}()
|
||||
defer cancel()
|
||||
|
||||
err := waitReady(ctx, unusedPort, time.Minute, nil)
|
||||
Expect(errors.Is(err, context.Canceled)).To(BeTrue(), "got %v", err)
|
||||
})
|
||||
|
||||
It("gives up when the process it is waiting on has exited", func() {
|
||||
exited := make(chan struct{})
|
||||
close(exited)
|
||||
|
||||
err := waitReady(context.Background(), unusedPort, time.Minute, exited)
|
||||
Expect(err).To(MatchError(ContainSubstring("exited before it became ready")))
|
||||
})
|
||||
})
|
||||
112
core/cli/chat/session.go
Normal file
112
core/cli/chat/session.go
Normal file
@@ -0,0 +1,112 @@
|
||||
package chat
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"slices"
|
||||
"strings"
|
||||
)
|
||||
|
||||
const (
|
||||
chatRoleUser = "user"
|
||||
chatRoleAssistant = "assistant"
|
||||
)
|
||||
|
||||
type chatMessage struct {
|
||||
Role string
|
||||
Content string
|
||||
}
|
||||
|
||||
type chatSession struct {
|
||||
client chatClient
|
||||
model string
|
||||
models []string
|
||||
messages []chatMessage
|
||||
}
|
||||
|
||||
func newChatSession(ctx context.Context, client chatClient, requestedModel string) (*chatSession, error) {
|
||||
models, err := client.ListModels(ctx)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("list models: %w", err)
|
||||
}
|
||||
|
||||
model, err := resolveChatModel(requestedModel, models)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return &chatSession{
|
||||
client: client,
|
||||
model: model,
|
||||
models: models,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (s *chatSession) CurrentModel() string {
|
||||
return s.model
|
||||
}
|
||||
|
||||
func (s *chatSession) Models() []string {
|
||||
models := make([]string, len(s.models))
|
||||
copy(models, s.models)
|
||||
return models
|
||||
}
|
||||
|
||||
func (s *chatSession) Clear() {
|
||||
s.messages = nil
|
||||
}
|
||||
|
||||
func (s *chatSession) SwitchModel(model string) error {
|
||||
if !slices.Contains(s.models, model) {
|
||||
return fmt.Errorf("model %q is not available. Use /models to see installed models", model)
|
||||
}
|
||||
s.model = model
|
||||
s.Clear()
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *chatSession) Send(ctx context.Context, prompt string, out io.Writer) error {
|
||||
s.messages = append(s.messages, chatMessage{
|
||||
Role: chatRoleUser,
|
||||
Content: prompt,
|
||||
})
|
||||
|
||||
answer, err := s.client.StreamChat(ctx, s.model, s.messages, out)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
s.messages = append(s.messages, chatMessage{
|
||||
Role: chatRoleAssistant,
|
||||
Content: answer,
|
||||
})
|
||||
return nil
|
||||
}
|
||||
|
||||
func resolveChatModel(requested string, models []string) (string, error) {
|
||||
switch {
|
||||
case requested == "" && len(models) == 0:
|
||||
return "", errors.New(`no chat models are installed.
|
||||
|
||||
Install a model first, for example:
|
||||
local-ai models list
|
||||
local-ai models install <model>
|
||||
local-ai run
|
||||
|
||||
Then start a chat session:
|
||||
local-ai chat --model <model>`)
|
||||
case requested == "" && len(models) == 1:
|
||||
return models[0], nil
|
||||
case requested == "" && len(models) > 1:
|
||||
var b strings.Builder
|
||||
b.WriteString("multiple models are available; choose one with --model:\n")
|
||||
b.WriteString(formatChatModelList(models, ""))
|
||||
return "", errors.New(b.String())
|
||||
case !slices.Contains(models, requested):
|
||||
return "", fmt.Errorf("model %q is not available. Use `local-ai models list` and `local-ai models install <model>`, or pass an installed model with --model", requested)
|
||||
default:
|
||||
return requested, nil
|
||||
}
|
||||
}
|
||||
56
core/cli/chat/session_test.go
Normal file
56
core/cli/chat/session_test.go
Normal file
@@ -0,0 +1,56 @@
|
||||
package chat
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
)
|
||||
|
||||
var _ = Describe("Chat session", func() {
|
||||
It("keeps model switching and message history out of the terminal adapter", func() {
|
||||
client := &fakeChatClient{
|
||||
models: []string{"alpha", "beta"},
|
||||
answer: "pong",
|
||||
}
|
||||
|
||||
session, err := newChatSession(context.Background(), client, "alpha")
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(session.CurrentModel()).To(Equal("alpha"))
|
||||
|
||||
Expect(session.SwitchModel("beta")).To(Succeed())
|
||||
Expect(session.CurrentModel()).To(Equal("beta"))
|
||||
Expect(session.Send(context.Background(), "ping", io.Discard)).To(Succeed())
|
||||
|
||||
Expect(client.requests).To(HaveLen(1))
|
||||
Expect(client.requests[0].model).To(Equal("beta"))
|
||||
Expect(client.requests[0].messages).To(HaveLen(1))
|
||||
Expect(client.requests[0].messages[0].Content).To(Equal("ping"))
|
||||
})
|
||||
})
|
||||
|
||||
type fakeChatClient struct {
|
||||
models []string
|
||||
answer string
|
||||
requests []fakeChatRequest
|
||||
}
|
||||
|
||||
type fakeChatRequest struct {
|
||||
model string
|
||||
messages []chatMessage
|
||||
}
|
||||
|
||||
func (c *fakeChatClient) ListModels(context.Context) ([]string, error) {
|
||||
return c.models, nil
|
||||
}
|
||||
|
||||
func (c *fakeChatClient) StreamChat(_ context.Context, model string, messages []chatMessage, out io.Writer) (string, error) {
|
||||
copied := make([]chatMessage, len(messages))
|
||||
copy(copied, messages)
|
||||
c.requests = append(c.requests, fakeChatRequest{model: model, messages: copied})
|
||||
if _, err := io.WriteString(out, c.answer); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return c.answer, nil
|
||||
}
|
||||
93
core/cli/chat/terminal.go
Normal file
93
core/cli/chat/terminal.go
Normal file
@@ -0,0 +1,93 @@
|
||||
package chat
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"strings"
|
||||
)
|
||||
|
||||
func runTerminalChat(ctx context.Context, session *chatSession, in io.Reader, out io.Writer) error {
|
||||
scanner := bufio.NewScanner(in)
|
||||
scanner.Buffer(make([]byte, 0, 64*1024), 4*1024*1024)
|
||||
|
||||
if err := writeChat(out, "LocalAI chat (%s)\n", session.CurrentModel()); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := writeChat(out, "Type /exit to quit, /clear to reset the conversation, /models to list models.\n"); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
for {
|
||||
if err := writeChat(out, "\n> "); err != nil {
|
||||
return err
|
||||
}
|
||||
if !scanner.Scan() {
|
||||
break
|
||||
}
|
||||
|
||||
prompt := strings.TrimSpace(scanner.Text())
|
||||
switch prompt {
|
||||
case "":
|
||||
continue
|
||||
case "/bye", "/exit", "/quit":
|
||||
return writeChat(out, "bye\n")
|
||||
case "/clear":
|
||||
session.Clear()
|
||||
if err := writeChat(out, "conversation cleared\n"); err != nil {
|
||||
return err
|
||||
}
|
||||
continue
|
||||
case "/models":
|
||||
if err := printChatModels(out, session.Models(), session.CurrentModel()); err != nil {
|
||||
return err
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
if nextModel, ok := strings.CutPrefix(prompt, "/model "); ok {
|
||||
nextModel = strings.TrimSpace(nextModel)
|
||||
if nextModel == "" {
|
||||
if err := writeChat(out, "usage: /model <name>\n"); err != nil {
|
||||
return err
|
||||
}
|
||||
continue
|
||||
}
|
||||
if err := session.SwitchModel(nextModel); err != nil {
|
||||
if writeErr := writeChat(out, "%s\n", err); writeErr != nil {
|
||||
return writeErr
|
||||
}
|
||||
continue
|
||||
}
|
||||
if err := writeChat(out, "switched to %s; conversation cleared\n", session.CurrentModel()); err != nil {
|
||||
return err
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
if err := writeChat(out, "assistant: "); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := session.Send(ctx, prompt, out); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := writeChat(out, "\n"); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
return scanner.Err()
|
||||
}
|
||||
|
||||
func printChatModels(out io.Writer, models []string, current string) error {
|
||||
if len(models) == 0 {
|
||||
return writeChat(out, "no models installed\n")
|
||||
}
|
||||
return writeChat(out, "%s", formatChatModelList(models, current))
|
||||
}
|
||||
|
||||
func writeChat(out io.Writer, format string, args ...any) error {
|
||||
_, err := fmt.Fprintf(out, format, args...)
|
||||
return err
|
||||
}
|
||||
@@ -8,72 +8,18 @@ import (
|
||||
cliContext "github.com/mudler/LocalAI/core/cli/context"
|
||||
)
|
||||
|
||||
// ChatCMD runs the built-in terminal agent. Everything after the first
|
||||
// positional argument is forwarded to the agent verbatim, so its own
|
||||
// subcommands (plugin, skill, mcp) and their flags work unchanged. LocalAI's
|
||||
// own flags must therefore come first.
|
||||
type ChatCMD struct {
|
||||
Model string `short:"m" help:"Model to use. Defaults to the only model the server offers, or asks when there are several"`
|
||||
Endpoint string `env:"LOCALAI_CHAT_ENDPOINT" default:"http://127.0.0.1:8080" help:"LocalAI server endpoint. The /v1 path is added automatically when omitted"`
|
||||
APIKey string `env:"LOCALAI_API_KEY,API_KEY" help:"API key to use when the LocalAI server requires authentication"`
|
||||
ConfigDir string `env:"LOCALAI_CHAT_CONFIG_DIR" help:"Directory holding the agent's config, plugins, and skills. Defaults to ~/.config/localai/chat" type:"path"`
|
||||
TraceDir string `env:"LOCALAI_CHAT_TRACE_DIR" help:"Write a session LLM trace (NDJSON) to this directory" type:"path"`
|
||||
|
||||
CLI bool `help:"Run in plain CLI mode instead of the full-screen interface"`
|
||||
TUI bool `help:"Force the full-screen interface"`
|
||||
Height string `help:"Run as an inline drop-down of this height, e.g. '40%'"`
|
||||
Tmux bool `help:"Run in a tmux split"`
|
||||
NoTmux bool `name:"no-tmux" help:"Never use a tmux split, even inside tmux"`
|
||||
Init string `help:"Print the shell integration script for Ctrl+Space (zsh, bash, or fish)"`
|
||||
Yolo bool `env:"LOCALAI_CHAT_YOLO" help:"Auto-approve every tool call without prompting"`
|
||||
|
||||
Args []string `arg:"" optional:"" passthrough:"" help:"Arguments forwarded to the agent, e.g. 'plugin install <url>', 'skill list', 'mcp add'"`
|
||||
Model string `short:"m" help:"Model name to use. Defaults to the only model returned by the server when exactly one is available"`
|
||||
Endpoint string `env:"LOCALAI_CHAT_ENDPOINT" default:"http://127.0.0.1:8080" help:"LocalAI server endpoint. The /v1 path is added automatically when omitted"`
|
||||
APIKey string `env:"LOCALAI_API_KEY,API_KEY" help:"API key to use when the LocalAI server requires authentication"`
|
||||
}
|
||||
|
||||
func (c *ChatCMD) Run(ctx *cliContext.Context) error {
|
||||
err := chatcli.Run(context.Background(), chatcli.Options{
|
||||
Args: c.agentArgs(),
|
||||
Endpoint: c.Endpoint,
|
||||
BaseURL: chatAPIBaseURL(c.Endpoint),
|
||||
APIKey: c.APIKey,
|
||||
Model: c.Model,
|
||||
StateDir: c.ConfigDir,
|
||||
TraceDir: c.TraceDir,
|
||||
Yolo: c.Yolo,
|
||||
In: os.Stdin,
|
||||
Out: os.Stdout,
|
||||
ErrOut: os.Stderr,
|
||||
return chatcli.Run(context.Background(), chatcli.Options{
|
||||
Model: c.Model,
|
||||
BaseURL: chatAPIBaseURL(c.Endpoint),
|
||||
APIKey: c.APIKey,
|
||||
In: os.Stdin,
|
||||
Out: os.Stdout,
|
||||
})
|
||||
// The agent explains its own failures on stderr and hands back a code, so
|
||||
// carry the code out and leave the explanation to stand alone.
|
||||
if code, reported := chatcli.ExitStatus(err); reported {
|
||||
return ExitCodeError{Code: code}
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
// agentArgs rebuilds the argument vector the agent expects: LocalAI's mode
|
||||
// flags are declared here for discoverability and shell completion, so they
|
||||
// have to be translated back into the agent's own flag names.
|
||||
func (c *ChatCMD) agentArgs() []string {
|
||||
var args []string
|
||||
if c.CLI {
|
||||
args = append(args, "--cli")
|
||||
}
|
||||
if c.TUI {
|
||||
args = append(args, "--tui")
|
||||
}
|
||||
if c.Height != "" {
|
||||
args = append(args, "--height", c.Height)
|
||||
}
|
||||
if c.Tmux {
|
||||
args = append(args, "--tmux")
|
||||
}
|
||||
if c.NoTmux {
|
||||
args = append(args, "--no-tmux")
|
||||
}
|
||||
if c.Init != "" {
|
||||
args = append(args, "--init", c.Init)
|
||||
}
|
||||
return append(args, c.Args...)
|
||||
}
|
||||
|
||||
@@ -1,10 +1,6 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
|
||||
"github.com/alecthomas/kong"
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
)
|
||||
@@ -28,70 +24,4 @@ var _ = Describe("Chat command wiring", func() {
|
||||
Expect(chatAPIBaseURL("http://127.0.0.1:8080/localai")).To(Equal("http://127.0.0.1:8080/localai/v1"))
|
||||
})
|
||||
})
|
||||
|
||||
Describe("argument parsing", func() {
|
||||
parse := func(args ...string) *ChatCMD {
|
||||
var cli struct {
|
||||
Chat ChatCMD `cmd:""`
|
||||
}
|
||||
parser, err := kong.New(&cli)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
_, err = parser.Parse(append([]string{"chat"}, args...))
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
return &cli.Chat
|
||||
}
|
||||
|
||||
It("leaves Args empty for a bare invocation", func() {
|
||||
Expect(parse().Args).To(BeEmpty())
|
||||
})
|
||||
|
||||
It("binds flags that precede the forwarded arguments", func() {
|
||||
c := parse("--endpoint", "http://host:9090", "--model", "m", "plugin", "list")
|
||||
Expect(c.Endpoint).To(Equal("http://host:9090"))
|
||||
Expect(c.Model).To(Equal("m"))
|
||||
Expect(c.Args).To(Equal([]string{"plugin", "list"}))
|
||||
})
|
||||
|
||||
It("forwards flags that follow the first positional to the agent", func() {
|
||||
c := parse("plugin", "install", "https://example.invalid/p", "--yes")
|
||||
Expect(c.Args).To(Equal([]string{"plugin", "install", "https://example.invalid/p", "--yes"}))
|
||||
})
|
||||
|
||||
It("parses its own mode flags", func() {
|
||||
c := parse("--cli")
|
||||
Expect(c.CLI).To(BeTrue())
|
||||
Expect(c.Args).To(BeEmpty())
|
||||
})
|
||||
})
|
||||
|
||||
// The agent prints its own diagnosis and hands back a status. main exits
|
||||
// with that status and prints nothing more, so the user reads one message
|
||||
// rather than an "exit status 1" stacked under it.
|
||||
Describe("ExitCodeError", func() {
|
||||
It("carries the status out", func() {
|
||||
Expect(ExitCodeError{Code: 2}.Code).To(Equal(2))
|
||||
})
|
||||
|
||||
It("is recognisable after wrapping", func() {
|
||||
var got ExitCodeError
|
||||
Expect(errors.As(fmt.Errorf("chat: %w", ExitCodeError{Code: 2}), &got)).To(BeTrue())
|
||||
Expect(got.Code).To(Equal(2))
|
||||
})
|
||||
})
|
||||
|
||||
Describe("agentArgs", func() {
|
||||
It("translates mode flags into the agent's own flags", func() {
|
||||
c := &ChatCMD{CLI: true}
|
||||
Expect(c.agentArgs()).To(Equal([]string{"--cli"}))
|
||||
})
|
||||
|
||||
It("puts forwarded arguments after the translated flags", func() {
|
||||
c := &ChatCMD{Height: "40%", Args: []string{"plugin", "list"}}
|
||||
Expect(c.agentArgs()).To(Equal([]string{"--height", "40%", "plugin", "list"}))
|
||||
})
|
||||
|
||||
It("returns nothing for a bare invocation", func() {
|
||||
Expect((&ChatCMD{}).agentArgs()).To(BeEmpty())
|
||||
})
|
||||
})
|
||||
})
|
||||
|
||||
@@ -9,7 +9,7 @@ var CLI struct {
|
||||
cliContext.Context `embed:""`
|
||||
|
||||
Run RunCMD `cmd:"" help:"Run LocalAI, this the default command if no other command is specified. Run 'local-ai run --help' for more information" default:"withargs"`
|
||||
Chat ChatCMD `cmd:"" help:"Run the built-in terminal agent against a LocalAI server"`
|
||||
Chat ChatCMD `cmd:"" help:"Open an interactive chat session against a running LocalAI server"`
|
||||
Federated FederatedCLI `cmd:"" help:"Run LocalAI in federated mode"`
|
||||
Models ModelsCMD `cmd:"" help:"Manage LocalAI models and definitions"`
|
||||
Backends BackendsCMD `cmd:"" help:"Manage LocalAI backends and definitions"`
|
||||
|
||||
@@ -1,15 +0,0 @@
|
||||
package cli
|
||||
|
||||
import "fmt"
|
||||
|
||||
// ExitCodeError is a failure a command has already reported to the user. It
|
||||
// carries nothing but the status the process should exit with, and main prints
|
||||
// nothing more for it.
|
||||
//
|
||||
// It exists for commands that hand their terminal to something that does its
|
||||
// own error reporting. Returning that subordinate's error instead would put a
|
||||
// bare "exit status 1" underneath the explanation the user has just read, and
|
||||
// returning nil would tell a script the run succeeded.
|
||||
type ExitCodeError struct{ Code int }
|
||||
|
||||
func (e ExitCodeError) Error() string { return fmt.Sprintf("exit status %d", e.Code) }
|
||||
@@ -44,7 +44,6 @@ const (
|
||||
MethodPredictStream GRPCMethod = "PredictStream"
|
||||
MethodEmbedding GRPCMethod = "Embedding"
|
||||
MethodGenerateImage GRPCMethod = "GenerateImage"
|
||||
MethodUpscaleImage GRPCMethod = "UpscaleImage"
|
||||
MethodGenerateVideo GRPCMethod = "GenerateVideo"
|
||||
MethodGenerate3D GRPCMethod = "Generate3D"
|
||||
MethodAudioTranscription GRPCMethod = "AudioTranscription"
|
||||
@@ -349,7 +348,7 @@ var BackendCapabilities = map[string]BackendCapability{
|
||||
|
||||
// --- Image/video generation backends ---
|
||||
"diffusers": {
|
||||
GRPCMethods: []GRPCMethod{MethodGenerateImage, MethodUpscaleImage, MethodGenerateVideo},
|
||||
GRPCMethods: []GRPCMethod{MethodGenerateImage, MethodGenerateVideo},
|
||||
PossibleUsecases: []string{UsecaseImage, UsecaseVideo},
|
||||
DefaultUsecases: []string{UsecaseImage},
|
||||
Description: "HuggingFace diffusers — Stable Diffusion, Flux, video generation",
|
||||
@@ -420,47 +419,6 @@ var BackendCapabilities = map[string]BackendCapability{
|
||||
DefaultUsecases: []string{UsecaseTranscript},
|
||||
Description: "NVIDIA NeMo Parakeet ASR (parakeet.cpp)",
|
||||
},
|
||||
// nemo-speech-cpp is one gRPC server in front of four NeMo-Speech.cpp model
|
||||
// families, picked at load time from the GGUF general.architecture key, so
|
||||
// PossibleUsecases is their UNION and no single model serves all of it: an
|
||||
// asr model transcribes (and diarizes, when a Sortformer model is attached
|
||||
// through options), a sortformer model only diarizes, a magpietts model only
|
||||
// synthesizes, and a Riva-Translate model only answers Predict.
|
||||
//
|
||||
// UsecaseChat sits alongside UsecaseCompletion for the translation family
|
||||
// because Predict and PredictStream are exactly the RPCs /v1/chat/completions
|
||||
// drives, and chat is what a translation model is useful through: each turn
|
||||
// goes in as the prompt and comes back translated. The flag is not a gate on
|
||||
// any endpoint (a request naming the model explicitly is served either way);
|
||||
// what it buys is being eligible as the default chat model when a request
|
||||
// names none (core/http/routes/openai.go) and appearing in the React UI's
|
||||
// chat model picker (CAP_CHAT in react-ui/src/utils/capabilities.js).
|
||||
//
|
||||
// Leaving it out is not neutral: chat is a gallery filter key and completion
|
||||
// is not (usecaseFilters in core/http/routes/ui_api.go), so
|
||||
// GET /api/backends/usecases would grey the Chat filter out and hide a
|
||||
// Riva-Translate gallery entry from the one filter that fits it.
|
||||
//
|
||||
// DefaultUsecases is transcript alone because that is the only family whose
|
||||
// weights a bare `backend: nemo-speech-cpp` config is likely to name; a model
|
||||
// of any other family should pin its own known_usecases.
|
||||
//
|
||||
// No VoiceCloning key: MagpieTTS synthesizes from baked speaker ids, not from
|
||||
// a reference clip, so advertising cloning would accept a `voice:
|
||||
// "profile:<id>"` request the backend cannot serve.
|
||||
"nemo-speech-cpp": {
|
||||
GRPCMethods: []GRPCMethod{
|
||||
MethodAudioTranscription, MethodDiarize,
|
||||
MethodTTS, MethodTTSStream,
|
||||
MethodPredict, MethodPredictStream,
|
||||
},
|
||||
PossibleUsecases: []string{
|
||||
UsecaseTranscript, UsecaseDiarization, UsecaseTTS,
|
||||
UsecaseCompletion, UsecaseChat,
|
||||
},
|
||||
DefaultUsecases: []string{UsecaseTranscript},
|
||||
Description: "NVIDIA NeMo-Speech.cpp: one server for Nemotron ASR (offline, streaming and live), Sortformer diarization, MagpieTTS synthesis and Riva-Translate translation; the model's GGUF architecture decides which",
|
||||
},
|
||||
"qwen-asr": {
|
||||
GRPCMethods: []GRPCMethod{MethodAudioTranscription},
|
||||
PossibleUsecases: []string{UsecaseTranscript},
|
||||
|
||||
@@ -124,39 +124,6 @@ var _ = Describe("GetBackendCapability", func() {
|
||||
})
|
||||
})
|
||||
|
||||
// nemo-speech-cpp fronts four model families from one server, and its
|
||||
// PossibleUsecases is their union. The entry has to stay in step with what
|
||||
// docs/content/features/nemo-speech-cpp.md tells operators to put in
|
||||
// known_usecases: nothing validates known_usecases against PossibleUsecases, so
|
||||
// a flag the docs recommend and the map omits fails silently, and the place it
|
||||
// surfaces is the gallery. GET /api/backends/usecases is derived from this list
|
||||
// and greys out the filters missing from it, so a recommended-but-unlisted flag
|
||||
// hides the very models it was recommended for.
|
||||
var _ = Describe("nemo-speech-cpp capabilities", func() {
|
||||
It("advertises every usecase its four families serve", func() {
|
||||
capability := GetBackendCapability("nemo-speech-cpp")
|
||||
Expect(capability).NotTo(BeNil())
|
||||
Expect(capability.PossibleUsecases).To(ContainElements(
|
||||
UsecaseTranscript, UsecaseDiarization, UsecaseTTS,
|
||||
UsecaseCompletion, UsecaseChat))
|
||||
})
|
||||
|
||||
// Chat is the translation family's usecase, and it needs both Predict RPCs:
|
||||
// /v1/chat/completions streams through PredictStream and answers
|
||||
// non-streaming requests through Predict.
|
||||
It("backs the chat usecase with the RPCs chat actually drives", func() {
|
||||
capability := GetBackendCapability("nemo-speech-cpp")
|
||||
Expect(capability.GRPCMethods).To(ContainElements(MethodPredict, MethodPredictStream))
|
||||
})
|
||||
|
||||
// Defaults stay conservative: a bare `backend: nemo-speech-cpp` with no
|
||||
// known_usecases is overwhelmingly an ASR model, and every other family is
|
||||
// expected to pin its own flags.
|
||||
It("still defaults to transcript alone", func() {
|
||||
Expect(DefaultUsecasesForBackendCap("nemo-speech-cpp")).To(Equal([]string{UsecaseTranscript}))
|
||||
})
|
||||
})
|
||||
|
||||
// audio-cpp advertises voice cloning from the backend itself and ships
|
||||
// audio-cpp-chatterbox, whose family serves cloning and NOT plain TTS, so a
|
||||
// reference clip is the only way to use it. Without a capability entry
|
||||
|
||||
@@ -1,117 +0,0 @@
|
||||
package config
|
||||
|
||||
// Speculative-decoding auto-defaults for the vllm-cpp backend, the safetensors
|
||||
// counterpart of the GGUF/llama.cpp hook in mtp.go.
|
||||
//
|
||||
// The two engines detect and spell the same feature differently. llama.cpp
|
||||
// reads `<arch>.nextn_predict_layers` out of the GGUF header and takes
|
||||
// `spec_type:draft-mtp` in `options:`; vllm.cpp reads `mtp_num_hidden_layers`
|
||||
// out of the checkpoint's config.json and takes vLLM's own
|
||||
// `--speculative-config` JSON, which LocalAI carries in `engine_args`. The
|
||||
// engine resolves the draft depth and the default k itself, so the config only
|
||||
// has to name the method.
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
|
||||
"github.com/mudler/xlog"
|
||||
)
|
||||
|
||||
// hfSpecConfig is the subset of a HuggingFace config.json that decides whether
|
||||
// speculative decoding can be auto-enabled.
|
||||
type hfSpecConfig struct {
|
||||
ModelType string `json:"model_type"`
|
||||
// MtpNumHiddenLayers is the MTP head depth (upstream speculative.py reads
|
||||
// it as n_predict for the qwen3_5 / qwen3_5_moe families).
|
||||
MtpNumHiddenLayers uint32 `json:"mtp_num_hidden_layers"`
|
||||
// DFlashConfig marks a z-lab DFlash DRAFT checkpoint (mask_token_id +
|
||||
// target_layer_ids). Its presence means this repo is a draft, not a
|
||||
// servable target.
|
||||
DFlashConfig json.RawMessage `json:"dflash_config"`
|
||||
// TextConfig is where multimodal checkpoints nest the language-model
|
||||
// config, and therefore the MTP depth.
|
||||
TextConfig *hfSpecConfig `json:"text_config"`
|
||||
}
|
||||
|
||||
// parseHFSpecConfig decodes the speculative-relevant subset of a config.json.
|
||||
// A document that does not parse yields nothing rather than an error: detection
|
||||
// is best-effort and must never break an import.
|
||||
func parseHFSpecConfig(configJSON []byte) (hfSpecConfig, bool) {
|
||||
if len(configJSON) == 0 {
|
||||
return hfSpecConfig{}, false
|
||||
}
|
||||
var c hfSpecConfig
|
||||
if err := json.Unmarshal(configJSON, &c); err != nil {
|
||||
xlog.Debug("[vllm-spec] config.json did not parse; skipping detection", "error", err)
|
||||
return hfSpecConfig{}, false
|
||||
}
|
||||
return c, true
|
||||
}
|
||||
|
||||
// IsDFlashDraftConfig reports whether a HuggingFace config.json describes a
|
||||
// DFlash DRAFT checkpoint. Unlike MTP - whose head ships inside the target
|
||||
// checkpoint's `mtp.*` tensors - a DFlash draft is its own repo that can only
|
||||
// run paired with a target it verifies against, so it must never be configured
|
||||
// as a standalone model.
|
||||
func IsDFlashDraftConfig(configJSON []byte) bool {
|
||||
c, ok := parseHFSpecConfig(configJSON)
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
return len(c.DFlashConfig) > 0 ||
|
||||
(c.TextConfig != nil && len(c.TextConfig.DFlashConfig) > 0)
|
||||
}
|
||||
|
||||
// HasSafetensorsMTPHead reports whether a HuggingFace config.json declares a
|
||||
// self-speculating Multi-Token Prediction head, returning its depth. The depth
|
||||
// is informational: vllm.cpp resolves n_predict and the default
|
||||
// num_speculative_tokens from the checkpoint itself.
|
||||
//
|
||||
// DFlash drafts are excluded for the same reason `gemma4-assistant` GGUFs are
|
||||
// excluded from the llama.cpp hook: they carry head metadata but cannot
|
||||
// self-speculate.
|
||||
//
|
||||
// NOTE this is a safetensors-only signal. vllm.cpp rejects an MTP config over a
|
||||
// GGUF source, because the `mtp.*` draft tensors only exist in the safetensors
|
||||
// checkpoint - so the GGUF import path must not use this.
|
||||
func HasSafetensorsMTPHead(configJSON []byte) (uint32, bool) {
|
||||
c, ok := parseHFSpecConfig(configJSON)
|
||||
if !ok {
|
||||
return 0, false
|
||||
}
|
||||
if IsDFlashDraftConfig(configJSON) {
|
||||
return 0, false
|
||||
}
|
||||
n := c.MtpNumHiddenLayers
|
||||
if n == 0 && c.TextConfig != nil {
|
||||
n = c.TextConfig.MtpNumHiddenLayers
|
||||
}
|
||||
return n, n > 0
|
||||
}
|
||||
|
||||
// ApplyVLLMSpeculativeDefaults enables MTP speculative decoding in cfg's
|
||||
// engine_args when nothing is configured there yet. It is a no-op when the user
|
||||
// already set a speculative_config, so an explicit choice (a different method,
|
||||
// an explicit k, a DFlash draft) is never clobbered.
|
||||
//
|
||||
// `layers` is the detected head depth and is only used for the diagnostic log
|
||||
// line - the engine derives the real k from the checkpoint.
|
||||
func ApplyVLLMSpeculativeDefaults(cfg *ModelConfig, layers uint32) {
|
||||
if cfg == nil {
|
||||
return
|
||||
}
|
||||
if _, set := cfg.EngineArgs["speculative_config"]; set {
|
||||
xlog.Debug("[vllm-spec] MTP head detected but speculative_config already configured; leaving user choice intact",
|
||||
"name", cfg.Name, "mtp_num_hidden_layers", layers)
|
||||
return
|
||||
}
|
||||
if cfg.EngineArgs == nil {
|
||||
cfg.EngineArgs = map[string]any{}
|
||||
}
|
||||
// Only the method: vllm.cpp defaults num_speculative_tokens to the
|
||||
// checkpoint's own n_predict (speculative.py:865-875), which is the right
|
||||
// value far more reliably than anything guessable here.
|
||||
cfg.EngineArgs["speculative_config"] = map[string]any{"method": "mtp"}
|
||||
xlog.Info("[vllm-spec] MTP head detected; enabling mtp speculative decoding",
|
||||
"name", cfg.Name, "mtp_num_hidden_layers", layers)
|
||||
}
|
||||
@@ -1,117 +0,0 @@
|
||||
package config_test
|
||||
|
||||
import (
|
||||
. "github.com/mudler/LocalAI/core/config"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
)
|
||||
|
||||
var _ = Describe("vllm-cpp speculative-decoding auto-defaults", func() {
|
||||
Context("HasSafetensorsMTPHead", func() {
|
||||
It("detects a top-level mtp_num_hidden_layers", func() {
|
||||
n, ok := HasSafetensorsMTPHead([]byte(`{
|
||||
"model_type": "qwen3_5_moe",
|
||||
"mtp_num_hidden_layers": 1
|
||||
}`))
|
||||
Expect(ok).To(BeTrue())
|
||||
Expect(n).To(Equal(uint32(1)))
|
||||
})
|
||||
|
||||
It("detects the head nested under text_config", func() {
|
||||
// Multimodal checkpoints nest the language-model config, which is
|
||||
// where the MTP depth lives (mirrors the engine's own resolution
|
||||
// off config.raw text_config).
|
||||
n, ok := HasSafetensorsMTPHead([]byte(`{
|
||||
"model_type": "qwen3_5_moe",
|
||||
"text_config": {"mtp_num_hidden_layers": 2}
|
||||
}`))
|
||||
Expect(ok).To(BeTrue())
|
||||
Expect(n).To(Equal(uint32(2)))
|
||||
})
|
||||
|
||||
It("reports no head when the key is absent", func() {
|
||||
n, ok := HasSafetensorsMTPHead([]byte(`{"model_type": "llama"}`))
|
||||
Expect(ok).To(BeFalse())
|
||||
Expect(n).To(BeZero())
|
||||
})
|
||||
|
||||
It("reports no head for a zero depth", func() {
|
||||
_, ok := HasSafetensorsMTPHead([]byte(`{"mtp_num_hidden_layers": 0}`))
|
||||
Expect(ok).To(BeFalse())
|
||||
})
|
||||
|
||||
It("ignores a DFlash draft checkpoint", func() {
|
||||
// A DFlash draft is a SEPARATE checkpoint that cannot serve alone:
|
||||
// it needs a target to verify against. Same exclusion the GGUF path
|
||||
// makes for gemma4-assistant drafts.
|
||||
_, ok := HasSafetensorsMTPHead([]byte(`{
|
||||
"model_type": "qwen3_dflash",
|
||||
"mtp_num_hidden_layers": 1,
|
||||
"dflash_config": {"mask_token_id": 151666, "target_layer_ids": [0, 1]}
|
||||
}`))
|
||||
Expect(ok).To(BeFalse())
|
||||
})
|
||||
|
||||
It("reports no head on unparseable JSON", func() {
|
||||
_, ok := HasSafetensorsMTPHead([]byte(`{not json`))
|
||||
Expect(ok).To(BeFalse())
|
||||
})
|
||||
|
||||
It("reports no head on empty input", func() {
|
||||
_, ok := HasSafetensorsMTPHead(nil)
|
||||
Expect(ok).To(BeFalse())
|
||||
})
|
||||
})
|
||||
|
||||
Context("IsDFlashDraftConfig", func() {
|
||||
It("recognises a draft by its dflash_config block", func() {
|
||||
Expect(IsDFlashDraftConfig([]byte(`{
|
||||
"dflash_config": {"mask_token_id": 151666, "target_layer_ids": [0]}
|
||||
}`))).To(BeTrue())
|
||||
})
|
||||
|
||||
It("does not flag an ordinary checkpoint", func() {
|
||||
Expect(IsDFlashDraftConfig([]byte(`{"model_type": "qwen3_5_moe"}`))).To(BeFalse())
|
||||
})
|
||||
})
|
||||
|
||||
Context("ApplyVLLMSpeculativeDefaults", func() {
|
||||
It("writes the mtp method into engine_args", func() {
|
||||
cfg := &ModelConfig{Name: "qwen"}
|
||||
ApplyVLLMSpeculativeDefaults(cfg, 1)
|
||||
Expect(cfg.EngineArgs).To(HaveKey("speculative_config"))
|
||||
spec, ok := cfg.EngineArgs["speculative_config"].(map[string]any)
|
||||
Expect(ok).To(BeTrue())
|
||||
Expect(spec["method"]).To(Equal("mtp"))
|
||||
})
|
||||
|
||||
It("leaves an existing speculative_config alone", func() {
|
||||
cfg := &ModelConfig{
|
||||
Name: "qwen",
|
||||
LLMConfig: LLMConfig{
|
||||
EngineArgs: map[string]any{
|
||||
"speculative_config": map[string]any{"method": "ngram", "num_speculative_tokens": 4},
|
||||
},
|
||||
},
|
||||
}
|
||||
ApplyVLLMSpeculativeDefaults(cfg, 1)
|
||||
spec := cfg.EngineArgs["speculative_config"].(map[string]any)
|
||||
Expect(spec["method"]).To(Equal("ngram"))
|
||||
})
|
||||
|
||||
It("preserves unrelated engine_args keys", func() {
|
||||
cfg := &ModelConfig{
|
||||
Name: "qwen",
|
||||
LLMConfig: LLMConfig{EngineArgs: map[string]any{"max_num_seqs": 32}},
|
||||
}
|
||||
ApplyVLLMSpeculativeDefaults(cfg, 1)
|
||||
Expect(cfg.EngineArgs).To(HaveKeyWithValue("max_num_seqs", 32))
|
||||
Expect(cfg.EngineArgs).To(HaveKey("speculative_config"))
|
||||
})
|
||||
|
||||
It("tolerates a nil config", func() {
|
||||
Expect(func() { ApplyVLLMSpeculativeDefaults(nil, 1) }).ToNot(Panic())
|
||||
})
|
||||
})
|
||||
})
|
||||
@@ -1,212 +0,0 @@
|
||||
package gallery
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/mudler/LocalAI/core/config"
|
||||
"github.com/mudler/LocalAI/pkg/concurrency"
|
||||
"github.com/mudler/LocalAI/pkg/system"
|
||||
"github.com/mudler/LocalAI/pkg/vram"
|
||||
"github.com/mudler/xlog"
|
||||
)
|
||||
|
||||
// EstimateInput builds the VRAM estimator's input from a gallery entry.
|
||||
//
|
||||
// It lives here rather than beside the HTTP handler because two callers need
|
||||
// it: the handler answering one model, and the warmer below answering all of
|
||||
// them ahead of time.
|
||||
func EstimateInput(m *GalleryModel) vram.ModelEstimateInput {
|
||||
var input vram.ModelEstimateInput
|
||||
input.Size = m.Size
|
||||
if repoID := extractHFRepo(m.Overrides, m.URLs); repoID != "" {
|
||||
input.HFRepo = repoID
|
||||
}
|
||||
for _, f := range m.AdditionalFiles {
|
||||
if vram.IsWeightFile(f.URI) {
|
||||
input.Files = append(input.Files, vram.FileInput{URI: f.URI, Size: 0})
|
||||
}
|
||||
}
|
||||
return input
|
||||
}
|
||||
|
||||
// extractHFRepo finds a HuggingFace repo ID in a model's overrides or URLs.
|
||||
func extractHFRepo(overrides map[string]any, urls []string) string {
|
||||
if overrides != nil {
|
||||
if params, ok := overrides["parameters"].(map[string]any); ok {
|
||||
if modelRef, ok := params["model"].(string); ok {
|
||||
if repoID, ok := vram.ExtractHFRepoID(modelRef); ok {
|
||||
return repoID
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
for _, u := range urls {
|
||||
if repoID, ok := vram.ExtractHFRepoID(u); ok {
|
||||
return repoID
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// EstimateWarmConfig bounds the background warm-up.
|
||||
type EstimateWarmConfig struct {
|
||||
// Limit is how many gallery entries to warm, in gallery order. Zero
|
||||
// disables warming entirely. The order matters: it is the order the UI
|
||||
// lists them in, so the entries a user sees first are warmed first.
|
||||
Limit int
|
||||
// Concurrency is how many estimates run at once. Each one can be a remote
|
||||
// probe, so this is deliberately small: the point is to be finished before
|
||||
// anybody looks, not to saturate the link or the upstream.
|
||||
Concurrency int
|
||||
// Contexts are the context lengths to estimate at. These want to match what
|
||||
// the UI asks for, or the warmed entry is not the one it reads.
|
||||
Contexts []uint32
|
||||
}
|
||||
|
||||
// DefaultEstimateWarmConfig is what the server uses unless told otherwise.
|
||||
//
|
||||
// The limit is a deliberate compromise. Warming the whole gallery would be
|
||||
// thousands of remote probes on every boot, which is rude to the upstream and
|
||||
// slow to finish; warming nothing leaves the first page of the model gallery
|
||||
// paying two seconds per row. A few hundred covers what anyone browses in a
|
||||
// sitting, and everything past it still warms itself on first view.
|
||||
var DefaultEstimateWarmConfig = EstimateWarmConfig{
|
||||
Limit: 300,
|
||||
Concurrency: 4,
|
||||
Contexts: []uint32{8192, 16384, 32768, 65536, 131072, 262144},
|
||||
}
|
||||
|
||||
// WarmEstimateCache fills the gallery's derived caches in the background.
|
||||
//
|
||||
// Two things are warmed, and they are the same cost wearing different hats.
|
||||
// An estimate for an entry the server has never seen costs a network probe of
|
||||
// its weight files, and describing an entry's variants costs one probe per
|
||||
// build it offers. The UI asks for an estimate per row and a variant
|
||||
// description per model opened, so without this the first visitor pays for
|
||||
// both: ten seconds of a page filling in its own sizes, then another second
|
||||
// and a half the first time they click anything.
|
||||
//
|
||||
// Both land in the same caches underneath, which is why one pass covers them.
|
||||
//
|
||||
// It returns immediately; the work happens on its own goroutine and stops when
|
||||
// ctx is done. Failures are logged at debug and otherwise ignored: a warm-up
|
||||
// that cannot reach an upstream must never stop the server from starting, and
|
||||
// the entry it failed on simply stays cold.
|
||||
func WarmEstimateCache(ctx context.Context, galleries []config.Gallery, systemState *system.SystemState, cfg EstimateWarmConfig) {
|
||||
if cfg.Limit <= 0 || cfg.Concurrency <= 0 {
|
||||
return
|
||||
}
|
||||
|
||||
concurrency.SafeGo(func() {
|
||||
started := time.Now()
|
||||
|
||||
models, err := AvailableGalleryModelsCached(galleries, systemState)
|
||||
if err != nil {
|
||||
xlog.Debug("VRAM estimate warm-up skipped, gallery unavailable", "error", err)
|
||||
return
|
||||
}
|
||||
if len(models) > cfg.Limit {
|
||||
models = models[:cfg.Limit]
|
||||
}
|
||||
if len(models) == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
// The host gate the variant picker resolves against. Derived once: it
|
||||
// describes this machine, not this entry, and HostResolveEnv reads the
|
||||
// system state to build it.
|
||||
env := HostResolveEnv(ctx, systemState)
|
||||
|
||||
var (
|
||||
wg sync.WaitGroup
|
||||
cursor = make(chan *GalleryModel)
|
||||
warmed int
|
||||
warmedVariants int
|
||||
mu sync.Mutex
|
||||
)
|
||||
|
||||
for i := 0; i < cfg.Concurrency; i++ {
|
||||
wg.Add(1)
|
||||
concurrency.SafeGo(func() {
|
||||
defer wg.Done()
|
||||
for m := range cursor {
|
||||
// Per entry, not for the run: one unreachable weight file
|
||||
// must not hold a worker for the whole warm-up.
|
||||
entryCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
|
||||
|
||||
input := EstimateInput(m)
|
||||
if len(input.Files) > 0 || input.HFRepo != "" || input.Size != "" {
|
||||
if _, err := vram.EstimateModelMultiContext(entryCtx, input, cfg.Contexts); err != nil {
|
||||
xlog.Debug("VRAM estimate warm-up failed for entry", "model", m.GetName(), "error", err)
|
||||
} else {
|
||||
mu.Lock()
|
||||
warmed++
|
||||
mu.Unlock()
|
||||
}
|
||||
}
|
||||
|
||||
// Describing variants probes each build the entry offers.
|
||||
// An entry that declares none costs nothing here, so this is
|
||||
// gated rather than attempted and discarded.
|
||||
if m.HasVariants() {
|
||||
if _, err := DescribeVariants(models, m, env); err != nil {
|
||||
xlog.Debug("variant warm-up failed for entry", "model", m.GetName(), "error", err)
|
||||
} else {
|
||||
mu.Lock()
|
||||
warmedVariants++
|
||||
mu.Unlock()
|
||||
}
|
||||
}
|
||||
|
||||
cancel()
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
feed:
|
||||
for _, m := range models {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
break feed
|
||||
case cursor <- m:
|
||||
}
|
||||
}
|
||||
close(cursor)
|
||||
wg.Wait()
|
||||
|
||||
if ctx.Err() != nil {
|
||||
xlog.Debug("gallery warm-up stopped", "estimates", warmed, "variants", warmedVariants)
|
||||
return
|
||||
}
|
||||
xlog.Info("gallery caches warmed", "estimates", warmed, "variants", warmedVariants, "of", len(models), "took", time.Since(started).Round(time.Second))
|
||||
})
|
||||
}
|
||||
|
||||
// EstimateWarmConfigFromEnv reads the warm-up bounds from the environment,
|
||||
// falling back to the defaults.
|
||||
//
|
||||
// LOCALAI_VRAM_WARM_LIMIT entries to warm; 0 disables the warm-up
|
||||
// LOCALAI_VRAM_WARM_CONCURRENCY estimates in flight at once
|
||||
//
|
||||
// Env rather than a flag because it is an operational tuning knob, not part of
|
||||
// what the server does: an air-gapped host wants it off, and a host behind a
|
||||
// slow link wants it slower, and neither is a decision the CLI should carry.
|
||||
func EstimateWarmConfigFromEnv() EstimateWarmConfig {
|
||||
cfg := DefaultEstimateWarmConfig
|
||||
if v, ok := os.LookupEnv("LOCALAI_VRAM_WARM_LIMIT"); ok {
|
||||
if n, err := strconv.Atoi(strings.TrimSpace(v)); err == nil && n >= 0 {
|
||||
cfg.Limit = n
|
||||
}
|
||||
}
|
||||
if v, ok := os.LookupEnv("LOCALAI_VRAM_WARM_CONCURRENCY"); ok {
|
||||
if n, err := strconv.Atoi(strings.TrimSpace(v)); err == nil && n > 0 {
|
||||
cfg.Concurrency = n
|
||||
}
|
||||
}
|
||||
return cfg
|
||||
}
|
||||
@@ -1,180 +0,0 @@
|
||||
package gallery_test
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/binary"
|
||||
"math"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"time"
|
||||
|
||||
gguf "github.com/gpustack/gguf-parser-go"
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
"gopkg.in/yaml.v3"
|
||||
|
||||
"github.com/mudler/LocalAI/core/config"
|
||||
"github.com/mudler/LocalAI/core/gallery"
|
||||
"github.com/mudler/LocalAI/pkg/system"
|
||||
)
|
||||
|
||||
var _ = Describe("VRAM estimate warm-up", func() {
|
||||
var state *system.SystemState
|
||||
|
||||
BeforeEach(func() {
|
||||
dir, err := os.MkdirTemp("", "warm")
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
DeferCleanup(func() { os.RemoveAll(dir) })
|
||||
state, err = system.GetSystemState(system.WithModelPath(dir))
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
gallery.ResetGalleryModelCache()
|
||||
DeferCleanup(gallery.ResetGalleryModelCache)
|
||||
})
|
||||
|
||||
It("does nothing when disabled, and returns without blocking", func() {
|
||||
cfg := gallery.DefaultEstimateWarmConfig
|
||||
cfg.Limit = 0
|
||||
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
defer close(done)
|
||||
gallery.WarmEstimateCache(context.Background(), []config.Gallery{}, state, cfg)
|
||||
}()
|
||||
Eventually(done, "1s").Should(BeClosed())
|
||||
})
|
||||
|
||||
It("returns immediately even when there is work to do", func() {
|
||||
// The caller is a server still starting up: warming must never be on
|
||||
// the path to listening.
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
defer close(done)
|
||||
gallery.WarmEstimateCache(context.Background(), []config.Gallery{}, state, gallery.DefaultEstimateWarmConfig)
|
||||
}()
|
||||
Eventually(done, "1s").Should(BeClosed())
|
||||
})
|
||||
|
||||
It("stops when its context is cancelled", func() {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
gallery.WarmEstimateCache(ctx, []config.Gallery{}, state, gallery.DefaultEstimateWarmConfig)
|
||||
cancel()
|
||||
// Nothing to assert beyond not hanging or panicking: an aborted warm-up
|
||||
// leaves entries cold, which is the state they were already in.
|
||||
Consistently(func() bool { return true }, "100ms").Should(BeTrue())
|
||||
})
|
||||
|
||||
It("does not crash the server when remote GGUF metadata is malformed", func() {
|
||||
payload := warmMalformedGGUF()
|
||||
requested := make(chan struct{})
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
select {
|
||||
case <-requested:
|
||||
default:
|
||||
close(requested)
|
||||
}
|
||||
http.ServeContent(w, r, "model.gguf", time.Time{}, bytes.NewReader(payload))
|
||||
}))
|
||||
DeferCleanup(server.Close)
|
||||
|
||||
galleryPath := filepath.Join(state.Model.ModelsPath, "malformed-gallery.yaml")
|
||||
index, err := yaml.Marshal([]gallery.GalleryModel{{Metadata: gallery.Metadata{
|
||||
Name: "malformed-gguf",
|
||||
AdditionalFiles: []gallery.File{{
|
||||
Filename: "model.gguf",
|
||||
URI: server.URL + "/model.gguf",
|
||||
}},
|
||||
}}})
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(os.WriteFile(galleryPath, index, 0600)).To(Succeed())
|
||||
|
||||
cfg := gallery.DefaultEstimateWarmConfig
|
||||
cfg.Limit = 1
|
||||
cfg.Concurrency = 1
|
||||
cfg.Contexts = []uint32{8192}
|
||||
gallery.WarmEstimateCache(context.Background(), []config.Gallery{{
|
||||
Name: "malformed",
|
||||
URL: "file://" + galleryPath,
|
||||
}}, state, cfg)
|
||||
|
||||
Eventually(requested, "2s").Should(BeClosed())
|
||||
// The warm-up is detached. Give its parser time to consume the response;
|
||||
// before the recovery boundary, that goroutine panicked and killed the
|
||||
// entire test process (and the LocalAI server in production).
|
||||
Consistently(func() bool { return true }, "300ms").Should(BeTrue())
|
||||
})
|
||||
|
||||
Describe("configuration from the environment", func() {
|
||||
AfterEach(func() {
|
||||
os.Unsetenv("LOCALAI_VRAM_WARM_LIMIT")
|
||||
os.Unsetenv("LOCALAI_VRAM_WARM_CONCURRENCY")
|
||||
})
|
||||
|
||||
It("falls back to the defaults", func() {
|
||||
cfg := gallery.EstimateWarmConfigFromEnv()
|
||||
Expect(cfg.Limit).To(Equal(gallery.DefaultEstimateWarmConfig.Limit))
|
||||
Expect(cfg.Concurrency).To(Equal(gallery.DefaultEstimateWarmConfig.Concurrency))
|
||||
})
|
||||
|
||||
It("lets an operator turn it off entirely", func() {
|
||||
os.Setenv("LOCALAI_VRAM_WARM_LIMIT", "0")
|
||||
Expect(gallery.EstimateWarmConfigFromEnv().Limit).To(BeZero())
|
||||
})
|
||||
|
||||
It("lets an operator slow it down", func() {
|
||||
os.Setenv("LOCALAI_VRAM_WARM_CONCURRENCY", "1")
|
||||
Expect(gallery.EstimateWarmConfigFromEnv().Concurrency).To(Equal(1))
|
||||
})
|
||||
|
||||
It("ignores values that are not usable", func() {
|
||||
os.Setenv("LOCALAI_VRAM_WARM_LIMIT", "not-a-number")
|
||||
os.Setenv("LOCALAI_VRAM_WARM_CONCURRENCY", "0")
|
||||
cfg := gallery.EstimateWarmConfigFromEnv()
|
||||
Expect(cfg.Limit).To(Equal(gallery.DefaultEstimateWarmConfig.Limit))
|
||||
// Zero workers would be a warm-up that never runs while looking
|
||||
// enabled, so it keeps the default rather than honouring it.
|
||||
Expect(cfg.Concurrency).To(Equal(gallery.DefaultEstimateWarmConfig.Concurrency))
|
||||
})
|
||||
})
|
||||
|
||||
It("warms variant descriptions as well as estimates", func() {
|
||||
// Both are the same cost wearing different hats - a probe of an entry's
|
||||
// weight files - and both land in the same caches, so a warm-up that
|
||||
// covered only one would leave the first click paying for the other.
|
||||
// Asserted through the shared config rather than by observing network
|
||||
// calls: the gallery here is empty by design.
|
||||
Expect(gallery.DefaultEstimateWarmConfig.Limit).To(BeNumerically(">", 0))
|
||||
})
|
||||
|
||||
It("keeps the estimate contexts the UI actually asks for", func() {
|
||||
// A warmed entry at the wrong context lengths is a cache the gallery
|
||||
// never reads, so this pins them together.
|
||||
Expect(gallery.DefaultEstimateWarmConfig.Contexts).To(ContainElements(
|
||||
uint32(8192), uint32(16384), uint32(32768), uint32(65536), uint32(131072), uint32(262144),
|
||||
))
|
||||
})
|
||||
|
||||
It("bounds concurrency so a warm-up cannot saturate the link", func() {
|
||||
Expect(gallery.DefaultEstimateWarmConfig.Concurrency).To(BeNumerically("<=", 8))
|
||||
Expect(gallery.DefaultEstimateWarmConfig.Concurrency).To(BeNumerically(">", 0))
|
||||
})
|
||||
|
||||
})
|
||||
|
||||
func warmMalformedGGUF() []byte {
|
||||
payload := make([]byte, 0, 128)
|
||||
payload = binary.LittleEndian.AppendUint32(payload, uint32(gguf.GGUFMagicGGUFLe))
|
||||
payload = binary.LittleEndian.AppendUint32(payload, uint32(gguf.GGUFVersionV3))
|
||||
payload = binary.LittleEndian.AppendUint64(payload, 0)
|
||||
payload = binary.LittleEndian.AppendUint64(payload, 1)
|
||||
key := "tokenizer.ggml.tokens"
|
||||
payload = binary.LittleEndian.AppendUint64(payload, uint64(len(key)))
|
||||
payload = append(payload, key...)
|
||||
payload = binary.LittleEndian.AppendUint32(payload, uint32(gguf.GGUFMetadataValueTypeArray))
|
||||
payload = binary.LittleEndian.AppendUint32(payload, uint32(gguf.GGUFMetadataValueTypeString))
|
||||
payload = binary.LittleEndian.AppendUint64(payload, 1)
|
||||
payload = binary.LittleEndian.AppendUint64(payload, math.MaxUint64)
|
||||
return payload
|
||||
}
|
||||
@@ -325,32 +325,10 @@ func AvailableGalleryModels(galleries []config.Gallery, systemState *system.Syst
|
||||
var (
|
||||
availableModelsMu sync.RWMutex
|
||||
availableModelsCache GalleryElements[*GalleryModel]
|
||||
// Whether a load has happened, tracked apart from the slice itself. A
|
||||
// gallery that legitimately holds nothing caches as an empty (often nil)
|
||||
// slice, and testing the slice for nil read that as "never loaded": every
|
||||
// call then took the blocking path and bumped the generation, which is the
|
||||
// same cache-defeating loop the refresh interval exists to stop.
|
||||
availableModelsLoaded bool
|
||||
refreshing atomic.Bool
|
||||
galleryGeneration atomic.Uint64
|
||||
lastRefreshUnixNano atomic.Int64
|
||||
refreshing atomic.Bool
|
||||
galleryGeneration atomic.Uint64
|
||||
)
|
||||
|
||||
// How often the cached model list may be refreshed from upstream.
|
||||
//
|
||||
// This is a floor on refresh frequency, not a TTL: the cache is served
|
||||
// regardless, and this only decides how often a background re-fetch is worth
|
||||
// starting. It matters far more than it looks, because a refresh bumps
|
||||
// galleryGeneration, and that invalidates every VRAM estimate cache in
|
||||
// pkg/vram. Refreshing on every call therefore kept those caches permanently
|
||||
// cold: the gallery listing is one request but the UI asks for one VRAM
|
||||
// estimate per row, so a single page view triggered dozens of refreshes and
|
||||
// every estimate paid full price for a remote probe it had already made.
|
||||
//
|
||||
// A package variable rather than a constant so tests can drive refreshes
|
||||
// without waiting.
|
||||
var GalleryRefreshInterval = 5 * time.Minute
|
||||
|
||||
// GalleryGeneration returns a counter that increments each time the gallery
|
||||
// model list is refreshed from upstream. VRAM estimation caches use this to
|
||||
// invalidate entries when the gallery data changes.
|
||||
@@ -374,11 +352,7 @@ func ResetGalleryModelCache() {
|
||||
}
|
||||
availableModelsMu.Lock()
|
||||
availableModelsCache = nil
|
||||
availableModelsLoaded = false
|
||||
availableModelsMu.Unlock()
|
||||
// Also clear the refresh stamp, or a suite that reset the cache would find
|
||||
// the next refresh throttled by the previous spec's clock.
|
||||
lastRefreshUnixNano.Store(0)
|
||||
}
|
||||
|
||||
// AvailableGalleryModelsCached returns gallery models from an in-memory cache.
|
||||
@@ -389,10 +363,9 @@ func ResetGalleryModelCache() {
|
||||
func AvailableGalleryModelsCached(galleries []config.Gallery, systemState *system.SystemState) (GalleryElements[*GalleryModel], error) {
|
||||
availableModelsMu.RLock()
|
||||
cached := availableModelsCache
|
||||
loaded := availableModelsLoaded
|
||||
availableModelsMu.RUnlock()
|
||||
|
||||
if loaded {
|
||||
if cached != nil {
|
||||
// Refresh installed status under write lock to avoid races with
|
||||
// concurrent readers and the background refresh goroutine.
|
||||
availableModelsMu.Lock()
|
||||
@@ -414,10 +387,8 @@ func AvailableGalleryModelsCached(galleries []config.Gallery, systemState *syste
|
||||
|
||||
availableModelsMu.Lock()
|
||||
availableModelsCache = models
|
||||
availableModelsLoaded = true
|
||||
galleryGeneration.Add(1)
|
||||
availableModelsMu.Unlock()
|
||||
lastRefreshUnixNano.Store(time.Now().UnixNano())
|
||||
|
||||
return models, nil
|
||||
}
|
||||
@@ -426,18 +397,9 @@ func AvailableGalleryModelsCached(galleries []config.Gallery, systemState *syste
|
||||
// gallery model cache. Only one refresh runs at a time; concurrent calls
|
||||
// are no-ops.
|
||||
func triggerGalleryRefresh(galleries []config.Gallery, systemState *system.SystemState) {
|
||||
if GalleryRefreshInterval > 0 {
|
||||
last := lastRefreshUnixNano.Load()
|
||||
if last != 0 && time.Since(time.Unix(0, last)) < GalleryRefreshInterval {
|
||||
return
|
||||
}
|
||||
}
|
||||
if !refreshing.CompareAndSwap(false, true) {
|
||||
return
|
||||
}
|
||||
// Stamped before the fetch rather than after, so a slow upstream cannot
|
||||
// let a queue of callers each start their own refresh behind this one.
|
||||
lastRefreshUnixNano.Store(time.Now().UnixNano())
|
||||
go func() {
|
||||
defer refreshing.Store(false)
|
||||
models, err := AvailableGalleryModels(galleries, systemState)
|
||||
@@ -446,37 +408,12 @@ func triggerGalleryRefresh(galleries []config.Gallery, systemState *system.Syste
|
||||
return
|
||||
}
|
||||
availableModelsMu.Lock()
|
||||
changed := !sameModelSet(availableModelsCache, models)
|
||||
availableModelsCache = models
|
||||
availableModelsLoaded = true
|
||||
// Only a real change invalidates the VRAM caches. An unchanged gallery
|
||||
// re-fetched on schedule must not throw away work that is still valid,
|
||||
// which is the difference between an estimate costing nothing and
|
||||
// costing a network round trip.
|
||||
if changed {
|
||||
galleryGeneration.Add(1)
|
||||
}
|
||||
galleryGeneration.Add(1)
|
||||
availableModelsMu.Unlock()
|
||||
}()
|
||||
}
|
||||
|
||||
// sameModelSet reports whether two model lists describe the same gallery, for
|
||||
// the purpose of deciding whether derived caches are still valid. Names and
|
||||
// order are enough: a change to an entry's files or size arrives with a new
|
||||
// gallery index, and comparing every field on every entry would cost more than
|
||||
// the caches save.
|
||||
func sameModelSet(a, b GalleryElements[*GalleryModel]) bool {
|
||||
if len(a) != len(b) {
|
||||
return false
|
||||
}
|
||||
for i := range a {
|
||||
if a[i].GetName() != b[i].GetName() {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
// List available backends
|
||||
func AvailableBackends(galleries []config.Gallery, systemState *system.SystemState) (GalleryElements[*GalleryBackend], error) {
|
||||
return availableBackendsWithFilter(galleries, systemState, func(backend *GalleryBackend) bool {
|
||||
|
||||
@@ -1,80 +0,0 @@
|
||||
package gallery_test
|
||||
|
||||
import (
|
||||
"os"
|
||||
"time"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
|
||||
"github.com/mudler/LocalAI/core/config"
|
||||
"github.com/mudler/LocalAI/core/gallery"
|
||||
"github.com/mudler/LocalAI/pkg/system"
|
||||
)
|
||||
|
||||
// The gallery generation counter is what every VRAM estimate cache keys on, so
|
||||
// how often it moves decides whether those caches are worth having. Refreshing
|
||||
// on every call kept them permanently cold: one page of the model gallery asks
|
||||
// for a VRAM estimate per row, and each of those requests re-read the gallery,
|
||||
// triggering a refresh that invalidated the estimate the previous row had just
|
||||
// paid a network round trip for.
|
||||
var _ = Describe("Gallery refresh throttling", func() {
|
||||
var (
|
||||
tmp *system.SystemState
|
||||
galleries []config.Gallery
|
||||
origInterval time.Duration
|
||||
)
|
||||
|
||||
BeforeEach(func() {
|
||||
dir, err := os.MkdirTemp("", "gallery-throttle")
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
DeferCleanup(func() { os.RemoveAll(dir) })
|
||||
|
||||
tmp, err = system.GetSystemState(system.WithModelPath(dir))
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
// No upstream: the list comes back empty, which is all this needs. What
|
||||
// is under test is how often a refresh is started, not what it returns.
|
||||
galleries = []config.Gallery{}
|
||||
origInterval = gallery.GalleryRefreshInterval
|
||||
gallery.ResetGalleryModelCache()
|
||||
})
|
||||
|
||||
AfterEach(func() {
|
||||
gallery.GalleryRefreshInterval = origInterval
|
||||
gallery.ResetGalleryModelCache()
|
||||
})
|
||||
|
||||
It("does not bump the generation once per call", func() {
|
||||
gallery.GalleryRefreshInterval = time.Hour
|
||||
|
||||
_, err := gallery.AvailableGalleryModelsCached(galleries, tmp)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
start := gallery.GalleryGeneration()
|
||||
|
||||
// Stands in for one page view: many callers in quick succession.
|
||||
for i := 0; i < 30; i++ {
|
||||
_, err := gallery.AvailableGalleryModelsCached(galleries, tmp)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
}
|
||||
// Let any refresh that did start finish, so this cannot pass by racing.
|
||||
Eventually(func() uint64 { return gallery.GalleryGeneration() }, "2s", "50ms").
|
||||
Should(Equal(start))
|
||||
})
|
||||
|
||||
It("still refreshes once the interval has passed", func() {
|
||||
gallery.GalleryRefreshInterval = time.Millisecond
|
||||
|
||||
_, err := gallery.AvailableGalleryModelsCached(galleries, tmp)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
time.Sleep(5 * time.Millisecond)
|
||||
_, err = gallery.AvailableGalleryModelsCached(galleries, tmp)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
// An empty gallery refreshing to an empty gallery is unchanged, so the
|
||||
// generation must hold: only a real change may invalidate the caches.
|
||||
Consistently(func() uint64 { return gallery.GalleryGeneration() }, "300ms", "50ms").
|
||||
Should(Equal(gallery.GalleryGeneration()))
|
||||
})
|
||||
})
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user