diff --git a/CMakeLists.txt b/CMakeLists.txt index d50e7b46..e3dbd9f4 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1314,6 +1314,20 @@ audiocpp_add_model(granite5asr granite_speech5_ctc ) +audiocpp_add_model(audio8_asr + SOURCES + src/community_models/audio8_asr/assets.cpp + src/community_models/audio8_asr/projector.cpp + src/community_models/audio8_asr/thinker.cpp + src/community_models/audio8_asr/session.cpp + INCLUDES + engine/community_models/audio8_asr/session.h + LOADERS + engine::community_models::audio8_asr::make_audio8_asr_loader + DEPENDS + qwen3_asr +) + audiocpp_add_model(vevo2 SOURCES src/models/vevo2/ar.cpp @@ -2325,6 +2339,34 @@ if (ENGINE_BUILD_TESTS) target_link_libraries(test_granite5asr_golden_transcription PRIVATE OpenMP::OpenMP_CXX) endif() + if (audio8_asr IN_LIST AUDIOCPP_LINKED_MODELS) + add_executable(test_audio8_asr_units + tests/audio8_asr/test_audio8_asr_units.cpp + ) + target_link_libraries(test_audio8_asr_units PRIVATE engine_runtime ggml) + target_include_directories(test_audio8_asr_units PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/tests/unittests) + if (ENGINE_ENABLE_OPENMP) + target_link_libraries(test_audio8_asr_units PRIVATE OpenMP::OpenMP_CXX) + endif() + + add_test( + NAME test_audio8_asr_units + COMMAND test_audio8_asr_units + ) + + add_executable(test_audio8_asr_golden_transcription + tests/audio8_asr/test_audio8_asr_golden_transcription.cpp + ) + target_compile_definitions(test_audio8_asr_golden_transcription PRIVATE + ENGINE_REPO_ROOT="${CMAKE_CURRENT_SOURCE_DIR}" + ) + target_link_libraries(test_audio8_asr_golden_transcription PRIVATE engine_runtime ggml) + target_include_directories(test_audio8_asr_golden_transcription PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}) + if (ENGINE_ENABLE_OPENMP) + target_link_libraries(test_audio8_asr_golden_transcription PRIVATE OpenMP::OpenMP_CXX) + endif() + endif() + add_executable(torch_bin_parity tests/vibevoice/torch_bin_parity.cpp ) diff --git a/README.md b/README.md index 50f57624..c24e0951 100644 --- a/README.md +++ b/README.md @@ -144,6 +144,7 @@ Community model ports live under `community_models` to make the ownership bounda | Family | Task | Lang | Runtime | Contributor | What They Added | |---|---|---|---|---|---| +| **audio8_asr** | ASR | en, zh, yue, ja, ko, fr, de | GGUF Q8, Safetensors | [@0xShug0](https://github.com/0xShug0) | [Audio8-ASR-0.1B](docs/community_models/audio8_asr.md) compact multilingual autoregressive ASR reusing the Qwen3-ASR encoder with an MLP-tower adapter and an 8-layer Qwen2-style decoder (CC-BY-NC, local conversion only) | | **f5_tts** | TTS, Clone | en, ar (Habibi) | GGUF | [@tareko](https://github.com/tareko) | [F5-TTS](docs/community_models/f5_tts.md) flow-matching DiT synthesis and voice cloning, with Habibi Arabic aliases `habibi`/`habibi_tts` | | **glm_tts** | TTS, Clone | zh, en | GGUF | Mirek [@mirek190](https://github.com/mirek190) | [GLM-TTS](docs/community_models/glm_tts.md) zero-shot synthesis and voice cloning support | | **granite5asr** | ASR | en | GGUF Q8 | [@ampersandru](https://github.com/ampersandru) | [IBM Granite Speech 5.0 470M TurboCTC](docs/community_models/granite5asr.md) ultra-fast Conformer-CTC ASR with Shaw relative positional embeddings and ByteLevel BPE | diff --git a/docs/community_models/audio8_asr.md b/docs/community_models/audio8_asr.md new file mode 100644 index 00000000..10ca2f6f --- /dev/null +++ b/docs/community_models/audio8_asr.md @@ -0,0 +1,104 @@ +# Audio8-ASR-0.1B in audio.cpp + +[Audio8-ASR-0.1B](https://huggingface.co/Audio8/Audio8-ASR-0.1B) is a compact +multilingual autoregressive ASR model (en / zh / yue / ja / ko / fr / de): a +Qwen3-ASR audio encoder adapted by an MLP tower into an 8-layer Qwen2-style +decoder with only ~103M language-model parameters (324M end-to-end). The +checkpoint ships under **CC-BY-NC-4.0**, so audio.cpp loads it from locally +converted weights only; the converted GGUF must not be redistributed. + +## Architecture + +- **Audio frontend**: 16 kHz Whisper log-mel, 128 bins, hop 160, n_fft 400 + (shared with the `qwen3_asr` family). The reference processor emits + bfloat16 mel values; the audio8_asr frontend rounds to bfloat16 before + encoding to match. +- **Audio encoder**: Qwen3-ASR audio tower (d_model 896, 18 layers, 14 heads, + FFN 3584, 128 mel bins, output dim 1024) — bit-for-bit the same + architecture as `Qwen/Qwen3-ASR-0.6B`, loaded through the shared + `qwen3_asr` encoder implementation. +- **Adapter**: 4 pre-norm residual MLP blocks (1024 → 4096, erf GELU), a + final LayerNorm, then an adaptive average pool (merge factor 4 against the + mel-frame count) followed by LayerNorm + Linear(1024 → 512). +- **Decoder**: Qwen2-style causal LM, 8 layers, hidden 512, 8 heads, + head_dim 64, SwiGLU FFN 1408, tied embeddings over a 151,936 Qwen BPE + vocabulary, RoPE theta 1e6, RMSNorm eps 1e-6. +- **Prompt**: `<|user|><|begin_of_audio|><|audio|>×N<|end_of_audio|>Please + transcribe this audio.<|assistant|>` where + `N = max(floor(floor((floor(samples / hop) + 1) / 2) / merge_factor), 1)` + (all divisions integer). + +## Usage + +```bash +# Convert locally (requires the audiocpp_gguf tool and a self-downloaded +# checkpoint from the Audio8/Audio8-ASR-0.1B HF repository): +python tools/community_models/convert_audio8_asr.py \ + --checkpoint models/Audio8-ASR-0.1B-hf \ + --converter build/debug/bin/audiocpp_gguf \ + --type q8_0 \ + --output models/Audio8-ASR-0.1B-GGUF/audio8-asr-0.1b-q8_0.gguf + +# Transcribe +audiocpp_cli --task asr --family audio8_asr \ + --model models/Audio8-ASR-0.1B-GGUF/audio8-asr-0.1b-q8_0.gguf \ + --audio sample.wav + +# The safetensors package loads directly, no conversion required: +audiocpp_cli --task asr --family audio8_asr \ + --model models/Audio8-ASR-0.1B-hf --audio sample.wav +``` + +## Parity + +`tools/community_models/audio8_asr_reference.py` runs the Hugging Face +`trust_remote_code` reference (torch CPU, fp32) and +`tools/community_models/audio8_asr_stages.py` captures staged tensors (mel, +encoder output, projected audio embeddings) for comparison. + +Greedy transcription on the repo test clips matched the fp32 reference +exactly with the Q8_0 GGUF on both the Metal and CPU backends: + +| Audio | Reference (fp32) | audio.cpp (Q8_0 GGUF, Metal + CPU) | +|---|---|---| +| `assets/resources/a.wav` (5.95 s) | "This little work was finished in the year eighteen o three, and intended for immediate publication." | identical | +| `assets/resources/sample_16k.wav` (14.07 s) | "Some call me nature. Others call me Mother Nature. I've been here for over four point five billion years, twenty-two thousand five hundred times longer than you." | identical | + +A 61-second clip transcribed through rate-correct 30-second windows matched +the per-window reference transcripts (also verified by an independent review +pass). `test_audio8_asr_golden_transcription` asserts the first row +end-to-end — run it manually once weights exist (it is not part of ctest, +matching the granite5asr golden-test convention); `test_audio8_asr_units` +covers the token-count formula and bfloat16 rounding and runs under ctest. + +## Measured performance + +Release build (`-DCMAKE_BUILD_TYPE=Release`), Apple M4, Metal backend, +Q8_0 GGUF (345 MB weights): + +| Audio | Session wall (`session.wall_ms`) | Effective RTF | CLI wall (incl. ~2.2 s process + load) | +|---|---|---|---| +| 5.95 s | ~0.75 s | ~8x realtime | 2.2 s | +| 14.07 s | 757 ms | 18.6x realtime | 3.8 s | +| 61 s (3 windows) | 3.46 s | 17.6x realtime | 5.8 s | + +The audio encoder dominates (~64% of session time; its Metal shaders are +insensitive to build type, so Debug and Release measure within 3%). +Peak RSS is ~1.0 GB for a single clip and ~1.3 GB for a three-window +transcription (Metal buffers + per-window-shape graph pools on top of the +345 MB of weights; CPU backend measures the same ~1.0 GB). + +## Known limitations + +- **Offline only**: no streaming mode, no word timestamps, no language-id + output. +- **30-second windows**: audio longer than 30 s is transcribed in fixed + windows sized at the input sample rate (0.5 s minimum tail folds into the + previous window) and space-joined, without VAD segmentation. Unlike the + single-pass reference, each window is peak-normalized independently, so + relative loudness across a window boundary can shift. +- **CC-BY-NC-4.0**: non-commercial use only; convert locally, do not + redistribute the converted GGUF. There is no release GGUF package and no + WebUI catalog entry (the package manager cannot download + unsupported-license packages). +- Hotword logit boosting from the reference implementation is not ported. diff --git a/docs/community_models/models.md b/docs/community_models/models.md index 2ebc8340..fc6bc07c 100644 --- a/docs/community_models/models.md +++ b/docs/community_models/models.md @@ -16,6 +16,7 @@ Practical expectations: | Family | Task | Supported language(s) | Contributor | What They Added | |---|---|---|---|---| +| **audio8_asr** | ASR | en, zh, yue, ja, ko, fr, de | [@0xShug0](https://github.com/0xShug0) | [Audio8-ASR-0.1B](audio8_asr.md) compact multilingual autoregressive ASR reusing the Qwen3-ASR encoder with an MLP-tower adapter and an 8-layer Qwen2-style decoder (CC-BY-NC, local conversion only) | | **echo_tts** | TTS, voice cloning | en | Tym [@5uck1ess](https://github.com/5uck1ess), [@dignome](https://github.com/dignome) | [Echo-TTS](echo_tts.md) 44.1 kHz zero-shot voice cloning: 2.8B diffusion transformer in 80-D PCA space, decoded by the Fish S1-DAC autoencoder. Byte-level text, no phonemiser, no reference transcript | | **f5_tts** | TTS, voice cloning | en, ar (Habibi) | Community | [F5-TTS](f5_tts.md) flow-matching DiT — M0 scaffolding, aliases `habibi`/`habibi_tts` | | **glm_tts** | TTS, voice cloning | zh, en | Mirek [@mirek190](https://github.com/mirek190) | [GLM-TTS](glm_tts.md) zero-shot synthesis and voice cloning support | diff --git a/include/engine/community_models/audio8_asr/assets.h b/include/engine/community_models/audio8_asr/assets.h new file mode 100644 index 00000000..c4f8aec7 --- /dev/null +++ b/include/engine/community_models/audio8_asr/assets.h @@ -0,0 +1,31 @@ +#pragma once + +#include "engine/framework/assets/resource_bundle.h" +#include "engine/framework/tokenizers/llama_bpe.h" +#include "engine/models/qwen3_asr/assets.h" + +#include +#include + +#include "engine/community_models/audio8_asr/types.h" + +namespace engine::community_models::audio8_asr { + +struct Audio8ASRAssets { + // Bundle and config for this family (arkasr layout). + assets::ResourceBundle resources; + Audio8ASRConfig config; + + // Qwen2 BPE tokenizer shared with the Qwen3-ASR family tooling. + std::shared_ptr tokenizer; + + // A Qwen3-ASR view over the same bundle so the audio encoder and Whisper + // frontend implementations can be reused unchanged. The weights source in + // this view renames Audio8 tensors to the prefixes those implementations + // expect (`audio_encoder.*` -> `model.audio_tower.*` etc.). + std::shared_ptr encoder_assets; +}; + +std::shared_ptr load_audio8_asr_assets(const std::filesystem::path & model_path); + +} // namespace engine::community_models::audio8_asr diff --git a/include/engine/community_models/audio8_asr/projector.h b/include/engine/community_models/audio8_asr/projector.h new file mode 100644 index 00000000..460dd3dc --- /dev/null +++ b/include/engine/community_models/audio8_asr/projector.h @@ -0,0 +1,41 @@ +#pragma once + +#include "engine/framework/assets/tensor_source.h" +#include "engine/framework/core/execution_context.h" +#include "engine/community_models/audio8_asr/types.h" + +#include +#include + +namespace engine::community_models::audio8_asr { + +// Runs the Audio8 adapter tail over Qwen3-ASR encoder output: the residual +// MLP tower, an adaptive average pool down to the prompt's audio token count, +// and the LayerNorm + linear projector into the decoder hidden size. +class Audio8ProjectorRuntime { +public: + Audio8ProjectorRuntime( + std::shared_ptr weights_source, + const Audio8TowerConfig & config, + core::ExecutionContext & execution, + size_t graph_arena_bytes, + size_t weight_context_bytes, + assets::TensorStorageType weight_storage_type); + ~Audio8ProjectorRuntime(); + + Audio8ProjectorRuntime(const Audio8ProjectorRuntime &) = delete; + Audio8ProjectorRuntime & operator=(const Audio8ProjectorRuntime &) = delete; + + // input: [encoder_tokens, tower.input_size] float values (token-major); + // returns [audio_tokens, tower.output_size] float values (token-major). + Audio8ASRAudioEmbeddings project( + const std::vector & encoder_output, + int64_t encoder_tokens, + int64_t audio_tokens); + +private: + struct Impl; + std::unique_ptr impl_; +}; + +} // namespace engine::community_models::audio8_asr diff --git a/include/engine/community_models/audio8_asr/session.h b/include/engine/community_models/audio8_asr/session.h new file mode 100644 index 00000000..fc297498 --- /dev/null +++ b/include/engine/community_models/audio8_asr/session.h @@ -0,0 +1,66 @@ +#pragma once + +#include "engine/framework/runtime/model.h" +#include "engine/framework/runtime/session_base.h" +#include "engine/community_models/audio8_asr/assets.h" +#include "engine/community_models/audio8_asr/projector.h" +#include "engine/community_models/audio8_asr/thinker.h" +#include "engine/models/qwen3_asr/audio_encoder.h" +#include "engine/models/qwen3_asr/frontend_whisper.h" + +#include +#include + +namespace engine::community_models::audio8_asr { + +class Audio8ASRSession final + : public runtime::RuntimeSessionBase + , public runtime::IOfflineVoiceTaskSession { +public: + Audio8ASRSession( + runtime::TaskSpec task, + runtime::SessionOptions options, + std::shared_ptr assets); + ~Audio8ASRSession() override; + + std::string family() const override; + runtime::VoiceTaskKind task_kind() const override; + runtime::RunMode run_mode() const override; + void prepare(const runtime::SessionPreparationRequest & request) override; + runtime::TaskResult run(const runtime::TaskRequest & request) override; + +private: + std::string transcribe_clip(const runtime::AudioBuffer & audio); + runtime::Transcript transcribe_audio(const runtime::AudioBuffer & audio); + + runtime::TaskSpec task_; + std::shared_ptr assets_; + qwen3_asr::Qwen3ASRWhisperFrontend frontend_; + qwen3_asr::Qwen3ASRAudioEncoderRuntime audio_encoder_; + Audio8ProjectorRuntime projector_; + Audio8ThinkerRuntime thinker_; +}; + +class Audio8ASRLoadedModel final : public runtime::ILoadedVoiceModel { +public: + Audio8ASRLoadedModel( + runtime::ModelMetadata metadata, + runtime::CapabilitySet capabilities, + std::shared_ptr assets); + + const runtime::ModelMetadata & metadata() const noexcept override; + const runtime::CapabilitySet & capabilities() const noexcept override; + std::unique_ptr create_task_session( + const runtime::TaskSpec & task, + const runtime::SessionOptions & options) const override; + +private: + runtime::ModelMetadata metadata_; + runtime::CapabilitySet capabilities_; + std::shared_ptr assets_; +}; + +std::unique_ptr load_audio8_asr_model(const runtime::ModelLoadRequest & request); +std::shared_ptr make_audio8_asr_loader(); + +} // namespace engine::community_models::audio8_asr diff --git a/include/engine/community_models/audio8_asr/thinker.h b/include/engine/community_models/audio8_asr/thinker.h new file mode 100644 index 00000000..59c4e033 --- /dev/null +++ b/include/engine/community_models/audio8_asr/thinker.h @@ -0,0 +1,40 @@ +#pragma once + +#include "engine/framework/assets/tensor_source.h" +#include "engine/framework/core/execution_context.h" +#include "engine/community_models/audio8_asr/types.h" + +#include +#include + +namespace engine::community_models::audio8_asr { + +// Greedy causal decoder for the Audio8 8-layer Qwen2-style LM. Audio +// embeddings are injected into the token embedding sequence at the prompt's +// audio placeholder positions before prefill. +class Audio8ThinkerRuntime { +public: + Audio8ThinkerRuntime( + std::shared_ptr weights_source, + const Audio8ASRDecoderConfig & config, + core::ExecutionContext & execution, + size_t prefill_graph_arena_bytes, + size_t decode_graph_arena_bytes, + size_t weight_context_bytes, + assets::TensorStorageType weight_storage_type); + ~Audio8ThinkerRuntime(); + + Audio8ThinkerRuntime(const Audio8ThinkerRuntime &) = delete; + Audio8ThinkerRuntime & operator=(const Audio8ThinkerRuntime &) = delete; + + Audio8ASRGeneratedTokens generate( + const Audio8ASRPrompt & prompt, + const Audio8ASRAudioEmbeddings & audio_embeddings, + const Audio8ASRGenerationOptions & options); + +private: + struct Impl; + std::unique_ptr impl_; +}; + +} // namespace engine::community_models::audio8_asr diff --git a/include/engine/community_models/audio8_asr/types.h b/include/engine/community_models/audio8_asr/types.h new file mode 100644 index 00000000..373dae7f --- /dev/null +++ b/include/engine/community_models/audio8_asr/types.h @@ -0,0 +1,110 @@ +#pragma once + +#include "engine/models/qwen3_asr/assets.h" + +#include +#include +#include +#include +#include + +namespace engine::community_models::audio8_asr { + +namespace qwen3_asr = engine::models::qwen3_asr; + +struct Audio8ASRDecoderConfig { + int64_t vocab_size = 0; + int64_t hidden_size = 0; + int64_t intermediate_size = 0; + int64_t num_hidden_layers = 0; + int64_t num_attention_heads = 0; + int64_t num_key_value_heads = 0; + int64_t head_dim = 0; + int64_t max_position_embeddings = 32768; + int64_t audio_token_id = 0; + int64_t pad_token_id = 0; + int64_t max_new_tokens = 256; + std::vector eos_token_ids; + bool tie_word_embeddings = true; + float rms_norm_eps = 1.0e-6F; + float rope_theta = 1000000.0F; +}; + +// The MLP tower + projector that adapts the Qwen3-ASR encoder output to the +// Audio8 decoder hidden size: four pre-norm residual MLP blocks, a final norm, +// an adaptive average pool down to the prompt's audio token count, and a +// LayerNorm + linear projector into the decoder hidden size. +struct Audio8TowerConfig { + int64_t input_size = 0; + int64_t hidden_size = 0; + int64_t intermediate_size = 0; + int64_t layers = 4; + int64_t output_size = 0; + float norm_eps = 1.0e-5F; +}; + +struct Audio8ASRConfig { + int64_t merge_factor = 4; + int64_t max_audio_samples = 480000; + int64_t user_token_id = 0; + int64_t begin_audio_token_id = 0; + int64_t end_audio_token_id = 0; + int64_t assistant_token_id = 0; + qwen3_asr::Qwen3ASRFrontendConfig frontend; + qwen3_asr::Qwen3ASRAudioEncoderConfig audio_encoder; + Audio8TowerConfig tower; + Audio8ASRDecoderConfig text_decoder; + std::vector supported_languages; +}; + +struct Audio8ASRPrompt { + std::vector input_ids; + std::vector audio_token_positions; +}; + +struct Audio8ASRAudioEmbeddings { + std::vector values; + int64_t tokens = 0; + int64_t hidden_size = 0; +}; + +struct Audio8ASRGenerationOptions { + int64_t max_new_tokens = 256; +}; + +struct Audio8ASRGeneratedTokens { + std::vector token_ids; +}; + +// The reference processor hands the model a bfloat16 mel tensor; fp32 +// pipelines replicate that rounding before encoding (round to nearest even, +// like torch's .to(torch.bfloat16)). +inline float audio8_asr_round_f32_to_bf16(float value) { + uint32_t bits = 0; + std::memcpy(&bits, &value, sizeof(bits)); + const uint32_t lsb = (bits >> 16) & 1u; + bits += 0x7FFFu + lsb; + bits &= 0xFFFF0000u; + std::memcpy(&value, &bits, sizeof(value)); + return value; +} + +// Audio token count used by the Arkasr processor: +// mel_frames = samples // hop_length +// downsampled = (mel_frames + 1) // 2 +// tokens = max(downsampled // merge_factor, 1) +inline int64_t audio8_asr_prompt_audio_token_count( + int64_t mel_frames, + int64_t merge_factor) { + if (mel_frames <= 0) { + throw std::runtime_error("Audio8 ASR requires positive mel frame count"); + } + if (merge_factor <= 0) { + throw std::runtime_error("Audio8 ASR merge_factor must be positive"); + } + const int64_t downsampled = (mel_frames + 1) / 2; + const int64_t merged = downsampled / merge_factor; + return merged > 0 ? merged : 1; +} + +} // namespace engine::community_models::audio8_asr diff --git a/model_specs/audio8_asr.json b/model_specs/audio8_asr.json new file mode 100644 index 00000000..8628709d --- /dev/null +++ b/model_specs/audio8_asr.json @@ -0,0 +1,201 @@ +{ + "family": "audio8_asr", + "display_name": "Audio8-ASR-0.1B", + "description": "Audio8-ASR-0.1B compact multilingual ASR (en/zh/yue/ja/ko/fr/de): Qwen3-ASR audio encoder with an MLP-tower adapter and an 8-layer Qwen2-style decoder. 30-second offline transcription windows. Checkpoint license: CC-BY-NC-4.0; GGUF conversion is local-only, no public redistribution is approved.", + "category": "asr", + "status": "community", + "tasks": [ + "asr" + ], + "modes": [ + "offline" + ], + "languages": [ + "en", + "zh", + "yue", + "ja", + "ko", + "fr", + "de" + ], + "capabilities": { + "asr": [] + }, + "options": { + "request": [], + "session": [ + { + "name": "weight_type", + "type": "enum", + "description": "Decoder and adapter weight storage type.", + "preset": "weight_type_full", + "required": false, + "default": "native" + }, + { + "name": "audio_encoder_weight_type", + "type": "enum", + "description": "Audio encoder weight storage type (quantized storage is rejected for the encoder; already-quantized GGUF tensors load as native).", + "values": [ + "native", + "f32", + "f16" + ], + "required": false, + "default": "native" + }, + { + "name": "encoder_graph_arena_mb", + "type": "int", + "description": "Audio encoder graph arena size in MB.", + "required": false, + "min": 64, + "default": 128 + }, + { + "name": "projector_graph_arena_mb", + "type": "int", + "description": "Adapter (MLP tower + pooling + projector) graph arena size in MB.", + "required": false, + "min": 64, + "default": 256 + }, + { + "name": "prefill_graph_arena_mb", + "type": "int", + "description": "Decoder prefill graph arena size in MB.", + "required": false, + "min": 64, + "default": 256 + }, + { + "name": "decode_graph_arena_mb", + "type": "int", + "description": "Decoder step decode graph arena size in MB.", + "required": false, + "min": 64, + "default": 256 + } + ], + "load": [] + }, + "runtime": { + "tags": [ + "gguf", + "server", + "cuda", + "metal", + "cpu" + ] + }, + "ui": { + "recommended_package": "audio8_asr_0_1b_q8_0", + "tags": [ + "ASR", + "GGUF" + ], + "docs": [ + "docs/asr.md", + "docs/community_models/audio8_asr.md" + ], + "summary": "Audio8-ASR-0.1B compact multilingual offline transcription." + }, + "package_defaults": { + "download": { + "kind": "unsupported", + "reason": "CC-BY-NC-4.0 checkpoint: convert locally with audiocpp_gguf; no public audio.cpp GGUF distribution is approved." + } + }, + "packages": [ + { + "id": "audio8_asr_0_1b_q8_0", + "display_name": "Audio8-ASR-0.1B Q8_0 GGUF", + "default": true, + "format": "gguf", + "precision": "q8_0", + "target_directory": "Audio8-ASR-0.1B-GGUF", + "files": [ + "Audio8-ASR-0.1B-GGUF/audio8-asr-0.1b-q8_0.gguf" + ], + "strip_prefix": "Audio8-ASR-0.1B-GGUF" + }, + { + "id": "audio8_asr_0_1b_f16", + "display_name": "Audio8-ASR-0.1B F16 GGUF", + "format": "gguf", + "precision": "f16", + "target_directory": "Audio8-ASR-0.1B-GGUF", + "files": [ + "Audio8-ASR-0.1B-GGUF/audio8-asr-0.1b-f16.gguf" + ], + "strip_prefix": "Audio8-ASR-0.1B-GGUF" + }, + { + "id": "audio8_asr_0_1b_safetensors", + "display_name": "Audio8-ASR-0.1B HF Safetensors", + "format": "safetensors", + "precision": "bfloat16", + "target_directory": "Audio8-ASR-0.1B-hf", + "files": [ + "config.json", + "generation_config.json", + "preprocessor_config.json", + "tokenizer.json", + "tokenizer_config.json", + "vocab.json", + "merges.txt", + "model.safetensors" + ], + "download": { + "kind": "huggingface_snapshot", + "repo": "Audio8/Audio8-ASR-0.1B", + "revision": "main", + "gated": false + } + } + ], + "sources": [ + { + "format": "gguf", + "roots": { + "model": ".", + "weights": "$gguf" + }, + "files": { + "config": "model:config.json", + "generation_config": "model:generation_config.json", + "tokenizer_config": "model:tokenizer_config.json", + "preprocessor_config": "model:preprocessor_config.json" + }, + "optional_files": { + "vocab": "model:vocab.json", + "merges": "model:merges.txt", + "tokenizer_json": "model:tokenizer.json" + }, + "tensors": { + "weights": "weights:" + } + }, + { + "format": "safetensors", + "roots": { + "model": "." + }, + "files": { + "config": "model:config.json", + "generation_config": "model:generation_config.json", + "tokenizer_config": "model:tokenizer_config.json", + "preprocessor_config": "model:preprocessor_config.json" + }, + "optional_files": { + "vocab": "model:vocab.json", + "merges": "model:merges.txt", + "tokenizer_json": "model:tokenizer.json" + }, + "tensors": { + "weights": "model:model.safetensors" + } + } + ] +} diff --git a/src/community_models/audio8_asr/assets.cpp b/src/community_models/audio8_asr/assets.cpp new file mode 100644 index 00000000..4859bfaf --- /dev/null +++ b/src/community_models/audio8_asr/assets.cpp @@ -0,0 +1,286 @@ +#include "engine/community_models/audio8_asr/assets.h" + +#include "engine/framework/assets/tensor_source.h" +#include "engine/framework/io/json.h" +#include "engine/framework/model_spec/package.h" + +#include +#include +#include + +namespace engine::community_models::audio8_asr { +namespace json = engine::io::json; +namespace { + +// The Qwen3-ASR encoder/frontend implementations reused by this family address +// tensors with the official Qwen3-ASR prefixes (`model.audio_tower.*`, +// `model.multi_modal_projector.*`). Audio8 checkpoints keep the same encoder +// under `audio_encoder.*`; this source rewrites names in both directions so +// the shared implementations load Audio8 weights unchanged. +class RenamingTensorSource final : public assets::TensorSource { +public: + RenamingTensorSource( + std::shared_ptr inner, + std::vector> mappings) + : inner_(std::move(inner)), + mappings_(std::move(mappings)) { + // Longest `to` prefix first so the exact-inverse projection entries + // win over the catch-all tower mapping during reverse lookups. + reverse_ordered_ = mappings_; + std::stable_sort( + reverse_ordered_.begin(), + reverse_ordered_.end(), + [](const auto & left, const auto & right) { + return left.second.size() > right.second.size(); + }); + } + + const std::filesystem::path & source_path() const noexcept override { + return inner_->source_path(); + } + + bool has_tensor(std::string_view name) const noexcept override { + return inner_->has_tensor(rewrite(name)); + } + + assets::TensorMetadata require_metadata(std::string_view name) const override { + return inner_->require_metadata(rewrite(name)); + } + + std::vector tensors() const override { + auto tensors = inner_->tensors(); + for (auto & tensor : tensors) { + tensor.name = rewrite_back(tensor.name); + } + return tensors; + } + + void release_storage() const override { + inner_->release_storage(); + } + + assets::RawTensorData require_tensor_data(std::string_view name) const override { + return inner_->require_tensor_data(rewrite(name)); + } + + std::vector require_f32( + std::string_view name, + const std::optional> & expected_shape) const override { + return inner_->require_f32(rewrite(name), expected_shape); + } + + std::optional> optional_f32( + std::string_view name, + const std::optional> & expected_shape) const override { + return inner_->optional_f32(rewrite(name), expected_shape); + } + + int64_t require_i64_scalar(std::string_view name) const override { + return inner_->require_i64_scalar(rewrite(name)); + } + +private: + std::string rewrite(std::string_view name) const { + for (const auto & [from, to] : mappings_) { + if (name.rfind(from, 0) == 0) { + return to + std::string(name.substr(from.size())); + } + } + return std::string(name); + } + + std::string rewrite_back(const std::string & name) const { + // Reverse mapping is only used for diagnostics listings. Check the + // longest `to` prefixes first so proj1/proj2 map back to the + // multi-modal projector names instead of the tower prefix. + for (const auto & [from, to] : reverse_ordered_) { + if (name.rfind(to, 0) == 0) { + return from + name.substr(to.size()); + } + } + return name; + } + + std::shared_ptr inner_; + std::vector> mappings_; + std::vector> reverse_ordered_; +}; + +qwen3_asr::Qwen3ASRAudioEncoderConfig parse_audio_encoder_config(const json::Value & value) { + qwen3_asr::Qwen3ASRAudioEncoderConfig config; + config.num_mel_bins = json::require_i64(value, "num_mel_bins"); + config.encoder_layers = json::require_i64(value, "encoder_layers"); + config.encoder_attention_heads = json::require_i64(value, "encoder_attention_heads"); + config.encoder_ffn_dim = json::require_i64(value, "encoder_ffn_dim"); + config.d_model = json::require_i64(value, "d_model"); + config.max_source_positions = json::optional_i64(value, "max_source_positions", 1500); + config.n_window = json::require_i64(value, "n_window"); + config.n_window_infer = json::require_i64(value, "n_window_infer"); + config.conv_chunksize = json::require_i64(value, "conv_chunksize"); + config.downsample_hidden_size = json::require_i64(value, "downsample_hidden_size"); + config.output_dim = json::require_i64(value, "output_dim"); + config.activation_function = json::require_string(value, "activation_function"); + if (config.activation_function != "gelu") { + throw std::runtime_error("Audio8 ASR currently supports gelu audio activation"); + } + return config; +} + +int64_t require_added_token_id(const assets::ResourceBundle & resources, std::string_view content) { + const auto tokenizer = resources.parse_json("tokenizer_json"); + for (const auto & item : tokenizer.require("added_tokens").as_array()) { + const auto * token_content = item.find("content"); + const auto * token_id = item.find("id"); + if (token_content != nullptr && token_content->is_string() && + token_id != nullptr && token_id->is_number() && token_content->as_string() == content) { + return token_id->as_i64(); + } + } + throw std::runtime_error("Audio8 ASR tokenizer.json is missing token: " + std::string(content)); +} + +Audio8ASRConfig parse_config(const assets::ResourceBundle & resources) { + const auto root = resources.parse_json("config"); + + Audio8ASRConfig config; + const auto model_type = json::require_string(root, "model_type"); + if (model_type != "arkasr") { + throw std::runtime_error( + "Audio8 ASR requires model_type 'arkasr', got: " + model_type); + } + config.merge_factor = json::optional_i64(root, "merge_factor", 4); + if (config.merge_factor <= 0) { + throw std::runtime_error("Audio8 ASR merge_factor must be positive"); + } + + const auto & audio_config = root.require("qwen3_asr_audio_config"); + config.audio_encoder = parse_audio_encoder_config(audio_config); + + config.tower.input_size = config.audio_encoder.output_dim; + config.tower.hidden_size = config.audio_encoder.output_dim; + const auto tower_intermediate = json::optional_i64(root, "qwen3_asr_mlp_tower_hidden_size", 0); + config.tower.intermediate_size = + tower_intermediate > 0 ? tower_intermediate : config.tower.input_size * 4; + config.tower.layers = json::optional_i64(root, "qwen3_asr_mlp_tower_layers", 4); + config.tower.output_size = json::require_i64(root, "hidden_size"); + + auto & decoder = config.text_decoder; + decoder.vocab_size = json::require_i64(root, "vocab_size"); + decoder.hidden_size = config.tower.output_size; + decoder.intermediate_size = json::require_i64(root, "intermediate_size"); + decoder.num_hidden_layers = json::require_i64(root, "num_hidden_layers"); + decoder.num_attention_heads = json::require_i64(root, "num_attention_heads"); + decoder.num_key_value_heads = json::require_i64(root, "num_key_value_heads"); + decoder.head_dim = json::optional_i64( + root, "head_dim", decoder.hidden_size / decoder.num_attention_heads); + decoder.max_position_embeddings = json::optional_i64(root, "max_position_embeddings", 32768); + decoder.audio_token_id = json::require_i64(root, "audio_token_id"); + decoder.pad_token_id = json::require_i64(root, "pad_token_id"); + decoder.tie_word_embeddings = json::optional_bool(root, "tie_word_embeddings", true); + decoder.rms_norm_eps = json::optional_f32(root, "rms_norm_eps", decoder.rms_norm_eps); + decoder.rope_theta = json::optional_f32(root, "rope_theta", decoder.rope_theta); + + const auto generation = resources.parse_json("generation_config"); + config.max_audio_samples = json::optional_i64(generation, "max_audio_samples", config.max_audio_samples); + decoder.max_new_tokens = json::optional_i64(generation, "max_new_tokens", decoder.max_new_tokens); + decoder.pad_token_id = json::optional_i64(generation, "pad_token_id", decoder.pad_token_id); + decoder.eos_token_ids = json::require_i64_array_or_scalar(generation, "eos_token_id"); + + const auto processor = resources.parse_json("preprocessor_config"); + const auto * feature_extractor = processor.find("feature_extractor"); + const auto & frontend = feature_extractor != nullptr && feature_extractor->is_object() ? *feature_extractor : processor; + config.frontend.sample_rate = static_cast(json::optional_i64(frontend, "sampling_rate", config.frontend.sample_rate)); + config.frontend.feature_size = json::require_i64(frontend, "feature_size"); + config.frontend.hop_length = json::require_i64(frontend, "hop_length"); + config.frontend.n_fft = json::require_i64(frontend, "n_fft"); + if (config.frontend.feature_size != config.audio_encoder.num_mel_bins) { + throw std::runtime_error("Audio8 ASR frontend feature size does not match audio encoder config"); + } + + // Prompt special tokens live in tokenizer.json added_tokens; the audio + // token id from the config must match the tokenizer entry. + config.user_token_id = require_added_token_id(resources, "<|user|>"); + config.begin_audio_token_id = require_added_token_id(resources, "<|begin_of_audio|>"); + config.end_audio_token_id = require_added_token_id(resources, "<|end_of_audio|>"); + config.assistant_token_id = require_added_token_id(resources, "<|assistant|>"); + config.text_decoder.audio_token_id = require_added_token_id(resources, "<|audio|>"); + + config.supported_languages = { + "Chinese", "English", "Cantonese", "French", "German", "Japanese", "Korean"}; + return config; +} + +void validate_config(const Audio8ASRConfig & config) { + const auto & decoder = config.text_decoder; + if (decoder.vocab_size <= 0 || decoder.hidden_size <= 0 || decoder.intermediate_size <= 0 || + decoder.num_hidden_layers <= 0 || decoder.num_attention_heads <= 0 || + decoder.num_key_value_heads <= 0 || decoder.head_dim <= 0 || + decoder.hidden_size % decoder.num_attention_heads != 0) { + throw std::runtime_error("Audio8 ASR invalid decoder metadata"); + } + if (config.audio_encoder.output_dim <= 0 || config.tower.layers <= 0 || + config.tower.intermediate_size <= 0) { + throw std::runtime_error("Audio8 ASR invalid audio adapter metadata"); + } + if (config.frontend.sample_rate != 16000) { + throw std::runtime_error("Audio8 ASR requires 16 kHz audio frontend"); + } +} + +} // namespace + +std::shared_ptr load_audio8_asr_assets(const std::filesystem::path & model_path) { + auto resources = engine::model_spec::load_resource_bundle_for_family(model_path, "audio8_asr"); + if (!resources.has_file("preprocessor_config")) { + throw std::runtime_error("Audio8 ASR requires preprocessor_config.json"); + } + if (!resources.has_file("tokenizer_json")) { + throw std::runtime_error("Audio8 ASR requires tokenizer.json"); + } + + auto assets = std::make_shared(); + assets->resources = std::move(resources); + assets->config = parse_config(assets->resources); + validate_config(assets->config); + + const auto weights = assets->resources.open_tensor_source("weights"); + // Encoder view for the shared Qwen3-ASR implementations: expose the + // Audio8 `audio_encoder.*` namespace under the official Qwen3-ASR + // prefixes. The Audio8 encoder carries its input/output projections + // (proj1/proj2) inside its own namespace; the shared code addresses them + // via the multi-modal projector prefix with linear_1/linear_2 names. + auto encoder_source = std::make_shared( + weights, + std::vector>{ + {"model.audio_tower.", "audio_encoder."}, + {"model.multi_modal_projector.linear_1.", "audio_encoder.proj1."}, + {"model.multi_modal_projector.linear_2.", "audio_encoder.proj2."}, + {"model.multi_modal_projector.", "audio_encoder."}}); + + auto encoder_assets = std::make_shared(); + encoder_assets->config.hf_transformers_layout = true; + encoder_assets->config.frontend.sample_rate = assets->config.frontend.sample_rate; + encoder_assets->config.frontend.feature_size = assets->config.frontend.feature_size; + encoder_assets->config.frontend.hop_length = assets->config.frontend.hop_length; + encoder_assets->config.frontend.n_fft = assets->config.frontend.n_fft; + encoder_assets->config.audio_encoder = assets->config.audio_encoder; + encoder_assets->model_weights = std::move(encoder_source); + assets->encoder_assets = std::move(encoder_assets); + + engine::tokenizers::LlamaBpeTokenizerSpec tokenizer_spec; + tokenizer_spec.tokenizer_config_path = assets->resources.require_file("tokenizer_config"); + if (const auto * path = assets->resources.find_file("vocab")) { + tokenizer_spec.vocab_path = *path; + } + if (const auto * path = assets->resources.find_file("merges")) { + tokenizer_spec.merges_path = *path; + } + if (const auto * path = assets->resources.find_file("tokenizer_json")) { + tokenizer_spec.tokenizer_json_path = *path; + } + tokenizer_spec.pre_type = engine::tokenizers::LlamaBpePreTokenizer::Qwen2; + assets->tokenizer = engine::tokenizers::load_llama_bpe_tokenizer(tokenizer_spec); + return assets; +} + +} // namespace engine::community_models::audio8_asr diff --git a/src/community_models/audio8_asr/projector.cpp b/src/community_models/audio8_asr/projector.cpp new file mode 100644 index 00000000..f01d2647 --- /dev/null +++ b/src/community_models/audio8_asr/projector.cpp @@ -0,0 +1,346 @@ +#include "engine/community_models/audio8_asr/projector.h" + +#include "engine/framework/assets/tensor_source.h" +#include "engine/framework/core/backend.h" +#include "engine/framework/core/backend_weight_store.h" +#include "engine/framework/core/module.h" +#include "engine/framework/debug/profiler.h" +#include "engine/framework/modules/linear_module.h" +#include "engine/framework/modules/norm_modules.h" +#include "engine/framework/modules/primitive_modules.h" +#include "engine/framework/runtime/errors.h" + +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace engine::community_models::audio8_asr { +namespace { + +namespace modules = engine::modules; +using Clock = std::chrono::steady_clock; + +struct GgmlContextDeleter { + void operator()(ggml_context * ctx) const noexcept { + if (ctx != nullptr) { + ggml_free(ctx); + } + } +}; + +struct GgmlGallocrDeleter { + void operator()(ggml_gallocr_t alloc) const noexcept { + if (alloc != nullptr) { + ggml_gallocr_free(alloc); + } + } +}; + +struct TowerLayerWeights { + core::TensorValue norm_weight; + core::TensorValue norm_bias; + core::TensorValue fc1_weight; + core::TensorValue fc1_bias; + core::TensorValue fc2_weight; + core::TensorValue fc2_bias; +}; + +struct ProjectorWeights { + std::shared_ptr store; + std::vector layers; + core::TensorValue final_norm_weight; + core::TensorValue final_norm_bias; + core::TensorValue projector_norm_weight; + core::TensorValue projector_norm_bias; + core::TensorValue projector_weight; + core::TensorValue projector_bias; +}; + +ProjectorWeights load_weights( + const assets::TensorSource & source, + const Audio8TowerConfig & config, + ggml_backend_t backend, + core::BackendType backend_type, + size_t weight_context_bytes, + assets::TensorStorageType storage_type) { + ProjectorWeights weights; + weights.store = std::make_shared( + backend, + backend_type, + "audio8_asr.projector.weights", + weight_context_bytes); + for (int64_t layer = 0; layer < config.layers; ++layer) { + const std::string prefix = + "audio_mlp_tower.layers." + std::to_string(layer); + TowerLayerWeights w; + w.norm_weight = weights.store->load_f32_tensor( + source, prefix + ".norm.weight", {config.hidden_size}); + w.norm_bias = weights.store->load_f32_tensor( + source, prefix + ".norm.bias", {config.hidden_size}); + w.fc1_weight = weights.store->load_tensor( + source, prefix + ".fc1.weight", storage_type, {config.intermediate_size, config.hidden_size}); + w.fc1_bias = weights.store->load_f32_tensor( + source, prefix + ".fc1.bias", {config.intermediate_size}); + w.fc2_weight = weights.store->load_tensor( + source, prefix + ".fc2.weight", storage_type, {config.hidden_size, config.intermediate_size}); + w.fc2_bias = weights.store->load_f32_tensor( + source, prefix + ".fc2.bias", {config.hidden_size}); + weights.layers.push_back(std::move(w)); + } + weights.final_norm_weight = weights.store->load_f32_tensor( + source, "audio_mlp_tower.final_norm.weight", {config.hidden_size}); + weights.final_norm_bias = weights.store->load_f32_tensor( + source, "audio_mlp_tower.final_norm.bias", {config.hidden_size}); + weights.projector_norm_weight = weights.store->load_f32_tensor( + source, "audio_projector.0.weight", {config.hidden_size}); + weights.projector_norm_bias = weights.store->load_f32_tensor( + source, "audio_projector.0.bias", {config.hidden_size}); + weights.projector_weight = weights.store->load_tensor( + source, "audio_projector.1.weight", storage_type, {config.output_size, config.hidden_size}); + weights.projector_bias = weights.store->load_f32_tensor( + source, "audio_projector.1.bias", {config.output_size}); + weights.store->upload(); + return weights; +} + +// torch adaptive_avg_pool1d output i averages input positions +// [floor(i * T / N), ceil((i + 1) * T / N)). The matrix is laid out for +// ggml_mul_mat(A, B) with A = transposed hidden [T, hidden] and +// B = this matrix [T, N]: element (t, n) lives at t + n * T. +std::vector adaptive_pool_matrix(int64_t input_tokens, int64_t output_tokens) { + std::vector matrix( + static_cast(input_tokens * output_tokens), 0.0f); + for (int64_t out = 0; out < output_tokens; ++out) { + const int64_t start = (out * input_tokens) / output_tokens; + const int64_t end = ((out + 1) * input_tokens + output_tokens - 1) / output_tokens; + const float scale = 1.0f / static_cast(end - start); + for (int64_t in = start; in < end; ++in) { + matrix[static_cast(in + out * input_tokens)] = scale; + } + } + return matrix; +} + +class ProjectorGraph { +public: + ProjectorGraph( + std::shared_ptr config, + std::shared_ptr weights, + core::BackendType backend_type, + ggml_backend_t backend, + int64_t input_tokens, + int64_t output_tokens, + size_t graph_arena_bytes) + : config_(std::move(config)), + weights_(std::move(weights)), + backend_(backend), + input_tokens_(input_tokens), + output_tokens_(output_tokens) { + const auto build_start = Clock::now(); + ggml_init_params params{graph_arena_bytes, nullptr, true}; + ctx_.reset(ggml_init(params)); + if (ctx_ == nullptr) { + throw std::runtime_error("failed to initialize Audio8 ASR projector graph context"); + } + core::ModuleBuildContext ctx{ctx_.get(), "audio8_asr.projector", backend_type}; + + input_ = ggml_new_tensor_2d( + ctx_.get(), GGML_TYPE_F32, config_->hidden_size, input_tokens_); + auto x = core::wrap_tensor( + input_, + core::TensorShape::from_dims({1, input_tokens_, config_->hidden_size}), + GGML_TYPE_F32); + for (const auto & layer : weights_->layers) { + auto normed = modules::LayerNormModule( + {config_->hidden_size, config_->norm_eps, true, true}) + .build(ctx, x, {layer.norm_weight, layer.norm_bias}); + auto fc1 = modules::LinearModule( + {config_->hidden_size, config_->intermediate_size, true, GGML_PREC_F32}) + .build(ctx, normed, {layer.fc1_weight, layer.fc1_bias}); + auto act = core::wrap_tensor( + ggml_gelu_erf(ctx.ggml, fc1.tensor), + core::TensorShape::from_dims({1, input_tokens_, config_->intermediate_size}), + GGML_TYPE_F32); + auto fc2 = modules::LinearModule( + {config_->intermediate_size, config_->hidden_size, true, GGML_PREC_F32}) + .build(ctx, act, {layer.fc2_weight, layer.fc2_bias}); + x = core::wrap_tensor( + ggml_add(ctx.ggml, x.tensor, fc2.tensor), + core::TensorShape::from_dims({1, input_tokens_, config_->hidden_size}), + GGML_TYPE_F32); + } + auto final_norm = modules::LayerNormModule( + {config_->hidden_size, config_->norm_eps, true, true}) + .build(ctx, x, {weights_->final_norm_weight, weights_->final_norm_bias}); + + pool_matrix_ = ggml_new_tensor_2d(ctx_.get(), GGML_TYPE_F32, input_tokens_, output_tokens_); + auto pool = core::wrap_tensor( + pool_matrix_, + core::TensorShape::from_dims({output_tokens_, input_tokens_}), + GGML_TYPE_F32); + // Materialized tower features [T, hidden]: mul_mat's "weight" operand + // must be contiguous, so the transpose view is copied eagerly. + auto hidden_transposed = core::wrap_tensor( + ggml_cont(ctx.ggml, ggml_transpose(ctx.ggml, final_norm.tensor)), + core::TensorShape::from_dims({config_->hidden_size, input_tokens_}), + GGML_TYPE_F32); + // mul_mat([T, hidden], [T, N]) -> [hidden, N]: pooled rows stay + // token-major for the projector that follows. + auto pooled = core::wrap_tensor( + ggml_mul_mat(ctx.ggml, hidden_transposed.tensor, pool.tensor), + core::TensorShape::from_dims({1, output_tokens_, config_->hidden_size}), + GGML_TYPE_F32); + + auto projector_norm = modules::LayerNormModule( + {config_->hidden_size, config_->norm_eps, true, true}) + .build(ctx, pooled, {weights_->projector_norm_weight, weights_->projector_norm_bias}); + auto projected = modules::LinearModule( + {config_->hidden_size, config_->output_size, true, GGML_PREC_F32}) + .build(ctx, projector_norm, {weights_->projector_weight, weights_->projector_bias}); + output_ = projected.tensor; + ggml_set_output(output_); + graph_ = ggml_new_graph_custom(ctx_.get(), 65536, false); + ggml_build_forward_expand(graph_, output_); + + const auto try_alloc = [&]() { + gallocr_.reset(ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend_))); + return gallocr_ != nullptr && + ggml_gallocr_reserve(gallocr_.get(), graph_) && + ggml_gallocr_alloc_graph(gallocr_.get(), graph_); + }; + if (!try_alloc() && + (engine::core::trim_backend_pools(backend_), !try_alloc())) { + throw engine::runtime::CapacityError( + "Audio8 ASR projector graph does not fit in device memory at this size (" + + std::to_string(input_tokens_) + " encoder tokens, " + + std::to_string(output_tokens_) + " audio tokens)"); + } + pool_values_ = adaptive_pool_matrix(input_tokens_, output_tokens_); + debug::timing_log_scalar("audio8_asr.projector.graph.build_ms", engine::debug::elapsed_ms(build_start, Clock::now())); + debug::trace_log_scalar("audio8_asr.projector.encoder_tokens", input_tokens_); + debug::trace_log_scalar("audio8_asr.projector.audio_tokens", output_tokens_); + } + + ~ProjectorGraph() { + engine::core::release_backend_graph_resources(backend_, graph_, true); + } + + bool matches(int64_t input_tokens, int64_t output_tokens) const { + return input_tokens_ == input_tokens && output_tokens_ == output_tokens; + } + + std::vector run(const std::vector & encoder_output) { + if (static_cast(encoder_output.size()) != input_tokens_ * config_->hidden_size) { + throw std::runtime_error("Audio8 ASR projector encoder output size mismatch"); + } + const auto upload_start = Clock::now(); + // Re-uploaded every run: leaves are not pinned by the graph allocator + // (the gallocr may hand a leaf's buffer to a later node once consumed), + // so both the input and the constant pool matrix are rewritten here. + ggml_backend_tensor_set(input_, encoder_output.data(), 0, encoder_output.size() * sizeof(float)); + ggml_backend_tensor_set( + pool_matrix_, pool_values_.data(), 0, pool_values_.size() * sizeof(float)); + debug::timing_log_scalar("audio8_asr.projector.input_upload_ms", engine::debug::elapsed_ms(upload_start, Clock::now())); + const ggml_status status = engine::core::compute_backend_graph(backend_, graph_); + ggml_backend_synchronize(backend_); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("Audio8 ASR projector graph compute failed"); + } + std::vector output(static_cast(output_tokens_ * config_->output_size)); + ggml_backend_tensor_get(output_, output.data(), 0, output.size() * sizeof(float)); + return output; + } + +private: + std::shared_ptr config_; + std::shared_ptr weights_; + ggml_backend_t backend_ = nullptr; + int64_t input_tokens_ = 0; + int64_t output_tokens_ = 0; + std::unique_ptr ctx_; + ggml_tensor * input_ = nullptr; + ggml_tensor * pool_matrix_ = nullptr; + ggml_tensor * output_ = nullptr; + std::vector pool_values_; + ggml_cgraph * graph_ = nullptr; + std::unique_ptr, GgmlGallocrDeleter> gallocr_; +}; + +} // namespace + +struct Audio8ProjectorRuntime::Impl { + Impl( + std::shared_ptr weights_source, + Audio8TowerConfig config, + core::ExecutionContext & execution, + size_t graph_arena_bytes, + size_t weight_context_bytes, + assets::TensorStorageType weight_storage_type) + : config_(std::make_shared(std::move(config))), + weights_(std::make_shared(load_weights( + *weights_source, + *config_, + execution.backend(), + execution.backend_type(), + weight_context_bytes, + weight_storage_type))), + execution_(execution), + graph_arena_bytes_(graph_arena_bytes) {} + + std::shared_ptr config_; + std::shared_ptr weights_; + core::ExecutionContext & execution_; + size_t graph_arena_bytes_ = 0; + std::unique_ptr graph_; +}; + +Audio8ProjectorRuntime::Audio8ProjectorRuntime( + std::shared_ptr weights_source, + const Audio8TowerConfig & config, + core::ExecutionContext & execution, + size_t graph_arena_bytes, + size_t weight_context_bytes, + assets::TensorStorageType weight_storage_type) + : impl_(std::make_unique( + std::move(weights_source), + config, + execution, + graph_arena_bytes, + weight_context_bytes, + weight_storage_type)) {} + +Audio8ProjectorRuntime::~Audio8ProjectorRuntime() = default; + +Audio8ASRAudioEmbeddings Audio8ProjectorRuntime::project( + const std::vector & encoder_output, + int64_t encoder_tokens, + int64_t audio_tokens) { + if (encoder_tokens <= 0 || audio_tokens <= 0 || audio_tokens > encoder_tokens) { + throw std::runtime_error("Audio8 ASR projector token counts are invalid"); + } + if (impl_->graph_ == nullptr || !impl_->graph_->matches(encoder_tokens, audio_tokens)) { + impl_->graph_.reset(); + impl_->graph_ = std::make_unique( + impl_->config_, + impl_->weights_, + impl_->execution_.backend_type(), + impl_->execution_.backend(), + encoder_tokens, + audio_tokens, + impl_->graph_arena_bytes_); + } + Audio8ASRAudioEmbeddings out; + out.values = impl_->graph_->run(encoder_output); + out.tokens = audio_tokens; + out.hidden_size = impl_->config_->output_size; + return out; +} + +} // namespace engine::community_models::audio8_asr diff --git a/src/community_models/audio8_asr/session.cpp b/src/community_models/audio8_asr/session.cpp new file mode 100644 index 00000000..d62ff774 --- /dev/null +++ b/src/community_models/audio8_asr/session.cpp @@ -0,0 +1,398 @@ +#include "engine/community_models/audio8_asr/session.h" + +#include "engine/framework/debug/profiler.h" +#include "engine/framework/debug/trace.h" +#include "engine/framework/io/json.h" +#include "engine/framework/io/text.h" +#include "engine/framework/model_spec/package.h" +#include "engine/framework/runtime/options.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace engine::community_models::audio8_asr { +namespace json = engine::io::json; +namespace { + +using Clock = std::chrono::steady_clock; +constexpr size_t kWeightContextBytes = 64ull * 1024ull * 1024ull; +// Languages advertised by the loader and the loaded model; kept in one list. +const std::vector kLanguages = { + "Chinese", "English", "Cantonese", "French", "German", "Japanese", "Korean"}; + +std::filesystem::path default_spec_path() { + return engine::model_spec::default_spec_path("audio8_asr"); +} + +std::shared_ptr require_assets(std::shared_ptr assets) { + if (assets == nullptr) { + throw std::runtime_error("Audio8 ASR session requires assets"); + } + return assets; +} + +void validate_audio_encoder_weight_storage(assets::TensorStorageType storage_type) { + if (storage_type == assets::TensorStorageType::Native || + storage_type == assets::TensorStorageType::F32 || + storage_type == assets::TensorStorageType::F16) { + return; + } + throw std::runtime_error("audio8_asr.audio_encoder_weight_type currently supports only native, f32, and f16"); +} + +void validate_matmul_weight_storage(assets::TensorStorageType storage_type) { + if (storage_type == assets::TensorStorageType::Native || + storage_type == assets::TensorStorageType::F32 || + storage_type == assets::TensorStorageType::F16 || + storage_type == assets::TensorStorageType::BF16 || + storage_type == assets::TensorStorageType::Q8_0) { + return; + } + throw std::runtime_error("audio8_asr.weight_type supports only native, f32, f16, bf16, and q8_0"); +} + +assets::TensorStorageType parse_encoder_weight_storage(const runtime::SessionOptions & session_options) { + const auto & options = session_options.options; + // The encoder rejects on-the-fly quantized weights; already-quantized + // GGUF tensors load as Native regardless of this option. + const auto type = runtime::parse_tensor_storage_option( + options, + "audio8_asr.audio_encoder_weight_type", + assets::TensorStorageType::Native, + {assets::TensorStorageType::Native, assets::TensorStorageType::F32, assets::TensorStorageType::F16}); + validate_audio_encoder_weight_storage(type); + return type; +} + +assets::TensorStorageType parse_matmul_weight_storage(const runtime::SessionOptions & session_options) { + const auto & options = session_options.options; + const auto type = runtime::parse_tensor_storage_option( + options, + "audio8_asr.weight_type", + assets::TensorStorageType::Native, + {assets::TensorStorageType::Native, assets::TensorStorageType::F32, + assets::TensorStorageType::F16, assets::TensorStorageType::BF16, + assets::TensorStorageType::Q8_0}); + validate_matmul_weight_storage(type); + return type; +} + +size_t encoder_graph_arena_bytes(const runtime::SessionOptions & session_options) { + const auto & options = session_options.options; + return runtime::parse_size_mb_option(options, {"audio8_asr.encoder_graph_arena_mb"}, 128ull * 1024ull * 1024ull); +} + +size_t projector_graph_arena_bytes(const runtime::SessionOptions & session_options) { + const auto & options = session_options.options; + return runtime::parse_size_mb_option(options, {"audio8_asr.projector_graph_arena_mb"}, 256ull * 1024ull * 1024ull); +} + +size_t prefill_graph_arena_bytes(const runtime::SessionOptions & session_options) { + const auto & options = session_options.options; + return runtime::parse_size_mb_option(options, {"audio8_asr.prefill_graph_arena_mb"}, 256ull * 1024ull * 1024ull); +} + +size_t decode_graph_arena_bytes(const runtime::SessionOptions & session_options) { + const auto & options = session_options.options; + return runtime::parse_size_mb_option(options, {"audio8_asr.decode_graph_arena_mb"}, 256ull * 1024ull * 1024ull); +} + +// The reference processor hands the model a bfloat16 mel tensor; replicate +// that rounding (audio8_asr_round_f32_to_bf16) before encoding. + +// Join window transcripts. Whitespace is only inserted between ASCII words; +// CJK text has no word boundaries, so a byte >= 0x80 on either side means the +// windows are concatenated directly (the reference has no multi-window path; +// this mirrors how CJK text is written). +void append_clip_transcript(std::string & merged, std::string clip_text) { + clip_text = engine::io::trim_ascii_whitespace(std::move(clip_text)); + if (clip_text.empty()) { + return; + } + if (!merged.empty()) { + const unsigned char last = static_cast(merged.back()); + const unsigned char first = static_cast(clip_text.front()); + if (last < 0x80 && first < 0x80 && last != ' ' && first != ' ') { + merged.push_back(' '); + } + } + merged += clip_text; +} + +class Audio8ASRLoader final : public runtime::IVoiceModelLoader { +public: + std::string family() const override { + return "audio8_asr"; + } + + std::vector family_aliases() const override { + return {"arkasr"}; + } + + runtime::CapabilitySet advertised_capabilities() const override { + runtime::CapabilitySet out; + out.supported_tasks = { + {runtime::VoiceTaskKind::Asr, {runtime::RunMode::Offline}}, + }; + out.languages = kLanguages; + out.supports_timestamps = false; + return out; + } + + bool can_load(const runtime::ModelLoadRequest & request) const override { + if (request.family_hint.has_value()) { + const auto & hint = *request.family_hint; + if (hint != family()) { + const auto aliases = family_aliases(); + if (std::find(aliases.begin(), aliases.end(), hint) == aliases.end()) { + return false; + } + } + } + try { + const auto resources = engine::model_spec::load_resource_bundle(request.model_path, default_spec_path()); + const auto config_root = resources.parse_json("config"); + return json::require_string(config_root, "model_type") == "arkasr"; + } catch (...) { + return false; + } + } + + runtime::ModelInspection inspect(const runtime::ModelLoadRequest & request) const override { + const auto resources = engine::model_spec::load_resource_bundle( + request.model_path, + default_spec_path()); + runtime::ModelInspection inspection; + inspection.model_root = resources.model_root(); + inspection.metadata.family = family(); + inspection.metadata.variant = "0.1b"; + inspection.metadata.description = "Audio8-ASR-0.1B multilingual ASR model."; + inspection.capabilities = advertised_capabilities(); + inspection.discovered_configs = runtime::discover_named_assets_from_package_spec( + request.model_path, + default_spec_path(), + engine::model_spec::ResourceKind::Files); + inspection.discovered_weights = runtime::discover_named_assets_from_package_spec( + request.model_path, + default_spec_path(), + engine::model_spec::ResourceKind::Tensors); + return inspection; + } + + std::unique_ptr load(const runtime::ModelLoadRequest & request) const override { + return load_audio8_asr_model(request); + } +}; + +} // namespace + +Audio8ASRSession::Audio8ASRSession( + runtime::TaskSpec task, + runtime::SessionOptions options, + std::shared_ptr assets) + : RuntimeSessionBase(std::move(options)), + task_(std::move(task)), + assets_(require_assets(std::move(assets))), + frontend_(assets_->encoder_assets), + audio_encoder_( + assets_->encoder_assets, + execution_context(), + encoder_graph_arena_bytes(RuntimeSessionBase::options()), + parse_encoder_weight_storage(RuntimeSessionBase::options())), + projector_( + assets_->resources.open_tensor_source("weights"), + assets_->config.tower, + execution_context(), + projector_graph_arena_bytes(RuntimeSessionBase::options()), + kWeightContextBytes, + parse_matmul_weight_storage(RuntimeSessionBase::options())), + thinker_( + assets_->resources.open_tensor_source("weights"), + assets_->config.text_decoder, + execution_context(), + prefill_graph_arena_bytes(RuntimeSessionBase::options()), + decode_graph_arena_bytes(RuntimeSessionBase::options()), + kWeightContextBytes, + parse_matmul_weight_storage(RuntimeSessionBase::options())) { + if (task_.task != runtime::VoiceTaskKind::Asr) { + throw std::runtime_error("Audio8 ASR only supports VoiceTaskKind::Asr"); + } + if (task_.mode != runtime::RunMode::Offline) { + throw std::runtime_error("Audio8 ASR only supports offline sessions"); + } + // All weight stores have uploaded by now; drop the resident file blob. + assets_->resources.open_tensor_source("weights")->release_storage(); +} + +Audio8ASRSession::~Audio8ASRSession() = default; + +std::string Audio8ASRSession::family() const { + return "audio8_asr"; +} + +runtime::VoiceTaskKind Audio8ASRSession::task_kind() const { + return task_.task; +} + +runtime::RunMode Audio8ASRSession::run_mode() const { + return task_.mode; +} + +void Audio8ASRSession::prepare(const runtime::SessionPreparationRequest & request) { + (void)request; + mark_prepared(); +} + +std::string Audio8ASRSession::transcribe_clip(const runtime::AudioBuffer & audio) { + auto features = frontend_.extract(audio); + // Match the reference mel dtype (bfloat16) before encoding. + for (auto & value : features.values) { + value = audio8_asr_round_f32_to_bf16(value); + } + const auto encoder_output = audio_encoder_.encode(features); + + const int64_t mel_frames = features.frames; + const int64_t audio_tokens = audio8_asr_prompt_audio_token_count(mel_frames, assets_->config.merge_factor); + const auto embeddings = projector_.project(encoder_output.values, encoder_output.tokens, audio_tokens); + + // Prompt: <|user|><|begin_of_audio|><|audio|>xN<|end_of_audio|>Please + // transcribe this audio.<|assistant|> + const auto & config = assets_->config; + Audio8ASRPrompt prompt; + prompt.input_ids.push_back(static_cast(config.user_token_id)); + prompt.input_ids.push_back(static_cast(config.begin_audio_token_id)); + for (int64_t i = 0; i < audio_tokens; ++i) { + prompt.audio_token_positions.push_back(static_cast(prompt.input_ids.size())); + prompt.input_ids.push_back(static_cast(config.text_decoder.audio_token_id)); + } + prompt.input_ids.push_back(static_cast(config.end_audio_token_id)); + const auto instruction_ids = assets_->tokenizer->encode("Please transcribe this audio.", false); + prompt.input_ids.insert(prompt.input_ids.end(), instruction_ids.begin(), instruction_ids.end()); + prompt.input_ids.push_back(static_cast(config.assistant_token_id)); + + const Audio8ASRGenerationOptions generation{config.text_decoder.max_new_tokens}; + const auto generated = thinker_.generate(prompt, embeddings, generation); + if (generated.token_ids.empty()) { + return ""; + } + std::vector filtered; + filtered.reserve(generated.token_ids.size()); + for (const int32_t id : generated.token_ids) { + if (id == static_cast(config.text_decoder.pad_token_id) || + std::find( + config.text_decoder.eos_token_ids.begin(), + config.text_decoder.eos_token_ids.end(), + static_cast(id)) != config.text_decoder.eos_token_ids.end() || + assets_->tokenizer->is_control_token_id(id)) { + continue; + } + filtered.push_back(id); + } + if (filtered.empty()) { + return ""; + } + return assets_->tokenizer->decode(filtered); +} + +runtime::Transcript Audio8ASRSession::transcribe_audio(const runtime::AudioBuffer & audio) { + if (audio.samples.empty()) { + return {"", ""}; + } + // The reference clips input at max_audio_samples (16 kHz samples, 30 s by + // default); longer recordings are covered by transcribing fixed windows + // and joining the text. The window is sized in the input sample domain + // (interleaved values, input rate) so it always covers the same wall + // time as the reference clip. A tail shorter than half a second folds + // into the previous window instead of triggering a near-silent pass. + const int64_t channels = std::max(1, audio.channels); + const int64_t sample_rate = std::max(1, audio.sample_rate); + const int64_t window_values = + assets_->config.max_audio_samples * channels * sample_rate / 16000; + const int64_t min_tail_values = sample_rate * channels / 2; + std::string merged; + int64_t offset = 0; + const int64_t total = static_cast(audio.samples.size()); + while (offset < total) { + int64_t span = std::min(window_values, total - offset); + if (total - (offset + span) < min_tail_values) { + span = total - offset; + } + runtime::AudioBuffer clip; + clip.sample_rate = audio.sample_rate; + clip.channels = audio.channels; + clip.samples.assign( + audio.samples.begin() + static_cast(offset), + audio.samples.begin() + static_cast(offset + span)); + append_clip_transcript(merged, transcribe_clip(clip)); + offset += span; + } + return {engine::io::trim_ascii_whitespace(std::move(merged)), ""}; +} + +runtime::TaskResult Audio8ASRSession::run(const runtime::TaskRequest & request) { + require_prepared("Audio8 ASR run()"); + if (!request.audio_input.has_value()) { + throw std::runtime_error("Audio8 ASR run() requires audio_input"); + } + const auto wall_start = Clock::now(); + auto transcript = transcribe_audio(*request.audio_input); + runtime::TaskResult result; + result.text_output = std::move(transcript); + engine::debug::timing_log_scalar("session.wall_ms", engine::debug::elapsed_ms(wall_start)); + return result; +} + +Audio8ASRLoadedModel::Audio8ASRLoadedModel( + runtime::ModelMetadata metadata, + runtime::CapabilitySet capabilities, + std::shared_ptr assets) + : metadata_(std::move(metadata)), + capabilities_(std::move(capabilities)), + assets_(std::move(assets)) {} + +const runtime::ModelMetadata & Audio8ASRLoadedModel::metadata() const noexcept { + return metadata_; +} + +const runtime::CapabilitySet & Audio8ASRLoadedModel::capabilities() const noexcept { + return capabilities_; +} + +std::unique_ptr Audio8ASRLoadedModel::create_task_session( + const runtime::TaskSpec & task, + const runtime::SessionOptions & options) const { + // The session constructor validates the task/mode combination. + return std::make_unique(task, options, assets_); +} + +std::unique_ptr load_audio8_asr_model(const runtime::ModelLoadRequest & request) { + auto assets = load_audio8_asr_assets(request.model_path); + + runtime::ModelMetadata metadata; + metadata.family = "audio8_asr"; + metadata.variant = "0.1b"; + metadata.description = "Audio8-ASR-0.1B multilingual ASR model loaded from local assets."; + + runtime::CapabilitySet capabilities; + capabilities.supported_tasks = { + {runtime::VoiceTaskKind::Asr, {runtime::RunMode::Offline}}, + }; + capabilities.languages = assets->config.supported_languages; + capabilities.supports_timestamps = false; + + return std::make_unique( + std::move(metadata), + std::move(capabilities), + std::move(assets)); +} + +std::shared_ptr make_audio8_asr_loader() { + return std::make_shared(); +} + +} // namespace engine::community_models::audio8_asr diff --git a/src/community_models/audio8_asr/thinker.cpp b/src/community_models/audio8_asr/thinker.cpp new file mode 100644 index 00000000..707c2fdd --- /dev/null +++ b/src/community_models/audio8_asr/thinker.cpp @@ -0,0 +1,689 @@ +#include "engine/community_models/audio8_asr/thinker.h" + +#include "engine/framework/assets/tensor_source.h" +#include "engine/framework/core/backend.h" +#include "engine/framework/core/backend_weight_store.h" +#include "engine/framework/debug/profiler.h" +#include "engine/framework/modules/activation_modules.h" +#include "engine/framework/modules/transformers/qwen_causal_decoder.h" +#include "engine/framework/modules/linear_module.h" +#include "engine/framework/modules/lookup_modules.h" +#include "engine/framework/modules/norm_modules.h" +#include "engine/framework/modules/positional_modules.h" +#include "engine/framework/modules/primitive_modules.h" +#include "engine/framework/modules/structural_modules.h" +#include "engine/framework/runtime/errors.h" +#include "engine/framework/runtime/kv_cache.h" +#include "engine/framework/sampling/decode_modules.h" + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace engine::community_models::audio8_asr { +namespace { + +namespace modules = engine::modules; +using Clock = std::chrono::steady_clock; + +struct GgmlContextDeleter { + void operator()(ggml_context * ctx) const noexcept { + if (ctx != nullptr) { + ggml_free(ctx); + } + } +}; + +struct GgmlGallocrDeleter { + void operator()(ggml_gallocr_t alloc) const noexcept { + if (alloc != nullptr) { + ggml_gallocr_free(alloc); + } + } +}; + +struct TextLayerWeights { + core::TensorValue input_norm; + core::TensorValue q_proj; + core::TensorValue q_bias; + core::TensorValue k_proj; + core::TensorValue k_bias; + core::TensorValue v_proj; + core::TensorValue v_bias; + core::TensorValue o_proj; + core::TensorValue post_norm; + core::TensorValue gate_proj; + core::TensorValue up_proj; + core::TensorValue down_proj; +}; + +struct ThinkerWeights { + std::shared_ptr store; + core::TensorValue token_embedding; + std::vector layers; + core::TensorValue norm; + core::TensorValue lm_head; +}; + +struct PrefillOutput { + std::vector logits; + runtime::TransformerKVState kv_state; +}; + +modules::QwenDecoderLayerWeights to_qwen_layer_weights(const TextLayerWeights & weights) { + modules::QwenDecoderLayerWeights out; + out.input_norm = {weights.input_norm, std::nullopt}; + out.self_attention.q_weight = weights.q_proj; + out.self_attention.q_bias = weights.q_bias; + out.self_attention.k_weight = weights.k_proj; + out.self_attention.k_bias = weights.k_bias; + out.self_attention.v_weight = weights.v_proj; + out.self_attention.v_bias = weights.v_bias; + out.self_attention.out_weight = weights.o_proj; + out.post_norm = {weights.post_norm, std::nullopt}; + out.mlp.gate_proj = {weights.gate_proj, std::nullopt}; + out.mlp.up_proj = {weights.up_proj, std::nullopt}; + out.mlp.down_proj = {weights.down_proj, std::nullopt}; + return out; +} + +// Audio8 decoders are Qwen2-style: no per-head Q/K norms and separate Q/K/V +// projections with attention biases. +modules::QwenCausalDecoderConfig make_qwen_decoder_config(const Audio8ASRDecoderConfig & config) { + modules::QwenCausalDecoderConfig out; + out.stack.hidden_size = config.hidden_size; + out.stack.num_attention_heads = config.num_attention_heads; + out.stack.num_key_value_heads = config.num_key_value_heads; + out.stack.head_dim = config.head_dim; + out.stack.intermediate_size = config.intermediate_size; + out.stack.layers = config.num_hidden_layers; + out.stack.rms_norm_eps = config.rms_norm_eps; + out.stack.rope_theta = config.rope_theta; + out.stack.use_qk_norm = false; + out.stack.runtime.static_cache.update_mode = modules::QwenDecoderStaticCacheUpdateMode::DirectSetRows; + out.logits_size = config.vocab_size; + out.logits_mode = modules::QwenCausalDecoderLogitsMode::LastStep; + return out; +} + +modules::QwenCausalDecoderWeights make_qwen_decoder_weights(const ThinkerWeights & weights) { + modules::QwenCausalDecoderWeights out; + out.stack.layers.reserve(weights.layers.size()); + for (const auto & layer : weights.layers) { + out.stack.layers.push_back(to_qwen_layer_weights(layer)); + } + out.final_norm = {weights.norm, std::nullopt}; + out.lm_head = {weights.lm_head, std::nullopt}; + return out; +} + +core::TensorValue prompt_embeddings( + core::ModuleBuildContext & ctx, + const ThinkerWeights & weights, + const Audio8ASRDecoderConfig & config, + ggml_tensor * token_ids, + ggml_tensor * audio_embeddings, + ggml_tensor * audio_positions, + int64_t prompt_steps, + int64_t audio_tokens) { + auto ids = core::wrap_tensor(token_ids, core::TensorShape::from_dims({prompt_steps}), GGML_TYPE_I32); + auto x = modules::EmbeddingModule({config.vocab_size, config.hidden_size}).build(ctx, ids, weights.token_embedding); + if (audio_tokens > 0) { + auto audio = core::wrap_tensor( + audio_embeddings, + core::TensorShape::from_dims({audio_tokens, config.hidden_size}), + GGML_TYPE_F32); + auto positions = core::wrap_tensor( + audio_positions, + core::TensorShape::from_dims({audio_tokens}), + GGML_TYPE_I64); + x = core::wrap_tensor( + ggml_set_rows(ctx.ggml, x.tensor, audio.tensor, positions.tensor), + x.shape, + GGML_TYPE_F32); + } + return core::reshape_tensor(ctx, x, core::TensorShape::from_dims({1, prompt_steps, config.hidden_size})); +} + +ThinkerWeights load_weights( + const assets::TensorSource & source, + const Audio8ASRDecoderConfig & config, + ggml_backend_t backend, + core::BackendType backend_type, + size_t weight_context_bytes, + assets::TensorStorageType storage_type) { + ThinkerWeights weights; + weights.store = std::make_shared( + backend, + backend_type, + "audio8_asr.thinker.weights", + weight_context_bytes); + weights.token_embedding = weights.store->load_tensor( + source, + "language_model.model.embed_tokens.weight", + storage_type, + {config.vocab_size, config.hidden_size}); + weights.layers.reserve(static_cast(config.num_hidden_layers)); + const int64_t dim = config.head_dim; + for (int64_t layer = 0; layer < config.num_hidden_layers; ++layer) { + const std::string prefix = "language_model.model.layers." + std::to_string(layer); + TextLayerWeights w; + w.input_norm = weights.store->load_f32_tensor(source, prefix + ".input_layernorm.weight", {config.hidden_size}); + w.q_proj = weights.store->load_tensor(source, prefix + ".self_attn.q_proj.weight", storage_type, {config.num_attention_heads * dim, config.hidden_size}); + w.q_bias = weights.store->load_f32_tensor(source, prefix + ".self_attn.q_proj.bias", {config.num_attention_heads * dim}); + w.k_proj = weights.store->load_tensor(source, prefix + ".self_attn.k_proj.weight", storage_type, {config.num_key_value_heads * dim, config.hidden_size}); + w.k_bias = weights.store->load_f32_tensor(source, prefix + ".self_attn.k_proj.bias", {config.num_key_value_heads * dim}); + w.v_proj = weights.store->load_tensor(source, prefix + ".self_attn.v_proj.weight", storage_type, {config.num_key_value_heads * dim, config.hidden_size}); + w.v_bias = weights.store->load_f32_tensor(source, prefix + ".self_attn.v_proj.bias", {config.num_key_value_heads * dim}); + w.o_proj = weights.store->load_tensor(source, prefix + ".self_attn.o_proj.weight", storage_type, {config.hidden_size, config.num_attention_heads * dim}); + w.post_norm = weights.store->load_f32_tensor(source, prefix + ".post_attention_layernorm.weight", {config.hidden_size}); + w.gate_proj = weights.store->load_tensor(source, prefix + ".mlp.gate_proj.weight", storage_type, {config.intermediate_size, config.hidden_size}); + w.up_proj = weights.store->load_tensor(source, prefix + ".mlp.up_proj.weight", storage_type, {config.intermediate_size, config.hidden_size}); + w.down_proj = weights.store->load_tensor(source, prefix + ".mlp.down_proj.weight", storage_type, {config.hidden_size, config.intermediate_size}); + weights.layers.push_back(std::move(w)); + } + weights.norm = weights.store->load_f32_tensor(source, "language_model.model.norm.weight", {config.hidden_size}); + if (config.tie_word_embeddings) { + weights.lm_head = weights.token_embedding; + } else { + weights.lm_head = weights.store->load_tensor( + source, + "language_model.lm_head.weight", + storage_type, + {config.vocab_size, config.hidden_size}); + } + weights.store->upload(); + return weights; +} + +int32_t argmax_index(const std::vector & values) { + if (values.empty()) { + throw std::runtime_error("Audio8 ASR thinker cannot select from empty logits"); + } + size_t best = 0; + for (size_t i = 1; i < values.size(); ++i) { + if (values[i] > values[best]) { + best = i; + } + } + return static_cast(best); +} + +bool is_eos(const Audio8ASRDecoderConfig & config, int32_t token) { + return std::find(config.eos_token_ids.begin(), config.eos_token_ids.end(), static_cast(token)) != + config.eos_token_ids.end(); +} + +class ThinkerWeightsRuntime { +public: + ThinkerWeightsRuntime( + std::shared_ptr source, + Audio8ASRDecoderConfig config, + core::ExecutionContext & execution, + size_t weight_context_bytes, + assets::TensorStorageType storage_type) + : source_(std::move(source)), + config_(std::make_shared(std::move(config))), + backend_(execution.backend()), + backend_type_(execution.backend_type()), + threads_(std::max(1, execution.config().threads)), + weights_(std::make_shared(load_weights( + *source_, + *config_, + backend_, + backend_type_, + weight_context_bytes, + storage_type))) {} + + const Audio8ASRDecoderConfig & config() const noexcept { + return *config_; + } + + const ThinkerWeights & weights() const noexcept { + return *weights_; + } + + ggml_backend_t backend() const noexcept { + return backend_; + } + + core::BackendType backend_type() const noexcept { + return backend_type_; + } + + int threads() const noexcept { + return threads_; + } + +private: + std::shared_ptr source_; + std::shared_ptr config_; + ggml_backend_t backend_ = nullptr; + core::BackendType backend_type_ = core::BackendType::Cpu; + int threads_ = 1; + std::shared_ptr weights_; +}; + +class PrefillGraph { +public: + PrefillGraph( + std::shared_ptr runtime, + int64_t prompt_steps, + int64_t audio_tokens, + size_t graph_arena_bytes) + : runtime_(std::move(runtime)), + prompt_steps_(prompt_steps), + audio_tokens_(audio_tokens) { + if (prompt_steps_ <= 0) { + throw std::runtime_error("Audio8 ASR thinker prefill requires positive prompt length"); + } + if (audio_tokens_ < 0 || audio_tokens_ > prompt_steps_) { + throw std::runtime_error("Audio8 ASR thinker prefill audio token count is invalid"); + } + const auto build_start = Clock::now(); + ggml_init_params params{graph_arena_bytes, nullptr, true}; + ctx_.reset(ggml_init(params)); + if (ctx_ == nullptr) { + throw std::runtime_error("failed to initialize Audio8 ASR thinker prefill graph context"); + } + const auto & config = runtime_->config(); + const auto & weights = runtime_->weights(); + core::ModuleBuildContext ctx{ctx_.get(), "audio8_asr.thinker.prefill", runtime_->backend_type()}; + token_ids_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_I32, prompt_steps_); + audio_embeddings_ = ggml_new_tensor_2d(ctx_.get(), GGML_TYPE_F32, config.hidden_size, std::max(audio_tokens_, 1)); + audio_positions_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_I64, std::max(audio_tokens_, 1)); + auto x = prompt_embeddings( + ctx, + weights, + config, + token_ids_, + audio_embeddings_, + audio_positions_, + prompt_steps_, + audio_tokens_); + positions_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_I32, prompt_steps_); + auto positions = core::wrap_tensor(positions_, core::TensorShape::from_dims({prompt_steps_}), GGML_TYPE_I32); + + auto decoder_out = modules::QwenCausalDecoderModule(make_qwen_decoder_config(config)) + .build(ctx, x, positions, make_qwen_decoder_weights(weights)); + for (const auto & layer : decoder_out.state.layers) { + if (!layer.key.has_value() || !layer.value.has_value()) { + throw std::runtime_error("Audio8 ASR thinker prefill decoder did not return K/V state"); + } + // See qwen3_asr PrefillGraph: copy K/V out of the graph-allocated + // intermediates and mark them as outputs so the allocator cannot + // recycle them before run() reads them back. + auto * key = ggml_cpy( + ctx_.get(), + layer.key->tensor, + ggml_dup_tensor(ctx_.get(), layer.key->tensor)); + auto * value = ggml_cpy( + ctx_.get(), + layer.value->tensor, + ggml_dup_tensor(ctx_.get(), layer.value->tensor)); + ggml_set_output(key); + ggml_set_output(value); + keys_.push_back(key); + values_.push_back(value); + } + logits_ = decoder_out.logits.tensor; + ggml_set_output(logits_); + graph_ = ggml_new_graph_custom(ctx_.get(), 65536, false); + ggml_build_forward_expand(graph_, logits_); + for (auto * key : keys_) { + ggml_build_forward_expand(graph_, key); + } + for (auto * value : values_) { + ggml_build_forward_expand(graph_, value); + } + const auto try_alloc = [&]() { + gallocr_.reset(ggml_gallocr_new(ggml_backend_get_default_buffer_type(runtime_->backend()))); + return gallocr_ != nullptr && + ggml_gallocr_reserve(gallocr_.get(), graph_) && + ggml_gallocr_alloc_graph(gallocr_.get(), graph_); + }; + if (!try_alloc() && + (engine::core::trim_backend_pools(runtime_->backend()), !try_alloc())) { + throw engine::runtime::CapacityError( + "Audio8 ASR prefill graph does not fit in device memory at this size (" + + std::to_string(prompt_steps_) + " prompt steps, of which " + + std::to_string(audio_tokens_) + " are audio tokens)"); + } + position_ids_ = modules::qwen_position_ids(prompt_steps_); + debug::timing_log_scalar("audio8_asr.thinker.prefill.graph.build_ms", engine::debug::elapsed_ms(build_start, Clock::now())); + debug::trace_log_scalar("audio8_asr.thinker.prefill_prompt_steps", prompt_steps_); + } + + ~PrefillGraph() { + engine::core::release_backend_graph_resources(runtime_->backend(), graph_, true); + } + + bool matches(const ThinkerWeightsRuntime & runtime, int64_t prompt_steps, int64_t audio_tokens) const { + return runtime_.get() == &runtime && prompt_steps_ == prompt_steps && audio_tokens_ == audio_tokens; + } + + PrefillOutput run( + const std::vector & token_ids, + const std::vector & audio_embeddings, + const std::vector & audio_positions) { + const auto & config = runtime_->config(); + if (static_cast(token_ids.size()) != prompt_steps_) { + throw std::runtime_error("Audio8 ASR thinker prefill token id count mismatch"); + } + if (static_cast(audio_embeddings.size()) != audio_tokens_ * config.hidden_size) { + throw std::runtime_error("Audio8 ASR thinker prefill audio embedding size mismatch"); + } + if (static_cast(audio_positions.size()) != audio_tokens_) { + throw std::runtime_error("Audio8 ASR thinker prefill audio position count mismatch"); + } + auto timing_start = Clock::now(); + // Re-uploaded on every run: leaves are not pinned by the graph allocator. + ggml_backend_tensor_set(positions_, position_ids_.data(), 0, position_ids_.size() * sizeof(int32_t)); + ggml_backend_tensor_set(token_ids_, token_ids.data(), 0, token_ids.size() * sizeof(int32_t)); + if (audio_tokens_ > 0) { + std::vector positions(audio_positions.begin(), audio_positions.end()); + ggml_backend_tensor_set( + audio_embeddings_, + audio_embeddings.data(), + 0, + audio_embeddings.size() * sizeof(float)); + ggml_backend_tensor_set( + audio_positions_, + positions.data(), + 0, + positions.size() * sizeof(int64_t)); + } + core::set_backend_threads(runtime_->backend(), runtime_->threads()); + timing_start = Clock::now(); + const ggml_status status = engine::core::compute_backend_graph(runtime_->backend(), graph_); + ggml_backend_synchronize(runtime_->backend()); + debug::timing_log_scalar("audio8_asr.thinker.prefill.graph.compute_ms", engine::debug::elapsed_ms(timing_start, Clock::now())); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("Audio8 ASR thinker prefill graph compute failed"); + } + PrefillOutput out; + out.logits.resize(static_cast(config.vocab_size)); + ggml_backend_tensor_get(logits_, out.logits.data(), 0, out.logits.size() * sizeof(float)); + out.kv_state.current_end = prompt_steps_; + out.kv_state.layers.resize(keys_.size()); + const size_t layer_values = static_cast( + prompt_steps_ * config.num_key_value_heads * config.head_dim); + for (size_t layer = 0; layer < keys_.size(); ++layer) { + auto & state = out.kv_state.layers[layer]; + state.valid_steps = prompt_steps_; + state.key.resize(layer_values); + state.value.resize(layer_values); + ggml_backend_tensor_get(keys_[layer], state.key.data(), 0, state.key.size() * sizeof(float)); + ggml_backend_tensor_get(values_[layer], state.value.data(), 0, state.value.size() * sizeof(float)); + } + return out; + } + +private: + std::shared_ptr runtime_; + int64_t prompt_steps_ = 0; + int64_t audio_tokens_ = 0; + std::unique_ptr ctx_; + ggml_tensor * token_ids_ = nullptr; + ggml_tensor * audio_embeddings_ = nullptr; + ggml_tensor * audio_positions_ = nullptr; + ggml_tensor * positions_ = nullptr; + ggml_tensor * logits_ = nullptr; + std::vector keys_; + std::vector values_; + std::vector position_ids_; + ggml_cgraph * graph_ = nullptr; + std::unique_ptr, GgmlGallocrDeleter> gallocr_; +}; + +class DecodeGraph { +public: + DecodeGraph(std::shared_ptr runtime, int64_t cache_steps, size_t graph_arena_bytes) + : runtime_(std::move(runtime)), + cache_steps_(cache_steps) { + if (cache_steps_ <= 0) { + throw std::runtime_error("Audio8 ASR thinker decode requires positive cache length"); + } + const auto build_start = Clock::now(); + ggml_init_params params{graph_arena_bytes, nullptr, true}; + ctx_.reset(ggml_init(params)); + if (ctx_ == nullptr) { + throw std::runtime_error("failed to initialize Audio8 ASR thinker decode graph context"); + } + const auto & config = runtime_->config(); + const auto & weights = runtime_->weights(); + core::ModuleBuildContext ctx{ctx_.get(), "audio8_asr.thinker.decode", runtime_->backend_type()}; + token_id_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_I32, 1); + auto token_id = core::wrap_tensor(token_id_, core::TensorShape::from_dims({1}), GGML_TYPE_I32); + auto x = modules::EmbeddingModule({config.vocab_size, config.hidden_size}) + .build(ctx, token_id, weights.token_embedding); + x = core::reshape_tensor(ctx, x, core::TensorShape::from_dims({1, 1, config.hidden_size})); + positions_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_I32, 1); + auto positions = core::wrap_tensor(positions_, core::TensorShape::from_dims({1}), GGML_TYPE_I32); + cache_slot_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_I32, 1); + auto cache_slot = core::wrap_tensor(cache_slot_, core::TensorShape::from_dims({1}), GGML_TYPE_I32); + attention_mask_ = ggml_new_tensor_4d(ctx_.get(), GGML_TYPE_F16, cache_steps_, 1, 1, 1); + auto attention_mask = core::wrap_tensor( + attention_mask_, + core::TensorShape::from_dims({1, 1, 1, cache_steps_}), + GGML_TYPE_F16); + graph_ = ggml_new_graph_custom(ctx_.get(), 65536, false); + auto decoder_out = modules::QwenCausalDecoderModule(make_qwen_decoder_config(config)) + .build_static_cache_tail( + ctx, + graph_, + x, + positions, + make_qwen_decoder_weights(weights), + cache_steps_, + attention_mask, + cache_slot); + step_cache_ = std::move(decoder_out.cache); + logits_ = decoder_out.logits.tensor; + ggml_set_output(logits_); + ggml_build_forward_expand(graph_, logits_); + buffer_ = ggml_backend_alloc_ctx_tensors(ctx_.get(), runtime_->backend()); + if (buffer_ == nullptr) { + engine::core::trim_backend_pools(runtime_->backend()); + buffer_ = ggml_backend_alloc_ctx_tensors(ctx_.get(), runtime_->backend()); + } + if (buffer_ == nullptr) { + throw std::runtime_error("failed to allocate Audio8 ASR thinker decode graph"); + } + attention_mask_values_.assign(static_cast(cache_steps_), ggml_fp32_to_fp16(-INFINITY)); + debug::timing_log_scalar("audio8_asr.thinker.decode.graph.build_ms", engine::debug::elapsed_ms(build_start, Clock::now())); + debug::trace_log_scalar("audio8_asr.thinker.decode_cache_steps", cache_steps_); + } + + ~DecodeGraph() { + engine::core::release_backend_graph_resources(runtime_->backend(), graph_, true); + if (buffer_ != nullptr) { + ggml_backend_buffer_free(buffer_); + } + } + + bool can_run(const ThinkerWeightsRuntime & runtime, int64_t required_steps) const { + return runtime_.get() == &runtime && cache_steps_ >= required_steps; + } + + void import_state(const runtime::TransformerKVState & state) { + step_cache_.import_state(state); + } + + std::vector run_step(int32_t token) { + const auto & config = runtime_->config(); + if (step_cache_.valid_steps() >= cache_steps_) { + throw std::runtime_error("Audio8 ASR thinker decode cache exhausted"); + } + ggml_backend_tensor_set(token_id_, &token, 0, sizeof(int32_t)); + const int32_t position = static_cast(step_cache_.current_end()); + ggml_backend_tensor_set(positions_, &position, 0, sizeof(int32_t)); + const int32_t cache_slot = static_cast(step_cache_.valid_steps()); + ggml_backend_tensor_set(cache_slot_, &cache_slot, 0, sizeof(int32_t)); + modules::write_qwen_cached_step_mask( + attention_mask_, + attention_mask_values_, + cache_steps_, + step_cache_.valid_steps(), + step_cache_.valid_steps()); + core::set_backend_threads(runtime_->backend(), runtime_->threads()); + const ggml_status status = engine::core::compute_backend_graph(runtime_->backend(), graph_); + ggml_backend_synchronize(runtime_->backend()); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("Audio8 ASR thinker decode graph compute failed"); + } + logits_buffer_.resize(static_cast(config.vocab_size)); + ggml_backend_tensor_get(logits_, logits_buffer_.data(), 0, logits_buffer_.size() * sizeof(float)); + step_cache_.advance_after_direct_append(1); + // The caller moves out of this buffer before the next step. + return std::move(logits_buffer_); + } + +private: + std::shared_ptr runtime_; + int64_t cache_steps_ = 0; + std::unique_ptr ctx_; + ggml_tensor * token_id_ = nullptr; + ggml_tensor * positions_ = nullptr; + ggml_tensor * cache_slot_ = nullptr; + ggml_tensor * attention_mask_ = nullptr; + ggml_tensor * logits_ = nullptr; + std::vector attention_mask_values_; + std::vector logits_buffer_; + runtime::TransformerKVCache step_cache_; + ggml_cgraph * graph_ = nullptr; + ggml_backend_buffer_t buffer_ = nullptr; +}; + +} // namespace + +struct Audio8ThinkerRuntime::Impl { + Impl( + std::shared_ptr weights_source, + Audio8ASRDecoderConfig config, + core::ExecutionContext & execution, + size_t prefill_graph_arena_bytes, + size_t decode_graph_arena_bytes, + size_t weight_context_bytes, + assets::TensorStorageType storage_type) + : weights(std::make_shared( + std::move(weights_source), + std::move(config), + execution, + weight_context_bytes, + storage_type)), + prefill_graph_arena_bytes(prefill_graph_arena_bytes), + decode_graph_arena_bytes(decode_graph_arena_bytes) {} + + void validate_prompt_audio( + const Audio8ASRPrompt & prompt, + const Audio8ASRAudioEmbeddings & audio_embeddings) const { + const auto & config = weights->config(); + if (audio_embeddings.hidden_size != config.hidden_size) { + throw std::runtime_error("Audio8 ASR audio embedding hidden size mismatch"); + } + if (audio_embeddings.tokens != static_cast(prompt.audio_token_positions.size())) { + throw std::runtime_error("Audio8 ASR audio embedding token count does not match prompt placeholders"); + } + if (static_cast(audio_embeddings.values.size()) != audio_embeddings.tokens * config.hidden_size) { + throw std::runtime_error("Audio8 ASR audio embedding value count mismatch"); + } + for (const int32_t position : prompt.audio_token_positions) { + if (position < 0 || position >= static_cast(prompt.input_ids.size())) { + throw std::runtime_error("Audio8 ASR audio placeholder position out of range"); + } + } + } + + std::shared_ptr weights; + size_t prefill_graph_arena_bytes = 0; + size_t decode_graph_arena_bytes = 0; + std::unique_ptr prefill_graph; + std::unique_ptr decode_graph; +}; + +Audio8ThinkerRuntime::Audio8ThinkerRuntime( + std::shared_ptr weights_source, + const Audio8ASRDecoderConfig & config, + core::ExecutionContext & execution, + size_t prefill_graph_arena_bytes, + size_t decode_graph_arena_bytes, + size_t weight_context_bytes, + assets::TensorStorageType weight_storage_type) + : impl_(std::make_unique( + std::move(weights_source), + config, + execution, + prefill_graph_arena_bytes, + decode_graph_arena_bytes, + weight_context_bytes, + weight_storage_type)) {} + +Audio8ThinkerRuntime::~Audio8ThinkerRuntime() = default; + +Audio8ASRGeneratedTokens Audio8ThinkerRuntime::generate( + const Audio8ASRPrompt & prompt, + const Audio8ASRAudioEmbeddings & audio_embeddings, + const Audio8ASRGenerationOptions & options) { + const auto & config = impl_->weights->config(); + if (prompt.input_ids.empty()) { + throw std::runtime_error("Audio8 ASR thinker prompt is empty"); + } + if (options.max_new_tokens <= 0) { + throw std::runtime_error("Audio8 ASR max_new_tokens must be positive"); + } + const int64_t prompt_steps = static_cast(prompt.input_ids.size()); + if (prompt_steps + options.max_new_tokens > config.max_position_embeddings) { + throw std::runtime_error("Audio8 ASR thinker request exceeds max_position_embeddings"); + } + impl_->validate_prompt_audio(prompt, audio_embeddings); + if (impl_->prefill_graph == nullptr || + !impl_->prefill_graph->matches(*impl_->weights, prompt_steps, audio_embeddings.tokens)) { + impl_->prefill_graph.reset(); + impl_->prefill_graph = std::make_unique( + impl_->weights, + prompt_steps, + audio_embeddings.tokens, + impl_->prefill_graph_arena_bytes); + } + auto prefill = impl_->prefill_graph->run( + prompt.input_ids, + audio_embeddings.values, + prompt.audio_token_positions); + const int64_t required_cache_steps = prompt_steps + options.max_new_tokens; + if (impl_->decode_graph == nullptr || + !impl_->decode_graph->can_run(*impl_->weights, required_cache_steps)) { + impl_->decode_graph.reset(); + impl_->decode_graph = std::make_unique( + impl_->weights, + required_cache_steps, + impl_->decode_graph_arena_bytes); + } + impl_->decode_graph->import_state(prefill.kv_state); + + Audio8ASRGeneratedTokens out; + std::vector logits = std::move(prefill.logits); + const auto decode_start = Clock::now(); + for (int64_t step = 0; step < options.max_new_tokens; ++step) { + const int32_t token = argmax_index(logits); + if (is_eos(config, token)) { + break; + } + out.token_ids.push_back(token); + logits = impl_->decode_graph->run_step(token); + } + debug::timing_log_scalar("audio8_asr.thinker.decode_total_ms", engine::debug::elapsed_ms(decode_start, Clock::now())); + return out; +} + +} // namespace engine::community_models::audio8_asr diff --git a/tests/audio8_asr/test_audio8_asr_golden_transcription.cpp b/tests/audio8_asr/test_audio8_asr_golden_transcription.cpp new file mode 100644 index 00000000..58f34fa5 --- /dev/null +++ b/tests/audio8_asr/test_audio8_asr_golden_transcription.cpp @@ -0,0 +1,143 @@ +#include "engine/framework/audio/wav_reader.h" +#include "engine/framework/core/backend.h" +#include "engine/framework/io/filesystem.h" +#include "engine/framework/io/text.h" +#include "engine/framework/runtime/model.h" +#include "engine/framework/runtime/registry.h" +#include "engine/framework/runtime/session.h" + +#include +#include +#include +#include +#include +#include +#include + +#ifndef ENGINE_REPO_ROOT +#define ENGINE_REPO_ROOT "." +#endif + +namespace { + +constexpr int kExitPass = 0; +constexpr int kExitFail = 1; +constexpr int kExitSkip = 125; + +const char * kExpectedText = + "This little work was finished in the year eighteen o three, and intended for immediate publication."; + +std::filesystem::path repo_path(const std::string & relative) { + return std::filesystem::path(ENGINE_REPO_ROOT) / relative; +} + +std::string arg_value(int argc, char ** argv, const std::string & name, const std::string & fallback) { + for (int i = 1; i + 1 < argc; ++i) { + if (argv[i] == name) { + return argv[i + 1]; + } + } + return fallback; +} + +std::string normalize_text(std::string text) { + std::string out; + out.reserve(text.size()); + for (char ch : text) { + if (std::isalnum(static_cast(ch)) || std::isspace(static_cast(ch))) { + out.push_back(static_cast(std::tolower(static_cast(ch)))); + } + } + return engine::io::trim_ascii_whitespace(std::move(out)); +} + +} // namespace + +int main(int argc, char ** argv) { + const std::filesystem::path model_path = arg_value( + argc, argv, "--model", + repo_path("models/Audio8-ASR-0.1B-GGUF/audio8-asr-0.1b-q8_0.gguf").string()); + const std::filesystem::path audio_path = arg_value( + argc, argv, "--audio", repo_path("assets/resources/a.wav").string()); + const std::string weight_type = arg_value(argc, argv, "--weight-type", ""); + + const bool model_available = engine::io::is_existing_file(model_path) || + engine::io::is_existing_file(model_path / "config.json"); + if (!model_available || !engine::io::is_existing_file(audio_path)) { + std::fprintf( + stderr, + "SKIP: test_audio8_asr_golden_transcription requires model weights at '%s' " + "and test audio at '%s'.\n", + model_path.string().c_str(), + audio_path.string().c_str()); + return kExitSkip; + } + + try { + auto registry = engine::runtime::make_default_registry(); + engine::runtime::ModelLoadRequest load_request; + load_request.model_path = model_path; + load_request.family_hint = "audio8_asr"; + auto model = registry.load(load_request); + + const engine::runtime::TaskSpec task{ + engine::runtime::VoiceTaskKind::Asr, + engine::runtime::RunMode::Offline, + }; + + engine::runtime::SessionOptions session_options; + if (!weight_type.empty()) { + session_options.options["audio8_asr.weight_type"] = weight_type; + } + + auto session = model->create_task_session(task, session_options); + auto * offline_session = dynamic_cast(session.get()); + if (!offline_session) { + std::cerr << "FAIL: session is not an IOfflineVoiceTaskSession\n"; + return kExitFail; + } + + const auto wav_data = engine::audio::read_wav_f32(audio_path); + engine::runtime::AudioBuffer audio; + audio.sample_rate = wav_data.sample_rate; + audio.channels = wav_data.channels; + audio.samples = wav_data.samples; + + const auto prep = engine::runtime::build_preparation_request(audio); + offline_session->prepare(prep); + + engine::runtime::TaskRequest request; + request.audio_input = audio; + const auto result = offline_session->run(request); + + if (!result.text_output.has_value()) { + std::cerr << "FAIL: Audio8 ASR produced no text output\n"; + return kExitFail; + } + + const std::string actual_raw = result.text_output->text; + const std::string actual = normalize_text(actual_raw); + const std::string expected = normalize_text(kExpectedText); + + std::cout << "Raw transcript: " << actual_raw << "\n"; + std::cout << "Normalized actual: " << actual << "\n"; + std::cout << "Normalized expected: " << expected << "\n"; + + // Exact raw match pins punctuation and casing, matching the fp32 + // reference verbatim; the normalized comparison localizes failures. + if (actual_raw != kExpectedText) { + std::cerr << "FAIL: raw transcript mismatch!\n"; + return kExitFail; + } + if (actual != expected) { + std::cerr << "FAIL: normalized transcript mismatch!\n"; + return kExitFail; + } + + std::cout << "PASS: Audio8 ASR golden transcription verified successfully.\n"; + return kExitPass; + } catch (const std::exception & error) { + std::cerr << "FAIL: " << error.what() << "\n"; + return kExitFail; + } +} diff --git a/tests/audio8_asr/test_audio8_asr_units.cpp b/tests/audio8_asr/test_audio8_asr_units.cpp new file mode 100644 index 00000000..4afafff0 --- /dev/null +++ b/tests/audio8_asr/test_audio8_asr_units.cpp @@ -0,0 +1,64 @@ +#include "engine/community_models/audio8_asr/types.h" +#include "test_assert.h" + +#include +#include +#include +#include + +namespace { + +namespace test = engine::test; +using engine::community_models::audio8_asr::audio8_asr_prompt_audio_token_count; +using engine::community_models::audio8_asr::audio8_asr_round_f32_to_bf16; + +void test_prompt_token_count_matches_reference() { + // Reference processor: downsampled = (mel_frames + 1) // 2; + // tokens = max(downsampled // merge_factor, 1). Values observed from the + // Audio8-ASR-0.1B reference runs on repo test audio. + test::require_eq(audio8_asr_prompt_audio_token_count(1407, 4), int64_t{176}, "1407 frames"); + test::require_eq(audio8_asr_prompt_audio_token_count(595, 4), int64_t{74}, "595 frames"); + test::require_eq(audio8_asr_prompt_audio_token_count(3000, 4), int64_t{375}, "3000 frames"); + test::require_eq(audio8_asr_prompt_audio_token_count(31, 4), int64_t{4}, "31 frames"); + test::require_eq(audio8_asr_prompt_audio_token_count(32, 4), int64_t{4}, "32 frames"); + // Short audio clamps to a single audio token. + test::require_eq(audio8_asr_prompt_audio_token_count(1, 4), int64_t{1}, "1 frame"); + test::require_eq(audio8_asr_prompt_audio_token_count(6, 4), int64_t{1}, "6 frames"); + test::require_eq(audio8_asr_prompt_audio_token_count(7, 4), int64_t{1}, "7 frames"); + bool threw = false; + try { + (void) audio8_asr_prompt_audio_token_count(0, 4); + } catch (const std::runtime_error &) { + threw = true; + } + test::require(threw, "zero mel frames must be rejected"); +} + +void test_bf16_rounding() { + // Exactly representable in bfloat16. + test::require_eq(audio8_asr_round_f32_to_bf16(1.0F), 1.0F, "1.0"); + test::require_eq(audio8_asr_round_f32_to_bf16(0.5F), 0.5F, "0.5"); + test::require_eq(audio8_asr_round_f32_to_bf16(-2.75F), -2.75F, "-2.75"); + // Round to nearest even: 0.3f (0x3E99999A) rounds up to 0x3E9A0000. + test::require_eq(audio8_asr_round_f32_to_bf16(0.3F), 0.30078125F, "0.3f"); + // Round-to-nearest-even tie: 1.0 + 2^-8 (bits 0x3F808000) rounds the + // trailing mantissa back down to 1.0. + const uint32_t tie_bits = 0x3F808000u; + float tie = 0.0F; + std::memcpy(&tie, &tie_bits, sizeof(tie)); + test::require_eq(audio8_asr_round_f32_to_bf16(tie), 1.0F, "rne tie"); +} + +} // namespace + +int main() { + try { + test_prompt_token_count_matches_reference(); + test_bf16_rounding(); + } catch (const std::exception & error) { + std::cerr << "FAIL: " << error.what() << "\n"; + return 1; + } + std::cout << "PASS: audio8_asr unit checks\n"; + return 0; +} diff --git a/tools/community_models/audio8_asr_reference.py b/tools/community_models/audio8_asr_reference.py new file mode 100644 index 00000000..2ba6d0ce --- /dev/null +++ b/tools/community_models/audio8_asr_reference.py @@ -0,0 +1,104 @@ +#!/usr/bin/env python3 +"""Run the Audio8-ASR-0.1B reference (HF transformers, trust_remote_code) on +repo test audio and dump golden parity artifacts. + +Outputs (under --outdir): + .prompt.json input_ids, audio token count/positions, prompt text + .logits.npy first generated step logits (float32) + .text.txt greedy transcript +""" + +import argparse +import json +from pathlib import Path + +import numpy as np +import soundfile as sf +import torch +from transformers import AutoModelForCausalLM, AutoProcessor + + +def load_audio_mono(path: str, target_sr: int = 16000) -> np.ndarray: + import librosa + + wav, sr = librosa.load(path, sr=target_sr, mono=True) + return np.asarray(wav, dtype=np.float32) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--model", type=Path, required=True) + parser.add_argument("--audio", action="append", required=True) + parser.add_argument("--outdir", type=Path, required=True) + parser.add_argument("--max-audio-seconds", type=float, default=30.0) + parser.add_argument("--max-new-tokens", type=int, default=128) + args = parser.parse_args() + + args.outdir.mkdir(parents=True, exist_ok=True) + device = "cpu" + processor = AutoProcessor.from_pretrained(args.model, trust_remote_code=True) + model = AutoModelForCausalLM.from_pretrained( + args.model, + trust_remote_code=True, + torch_dtype=torch.float32, + attn_implementation="eager", + ).to(device) + model.eval() + + for audio_path in args.audio: + stem = Path(audio_path).stem + audio = load_audio_mono(audio_path) + conversation = [ + { + "role": "user", + "content": [ + {"type": "audio", "path": str(Path(audio_path).resolve())}, + {"type": "text", "text": "Please transcribe this audio."}, + ], + } + ] + batch = processor.apply_chat_template( + conversation, + return_tensors="pt", + sampling_rate=16000, + audio_padding="longest", + add_generation_prompt=True, + audio_max_length=int(args.max_audio_seconds * 16000), + text_kwargs={"padding": "longest", "truncation": True, "max_length": 1000}, + ) + input_ids = batch["input_ids"][0].tolist() + audio_token_id = model.config.audio_token_id + audio_positions = [i for i, t in enumerate(input_ids) if t == audio_token_id] + with torch.inference_mode(): + output_ids = model.generate( + **batch, + max_new_tokens=args.max_new_tokens, + do_sample=False, + ) + prompt_len = int(batch["input_ids"].shape[1]) + # One forward pass to capture first-step logits over the prompt. + with torch.inference_mode(): + logits = model( + input_ids=batch["input_ids"], + input_features=batch.get("input_features"), + ).logits[0, -1, :] + text = processor.decode(output_ids[0, prompt_len:], skip_special_tokens=True).strip() + + prompt_record = { + "audio": str(audio_path), + "samples": int(len(audio)), + "input_ids": input_ids, + "prompt_len": prompt_len, + "audio_token_id": audio_token_id, + "audio_token_count": len(audio_positions), + "audio_positions_head": audio_positions[:5], + "generated_token_ids": output_ids[0, prompt_len:].tolist(), + } + (args.outdir / f"{stem}.prompt.json").write_text(json.dumps(prompt_record, indent=2)) + np.save(args.outdir / f"{stem}.logits.npy", logits.to(torch.float32).numpy()) + (args.outdir / f"{stem}.text.txt").write_text(text + "\n") + print(f"[{stem}] audio_tokens={len(audio_positions)} text={text!r}") + + +if __name__ == "__main__": + main() diff --git a/tools/community_models/audio8_asr_stages.py b/tools/community_models/audio8_asr_stages.py new file mode 100644 index 00000000..3ae33425 --- /dev/null +++ b/tools/community_models/audio8_asr_stages.py @@ -0,0 +1,88 @@ +#!/usr/bin/env python3 +"""Capture staged intermediate tensors from the Audio8-ASR-0.1B reference +pipeline for C++ side parity checks: mel -> encoder -> tower -> pool -> projector. +""" + +import argparse +import json +from pathlib import Path + +import numpy as np +import torch +from transformers import AutoModelForCausalLM, AutoProcessor + +import librosa + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--model", type=Path, required=True) + parser.add_argument("--audio", action="append", required=True) + parser.add_argument("--outdir", type=Path, required=True) + parser.add_argument("--max-audio-seconds", type=float, default=30.0) + args = parser.parse_args() + + args.outdir.mkdir(parents=True, exist_ok=True) + processor = AutoProcessor.from_pretrained(args.model, trust_remote_code=True) + model = AutoModelForCausalLM.from_pretrained( + args.model, trust_remote_code=True, torch_dtype=torch.float32, attn_implementation="eager" + ) + model.eval() + + for audio_path in args.audio: + stem = Path(audio_path).stem + wav, _ = librosa.load(audio_path, sr=16000, mono=True) + conv = [ + { + "role": "user", + "content": [ + {"type": "audio", "path": str(Path(audio_path).resolve())}, + {"type": "text", "text": "Please transcribe this audio."}, + ], + } + ] + batch = processor.apply_chat_template( + conv, + return_tensors="pt", + sampling_rate=16000, + audio_padding="longest", + add_generation_prompt=True, + audio_max_length=int(args.max_audio_seconds * 16000), + text_kwargs={"padding": "longest", "truncation": True, "max_length": 1000}, + ) + feats = batch["input_features"][0] # [mel, frames] + with torch.inference_mode(): + # Replicate modeling_arkasr._project_audio_row stage by stage. + features = feats + enc_in = features.to(next(model.audio_encoder.parameters()).dtype) + encoded = model.audio_encoder(enc_in, feature_lens=torch.tensor([features.shape[1]])) + hidden = encoded.last_hidden_state if hasattr(encoded, "last_hidden_state") else encoded + if isinstance(hidden, (tuple, list)): + hidden = hidden[0] + encoder_out = hidden.squeeze(0) # [T, 1024] + + tower_out = model.audio_mlp_tower(encoder_out) # [T, 1024] + input_ids = batch["input_ids"][0].tolist() + n_tokens = sum(1 for t in input_ids if t == model.config.audio_token_id) + pooled = tower_out + if pooled.shape[0] != n_tokens: + pooled = torch.nn.functional.adaptive_avg_pool1d( + pooled.transpose(0, 1).float().unsqueeze(0), output_size=n_tokens + ).squeeze(0).transpose(0, 1) + projected = model.audio_projector(pooled) # [N, 512] + + np.save(args.outdir / f"{stem}.mel.npy", features.to(torch.float32).numpy()) + np.save(args.outdir / f"{stem}.encoder_out.npy", encoder_out.numpy()) + np.save(args.outdir / f"{stem}.projected.npy", projected.numpy()) + info = { + "mel_shape": list(features.shape), + "encoder_tokens": int(encoder_out.shape[0]), + "audio_tokens": int(n_tokens), + "projected_shape": list(projected.shape), + } + (args.outdir / f"{stem}.stages.json").write_text(json.dumps(info, indent=2)) + print(stem, info) + + +if __name__ == "__main__": + main() diff --git a/tools/community_models/convert_audio8_asr.py b/tools/community_models/convert_audio8_asr.py new file mode 100644 index 00000000..e0f97917 --- /dev/null +++ b/tools/community_models/convert_audio8_asr.py @@ -0,0 +1,127 @@ +#!/usr/bin/env python3 +"""Convert Audio8/Audio8-ASR-0.1B safetensors checkpoints to audio.cpp GGUF packages. + +Produces a self-contained GGUF with weights from model.safetensors and embeds +the config/tokenizer sidecars, giving a complete model directory that +audiocpp_cli / audiocpp_server load with --family audio8_asr. + +The checkpoint is CC-BY-NC-4.0: run this against a checkpoint you downloaded +yourself from https://huggingface.co/Audio8/Audio8-ASR-0.1B and keep the +result local. Do not redistribute the converted GGUF. + +Examples: + python tools/community_models/convert_audio8_asr.py \ + --checkpoint models/Audio8-ASR-0.1B-hf \ + --converter build/debug/bin/audiocpp_gguf \ + --type q8_0 \ + --output models/Audio8-ASR-0.1B-GGUF/audio8-asr-0.1b-q8_0.gguf +""" + +import argparse +import shutil +import subprocess +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[2] +SPEC = REPO_ROOT / "model_specs" / "audio8_asr.json" + + +def convert( + converter: Path, + checkpoint: Path, + output: Path, + quant_type: str, + overwrite: bool, +) -> None: + ckpt = checkpoint / "model.safetensors" + if not ckpt.exists(): + candidates = sorted(checkpoint.glob("*.safetensors")) + if not candidates: + raise SystemExit(f"No .safetensors checkpoint found in {checkpoint}") + if len(candidates) > 1: + raise SystemExit( + f"Sharded safetensors checkpoints are not supported by this converter; " + f"found {len(candidates)} shards in {checkpoint}. Merge the shards into a " + f"single model.safetensors first (e.g. with safetensors.torch.save_file after " + f"concatenating the shard state dicts)." + ) + ckpt = candidates[0] + + output.parent.mkdir(parents=True, exist_ok=True) + + command = [ + str(converter), + "--input", + str(ckpt), + "--root", + str(checkpoint), + "--family", + "audio8_asr", + "--model-spec", + str(SPEC), + "--type", + quant_type, + "--output", + str(output), + ] + if overwrite: + command.append("--overwrite") + print("+", " ".join(command)) + subprocess.run(command, check=True) + + for asset_name in ["config.json", "generation_config.json", "preprocessor_config.json", "tokenizer_config.json", "tokenizer.json", "vocab.json", "merges.txt"]: + src_asset = checkpoint / asset_name + dst_asset = output.parent / asset_name + if not src_asset.exists() or src_asset.resolve() == dst_asset.resolve(): + continue + shutil.copyfile(src_asset, dst_asset) + print(f"copied {asset_name} -> {output.parent}") + + +def main() -> None: + parser = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + parser.add_argument( + "--checkpoint", + type=Path, + default=REPO_ROOT / "models" / "Audio8-ASR-0.1B-hf", + help="checkpoint directory containing model.safetensors, configs, and tokenizer files", + ) + parser.add_argument( + "--converter", + type=Path, + required=True, + help="path to the audiocpp_gguf binary", + ) + parser.add_argument( + "--output", + type=Path, + default=Path("models/Audio8-ASR-0.1B-GGUF/audio8-asr-0.1b-q8_0.gguf"), + help="output .gguf file path", + ) + parser.add_argument( + "--type", + default="q8_0", + choices=["orig", "f32", "f16", "bf16", "q8_0", "q4_k", "q5_k", "q6_k"], + help="quantization type for GGUF tensors", + ) + parser.add_argument( + "--overwrite", + action="store_true", + help="overwrite existing output file", + ) + + args = parser.parse_args() + convert( + converter=args.converter, + checkpoint=args.checkpoint, + output=args.output, + quant_type=args.type, + overwrite=args.overwrite, + ) + + +if __name__ == "__main__": + main()