diff --git a/.gitignore b/.gitignore index 4c8fae3a8..7a1caeaa1 100644 --- a/.gitignore +++ b/.gitignore @@ -19,6 +19,7 @@ __pycache__/ /tests/omnivoice/outputs/ /unittests/ /models/ +/granite5asr/ /reference/ /resources/ /patches/ diff --git a/CMakeLists.txt b/CMakeLists.txt index b50733ac5..e79b1085a 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1271,6 +1271,22 @@ audiocpp_add_model(citrinet_asr engine::models::citrinet_asr::make_citrinet_asr_loader ) +audiocpp_add_model(granite5asr + SOURCES + src/community_models/granite5asr/assets.cpp + src/community_models/granite5asr/frontend.cpp + src/community_models/granite5asr/encoder.cpp + src/community_models/granite5asr/session.cpp + INCLUDES + engine/community_models/granite5asr/session.h + LOADERS + engine::community_models::granite5asr::make_granite5asr_loader + ALIASES + granite_speech5_asr + granite_speech + granite_speech5_ctc +) + audiocpp_add_model(vevo2 SOURCES src/models/vevo2/ar.cpp @@ -1801,6 +1817,31 @@ add_executable(miocodec_wavlm_parity EXCLUDE_FROM_ALL target_link_libraries(miocodec_wavlm_parity PRIVATE engine_runtime ggml) +add_executable(test_granite5asr_golden_transcription + tests/granite5asr/test_granite5asr_golden_transcription.cpp +) +target_compile_definitions(test_granite5asr_golden_transcription PRIVATE + ENGINE_REPO_ROOT="${CMAKE_CURRENT_SOURCE_DIR}" +) +target_link_libraries(test_granite5asr_golden_transcription PRIVATE engine_runtime ggml) +target_include_directories(test_granite5asr_golden_transcription PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}) +if (ENGINE_ENABLE_OPENMP) + target_link_libraries(test_granite5asr_golden_transcription PRIVATE OpenMP::OpenMP_CXX) +endif() + +add_executable(granite5asr_warm_bench + tests/granite5asr/granite5asr_warm_bench.cpp +) +target_compile_definitions(granite5asr_warm_bench PRIVATE + ENGINE_REPO_ROOT="${CMAKE_CURRENT_SOURCE_DIR}" +) +target_link_libraries(granite5asr_warm_bench PRIVATE engine_runtime ggml) +target_include_directories(granite5asr_warm_bench PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}) +if (ENGINE_ENABLE_OPENMP) + target_link_libraries(granite5asr_warm_bench PRIVATE OpenMP::OpenMP_CXX) +endif() + + # F5/Habibi tests: parity harnesses + e2e sample generator. Only when the # f5_tts model is linked (they call model-internal symbols) AND test hooks # are explicitly enabled — F5_MEL_TEST must never leak into production @@ -1930,6 +1971,7 @@ if (ENGINE_BUILD_WARMBENCH) add_engine_warmbench(confucius4_tts_warm_bench tests/confucius4_tts/confucius4_tts_warm_bench.cpp) add_engine_warmbench(controlfoley_warm_bench tests/controlfoley/controlfoley_warm_bench.cpp) add_engine_warmbench(fun_asr_nano_warm_bench tests/fun_asr_nano/fun_asr_nano_warm_bench.cpp) + add_engine_warmbench(granite5asr_warm_bench tests/granite5asr/granite5asr_warm_bench.cpp) add_engine_warmbench(higgs_audio_stt_warm_bench tests/higgs_audio_stt/higgs_audio_stt_warm_bench.cpp) add_engine_warmbench(higgs_audio_tts_warm_bench tests/higgs_audio_tts/higgs_audio_tts_warm_bench.cpp) add_engine_warmbench(hviske_asr_warm_bench tests/hviske_asr/hviske_asr_warm_bench.cpp) diff --git a/docs/community_models/granite5asr.md b/docs/community_models/granite5asr.md new file mode 100644 index 000000000..ce1f524ff --- /dev/null +++ b/docs/community_models/granite5asr.md @@ -0,0 +1,38 @@ +# IBM Granite Speech 5.0 470M TurboCTC in audio.cpp + +IBM Granite Speech 5.0 TurboCTC is a compact 470-million-parameter English Automatic Speech Recognition (ASR) model delivering state-of-the-art transcription accuracy with ultra-low latency. + +## Architecture + +- **Audio Frontend**: 16 kHz, 80-bin HTK mel spectrogram with 8.0 dB dynamic flooring, first-order deltas, and 2x frame stacking (320-dim features). +- **Acoustic Conformer**: 16 layers ($d_{model}=1024, d_{ff}=4096, d_{conv}=2048, heads=8, head\_dim=128$) with block self-attention ($context\_size=128$), Shaw relative positional embeddings, depthwise convolution with folded batch-norm, and mid-layer CTC self-conditioning at layer 8. +- **Decoder**: Non-autoregressive Connectionist Temporal Classification (CTC) with greedy decoding and 16,384 BPE vocabulary via fast HuggingFace tokenizer. + +## CLI Usage + +### Offline Transcription + +```bash +# Transcribe audio using native safetensors directory +audiocpp_cli asr --model granite5asr --input sample.wav + +# Or explicitly pass the family +audiocpp_cli asr --family granite5asr --model path/to/checkpoint --input sample.wav +``` + +### Quantized GGUF Loading + +```bash +# Transcribe using a converted Q8_0 GGUF package +audiocpp_cli asr --model models/granite-speech-5.0-470m-turboctc-gguf --input sample.wav +``` + +### Audio Chunking & VAD + +```bash +# Long audio with automatic VAD segmentation +audiocpp_cli asr --model granite5asr --input long_audio.wav --option audio_chunk_mode=auto + +# Fixed 30-second chunking +audiocpp_cli asr --model granite5asr --input long_audio.wav --option audio_chunk_mode=fixed --option audio_chunk_duration_sec=30 +``` diff --git a/docs/community_models/models.md b/docs/community_models/models.md index 47e318db4..f6cad045c 100644 --- a/docs/community_models/models.md +++ b/docs/community_models/models.md @@ -18,6 +18,7 @@ Practical expectations: |---|---|---|---|---| | **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 | +| **granite5asr** | ASR | en | Community | [IBM Granite Speech 5.0 470M TurboCTC](granite5asr.md) ultra-fast Conformer-CTC ASR with Shaw relative positional embeddings and ByteLevel BPE | | **inflect_v2** | TTS | en | Community | [Inflect Micro v2 and Nano v2](inflect_v2.md) native FP32 offline synthesis | | **kroko_asr** | ASR | de, en, es, fr, it, he, nl, pt, sv, tr | Mirek [@mirek190](https://github.com/mirek190) | [Kroko Community ASR](kroko_asr.md) native offline/streaming Zipformer2/RNN-T transcription with word timestamps | | **mms_forced_aligner** | Align | nl (nld), en (eng); pre-romanized Latin | Community | [MMS-300M-1130 Forced Aligner](mms_forced_aligner.md) word-timestamp alignment from a wav2vec2 CTC checkpoint (safetensors or local GGUF) | diff --git a/include/engine/community_models/granite5asr/assets.h b/include/engine/community_models/granite5asr/assets.h new file mode 100644 index 000000000..b0943c428 --- /dev/null +++ b/include/engine/community_models/granite5asr/assets.h @@ -0,0 +1,66 @@ +#pragma once + +#include "engine/framework/assets/tensor_source.h" +#include "engine/framework/model_spec/package.h" +#include "engine/framework/tokenizers/hf_tokenizer_json.h" + +#include +#include +#include +#include +#include + +namespace engine::community_models::granite5asr { + +struct Granite5FrontendConfig { + int64_t sample_rate = 16000; + int64_t n_fft = 512; + int64_t win_length = 400; + int64_t hop_length = 160; + int64_t n_mels = 80; + int64_t stack_factor = 2; + bool deltas = true; + int64_t delta_win_length = 3; + float logmel_floor_db = 8.0f; +}; + +struct Granite5EncoderConfig { + int64_t hidden_size = 1024; + int64_t intermediate_size = 4096; + int64_t num_layers = 16; + int64_t num_attention_heads = 8; + int64_t num_key_value_heads = 8; + int64_t head_dim = 128; + int64_t context_size = 128; + int64_t conv_kernel_size = 7; + int64_t conv_expansion_factor = 2; + int64_t max_position_embeddings = 512; + int64_t num_mel_bins = 80; + int64_t input_features = 320; + int64_t vocab_size = 16384; + std::vector subsample_layers = {0, 1}; +}; + +struct Granite5ASRConfig { + std::string model_type = "granite_speech5_ctc"; + int64_t vocab_size = 16384; + int64_t pad_token_id = 0; + int64_t blank_token_id = 0; + Granite5FrontendConfig frontend; + Granite5EncoderConfig encoder; +}; + +struct Granite5ASRAssets { + assets::ResourceBundle resources; + std::shared_ptr source; + Granite5ASRConfig config; + std::shared_ptr tokenizer; + std::vector special_token_ids; +}; + +Granite5ASRConfig parse_granite5asr_config(const std::string & json_text); + +std::shared_ptr load_granite5asr_assets( + const std::filesystem::path & model_path); + +} // namespace engine::community_models::granite5asr diff --git a/include/engine/community_models/granite5asr/encoder.h b/include/engine/community_models/granite5asr/encoder.h new file mode 100644 index 000000000..141639a64 --- /dev/null +++ b/include/engine/community_models/granite5asr/encoder.h @@ -0,0 +1,72 @@ +#pragma once + +#include "engine/framework/core/backend.h" +#include "engine/framework/core/backend_weight_store.h" +#include "engine/framework/core/execution_context.h" +#include "engine/framework/modules/linear_module.h" +#include "engine/framework/modules/norm_modules.h" +#include "engine/community_models/granite5asr/assets.h" +#include "engine/community_models/granite5asr/frontend.h" + +#include +#include +#include + +namespace engine::community_models::granite5asr { + +struct Granite5LayerWeights { + modules::NormWeights ffn1_norm; + modules::LinearWeights ffn1_fc1; + modules::LinearWeights ffn1_fc2; + + modules::NormWeights norm_self_att; + core::TensorValue q_proj; + core::TensorValue k_proj; + core::TensorValue v_proj; + modules::LinearWeights o_proj; + core::TensorValue rel_pos_emb; + + modules::NormWeights norm_conv; + modules::LinearWeights conv_pw1; + core::TensorValue conv_dw_weight; + core::TensorValue conv_dw_bias; + modules::LinearWeights conv_pw2; + + modules::NormWeights ffn2_norm; + modules::LinearWeights ffn2_fc1; + modules::LinearWeights ffn2_fc2; + + modules::NormWeights norm_out; + + bool is_subsample = false; +}; + +struct Granite5EncoderWeights { + modules::LinearWeights input_linear; + std::vector layers; + modules::LinearWeights out; + modules::LinearWeights out_mid; +}; + +class Granite5EncoderRuntime { +public: + Granite5EncoderRuntime( + std::shared_ptr assets, + engine::core::ExecutionContext & execution_context, + assets::TensorStorageType storage_type, + size_t graph_arena_bytes = 1024ull * 1024ull * 1024ull); + + std::vector transcribe_features( + const Granite5FrontendFeatures & features); + + const Granite5ASRAssets & assets() const noexcept { return *assets_; } + +private: + std::shared_ptr assets_; + engine::core::ExecutionContext * execution_context_ = nullptr; + engine::core::BackendWeightStore weight_store_; + Granite5EncoderWeights weights_; + size_t graph_arena_bytes_; +}; + +} // namespace engine::community_models::granite5asr diff --git a/include/engine/community_models/granite5asr/frontend.h b/include/engine/community_models/granite5asr/frontend.h new file mode 100644 index 000000000..ec7d7777b --- /dev/null +++ b/include/engine/community_models/granite5asr/frontend.h @@ -0,0 +1,34 @@ +#pragma once + +#include "engine/framework/audio/dsp.h" +#include "engine/framework/runtime/session.h" +#include "engine/community_models/granite5asr/assets.h" + +#include +#include + +namespace engine::community_models::granite5asr { + +struct Granite5FrontendFeatures { + std::vector values; + int64_t frames = 0; + int64_t feature_dim = 320; +}; + +class Granite5Frontend { +public: + explicit Granite5Frontend(std::shared_ptr assets); + + Granite5FrontendFeatures extract(const runtime::AudioBuffer & audio) const; + Granite5FrontendFeatures extract_waveform(const std::vector & waveform) const; + std::vector prepare_waveform(const runtime::AudioBuffer & audio) const; + + const Granite5FrontendConfig & config() const noexcept { return assets_->config.frontend; } + +private: + std::shared_ptr assets_; + audio::AudioTensor mel_filterbank_; + std::vector window_; +}; + +} // namespace engine::community_models::granite5asr diff --git a/include/engine/community_models/granite5asr/session.h b/include/engine/community_models/granite5asr/session.h new file mode 100644 index 000000000..cca058ad5 --- /dev/null +++ b/include/engine/community_models/granite5asr/session.h @@ -0,0 +1,111 @@ +#pragma once + +#include "engine/framework/audio/chunking.h" +#include "engine/framework/runtime/model.h" +#include "engine/framework/runtime/session_base.h" +#include "engine/community_models/granite5asr/assets.h" +#include "engine/community_models/granite5asr/encoder.h" +#include "engine/community_models/granite5asr/frontend.h" + +#include +#include +#include + +namespace engine::community_models::granite5asr { + +class Granite5ASRSessionBase : public runtime::RuntimeSessionBase { +public: + Granite5ASRSessionBase( + runtime::TaskSpec task, + runtime::SessionOptions options, + std::shared_ptr assets); + ~Granite5ASRSessionBase() override; + +protected: + std::string family_impl() const; + runtime::VoiceTaskKind task_kind_impl() const; + runtime::RunMode run_mode_impl() const; + + runtime::Transcript transcribe_audio( + const runtime::AudioBuffer & audio, + const std::unordered_map & options); + + runtime::IOfflineVoiceTaskSession & vad_session(); + + runtime::TaskSpec task_; + std::shared_ptr assets_; + Granite5Frontend frontend_; + std::unique_ptr encoder_; + std::string vad_model_path_; + std::unique_ptr vad_model_; + std::unique_ptr vad_session_; +}; + +class Granite5ASROfflineSession final + : public Granite5ASRSessionBase + , public runtime::IOfflineVoiceTaskSession { +public: + Granite5ASROfflineSession( + runtime::TaskSpec task, + runtime::SessionOptions options, + std::shared_ptr assets); + + 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; +}; + +class Granite5ASRStreamingSession final + : public Granite5ASRSessionBase + , public runtime::IStreamingVoiceTaskSession { +public: + Granite5ASRStreamingSession( + runtime::TaskSpec task, + runtime::SessionOptions options, + std::shared_ptr assets); + + 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::StreamingPolicy streaming_policy() const override; + void start_stream(const runtime::TaskRequest & request) override; + void set_stream_event_sink(runtime::StreamEventCallback sink) override; + void reset() override; + runtime::StreamEvent process_audio_chunk(const runtime::AudioChunk & chunk) override; + runtime::TaskResult finish_stream() override; + runtime::TaskResult finalize() override; + +private: + runtime::StreamEventCallback stream_event_sink_; + runtime::AudioBuffer streaming_audio_; + runtime::TaskRequest streaming_request_; +}; + +class Granite5ASRLoadedModel final : public runtime::ILoadedVoiceModel { +public: + Granite5ASRLoadedModel( + 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_granite5asr_model( + const runtime::ModelLoadRequest & request); + +std::shared_ptr make_granite5asr_loader(); + +} // namespace engine::community_models::granite5asr diff --git a/include/engine/framework/tokenizers/hf_tokenizer_json.h b/include/engine/framework/tokenizers/hf_tokenizer_json.h index afb24664a..0c60ba76d 100644 --- a/include/engine/framework/tokenizers/hf_tokenizer_json.h +++ b/include/engine/framework/tokenizers/hf_tokenizer_json.h @@ -10,7 +10,11 @@ namespace engine::tokenizers { class HuggingFaceTokenizerJson { public: - HuggingFaceTokenizerJson(std::vector id_to_token, std::string metaspace_replacement, bool trim_leading_space); + HuggingFaceTokenizerJson( + std::vector id_to_token, + std::string metaspace_replacement, + bool trim_leading_space, + bool byte_level = false); const std::vector & id_to_token() const noexcept; std::string decode_ids(const std::vector & ids) const; @@ -20,6 +24,7 @@ class HuggingFaceTokenizerJson { std::vector id_to_token_; std::string metaspace_replacement_; bool trim_leading_space_ = false; + bool byte_level_ = false; }; std::shared_ptr load_huggingface_tokenizer_json( diff --git a/model_specs/granite5asr.json b/model_specs/granite5asr.json new file mode 100644 index 000000000..9427da555 --- /dev/null +++ b/model_specs/granite5asr.json @@ -0,0 +1,172 @@ +{ + "family": "granite5asr", + "display_name": "Granite Speech 5.0 470M TurboCTC", + "description": "IBM Granite Speech 5.0 TurboCTC 470M parameter compact and ultra-fast English ASR model using a Conformer CTC architecture with block self-attention and self-conditioning.", + "category": "asr", + "status": "supported", + "tasks": [ + "asr" + ], + "modes": [ + "offline", + "streaming" + ], + "languages": [ + "en" + ], + "capabilities": { + "asr": [ + "vad_chunking", + "partial_results" + ] + }, + "options": { + "request": [ + { + "name": "language", + "type": "string", + "description": "Recognition language (currently English).", + "required": false, + "default": "en" + }, + { + "name": "audio_chunk_mode", + "type": "enum", + "description": "Audio chunking mode: auto, fixed, or none.", + "values": [ + "auto", + "fixed", + "none" + ], + "required": false, + "default": "auto" + }, + { + "name": "audio_chunk_duration_sec", + "type": "float", + "description": "Fixed chunk duration in seconds when not using VAD segmentation.", + "required": false, + "min": 0.001, + "default": 30 + } + ], + "session": [ + { + "name": "weight_type", + "type": "enum", + "description": "Shared model weight storage type.", + "preset": "weight_type_full", + "required": false, + "default": "native" + }, + { + "name": "encoder_graph_arena_mb", + "type": "int", + "description": "Encoder graph arena size in MB.", + "required": false, + "min": 64, + "default": 1024 + }, + { + "name": "vad_model_path", + "type": "string", + "description": "Path to the Silero VAD model directory used by automatic audio chunking.", + "required": false, + "default": "assets/framework/models/silero_vad" + } + ], + "load": [] + }, + "runtime": { + "tags": [ + "gguf", + "server", + "stream", + "cuda", + "metal", + "cpu" + ] + }, + "ui": { + "recommended_package": "granite5asr_q8_0", + "tags": [ + "ASR", + "GGUF", + "Stream" + ], + "docs": [ + "docs/asr.md", + "docs/gguf.md" + ], + "summary": "IBM Granite Speech 5.0 470M TurboCTC English transcription." + }, + "package_defaults": { + "download": { + "kind": "huggingface_snapshot", + "repo": "ibm-granite/granite-speech-5.0-470m-turboctc", + "revision": "main", + "gated": false + } + }, + "packages": [ + { + "id": "granite5asr_q8_0", + "display_name": "Granite Speech 5.0 470M TurboCTC Q8_0 GGUF", + "default": true, + "format": "gguf", + "precision": "q8_0", + "target_directory": "granite-speech-5.0-470m-turboctc-gguf", + "files": [ + "granite-speech-5.0-470m-turboctc-q8_0.gguf" + ] + }, + { + "id": "granite5asr_safetensors", + "display_name": "Granite Speech 5.0 470M TurboCTC Native Safetensors", + "format": "safetensors", + "precision": "bfloat16", + "target_directory": "granite5asr", + "files": [ + "config.json", + "preprocessor_config.json", + "tokenizer.json", + "model.safetensors" + ] + } + ], + "sources": [ + { + "format": "gguf", + "roots": { + "model": ".", + "weights": "$gguf" + }, + "files": { + "config": "model:config.json", + "tokenizer_json": "model:tokenizer.json" + }, + "optional_files": { + "preprocessor_config": "model:preprocessor_config.json" + }, + "tensors": { + "weights": "weights:" + } + }, + { + "format": "safetensors", + "roots": { + "model": "." + }, + "files": { + "config": "model:config.json", + "tokenizer_json": "model:tokenizer.json" + }, + "optional_files": { + "preprocessor_config": "model:preprocessor_config.json" + }, + "tensors": { + "weights": "model:model.safetensors" + } + } + ] +} diff --git a/src/community_models/granite5asr/assets.cpp b/src/community_models/granite5asr/assets.cpp new file mode 100644 index 000000000..337f9f34f --- /dev/null +++ b/src/community_models/granite5asr/assets.cpp @@ -0,0 +1,118 @@ +#include "engine/community_models/granite5asr/assets.h" + +#include "engine/framework/io/json.h" +#include "engine/framework/model_spec/package.h" + +#include +#include +#include + +namespace engine::community_models::granite5asr { +namespace json = engine::io::json; +namespace { + +void validate_config(const Granite5ASRConfig & config) { + if (config.vocab_size <= 0) { + throw std::runtime_error("Granite 5 ASR invalid vocab size"); + } + if (config.encoder.hidden_size <= 0 || config.encoder.num_layers <= 0 || + config.encoder.num_attention_heads <= 0 || + config.encoder.hidden_size % config.encoder.num_attention_heads != 0) { + throw std::runtime_error("Granite 5 ASR invalid encoder metadata"); + } + if (config.frontend.sample_rate != 16000 || config.frontend.n_mels <= 0 || + config.frontend.n_fft <= 0 || config.frontend.win_length <= 0 || + config.frontend.hop_length <= 0) { + throw std::runtime_error("Granite 5 ASR invalid frontend metadata"); + } +} + +std::vector parse_special_token_ids( + const std::filesystem::path & tokenizer_json, + int64_t vocab_size) { + std::vector special(static_cast(vocab_size), 0); + const auto root = json::parse_file(tokenizer_json); + if (const auto * added = root.find("added_tokens"); added != nullptr && added->is_array()) { + for (const auto & item : added->as_array()) { + if (!json::optional_bool(item, "special", false)) { + continue; + } + const int64_t id = json::require_i64(item, "id"); + if (id >= 0 && id < vocab_size) { + special[static_cast(id)] = 1; + } + } + } + return special; +} + +Granite5ASRConfig parse_config(const assets::ResourceBundle & resources) { + const auto config_root = resources.parse_json("config"); + + Granite5ASRConfig config; + config.model_type = json::optional_string(config_root, "model_type", "granite_speech5_ctc"); + config.vocab_size = json::optional_i64(config_root, "vocab_size", 16384); + config.pad_token_id = json::optional_i64(config_root, "pad_token_id", 0); + config.blank_token_id = json::optional_i64(config_root, "pad_token_id", 0); + + if (const auto * enc = config_root.find("encoder_config"); enc != nullptr && enc->is_object()) { + config.encoder.hidden_size = json::optional_i64(*enc, "hidden_size", 1024); + config.encoder.intermediate_size = json::optional_i64(*enc, "intermediate_size", 4096); + config.encoder.num_layers = json::optional_i64(*enc, "num_hidden_layers", 16); + config.encoder.num_attention_heads = json::optional_i64(*enc, "num_attention_heads", 8); + config.encoder.num_key_value_heads = json::optional_i64(*enc, "num_key_value_heads", 8); + config.encoder.head_dim = json::optional_i64(*enc, "head_dim", 128); + config.encoder.context_size = json::optional_i64(*enc, "context_size", 128); + config.encoder.conv_kernel_size = json::optional_i64(*enc, "conv_kernel_size", 7); + config.encoder.conv_expansion_factor = json::optional_i64(*enc, "conv_expansion_factor", 2); + config.encoder.max_position_embeddings = json::optional_i64(*enc, "max_position_embeddings", 512); + config.encoder.num_mel_bins = json::optional_i64(*enc, "num_mel_bins", 80); + config.encoder.vocab_size = json::optional_i64(*enc, "vocab_size", config.vocab_size); + + if (const auto * subsample = enc->find("subsample_layers"); subsample != nullptr && subsample->is_array()) { + config.encoder.subsample_layers.clear(); + for (const auto & val : subsample->as_array()) { + if (val.is_number()) { + config.encoder.subsample_layers.push_back(val.as_i64()); + } + } + } + } + + if (resources.has_file("preprocessor_config")) { + const auto preproc_root = resources.parse_json("preprocessor_config"); + config.frontend.sample_rate = json::optional_i64(preproc_root, "sample_rate", 16000); + config.frontend.n_fft = json::optional_i64(preproc_root, "n_fft", 512); + config.frontend.win_length = json::optional_i64(preproc_root, "win_length", 400); + config.frontend.hop_length = json::optional_i64(preproc_root, "hop_length", 160); + config.frontend.n_mels = json::optional_i64(preproc_root, "n_mels", 80); + config.frontend.stack_factor = json::optional_i64(preproc_root, "stack_factor", 2); + config.frontend.deltas = json::optional_bool(preproc_root, "deltas", true); + config.frontend.delta_win_length = json::optional_i64(preproc_root, "delta_win_length", 3); + config.frontend.logmel_floor_db = json::optional_f32(preproc_root, "logmel_floor_db", 8.0f); + } + + config.encoder.input_features = config.frontend.n_mels * (config.frontend.deltas ? 2 : 1) * config.frontend.stack_factor; + validate_config(config); + return config; +} + +} // namespace + +std::shared_ptr load_granite5asr_assets( + const std::filesystem::path & model_path) { + auto resources = engine::model_spec::load_resource_bundle( + model_path, + engine::model_spec::default_spec_path("granite5asr")); + auto assets = std::make_shared(); + assets->resources = std::move(resources); + assets->source = assets->resources.open_tensor_source("weights"); + assets->config = parse_config(assets->resources); + const auto & tokenizer_json = assets->resources.require_file("tokenizer_json"); + assets->tokenizer = engine::tokenizers::load_huggingface_tokenizer_json(tokenizer_json); + assets->special_token_ids = + parse_special_token_ids(tokenizer_json, assets->config.vocab_size); + return assets; +} + +} // namespace engine::community_models::granite5asr diff --git a/src/community_models/granite5asr/encoder.cpp b/src/community_models/granite5asr/encoder.cpp new file mode 100644 index 000000000..7fd90adb3 --- /dev/null +++ b/src/community_models/granite5asr/encoder.cpp @@ -0,0 +1,561 @@ +#include "engine/community_models/granite5asr/encoder.h" + +#include "engine/framework/core/backend.h" +#include "engine/framework/core/backend_weight_store.h" +#include "engine/framework/core/execution_context.h" +#include "engine/framework/debug/profiler.h" +#include "engine/framework/modules/activation_modules.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/modules/streaming_conv_modules.h" +#include "engine/framework/modules/structural_modules.h" + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace engine::community_models::granite5asr { +namespace { + +constexpr size_t kEncoderGraphNodes = 1048576; +constexpr float kLayerNormEpsilon = 1.0e-5f; + +modules::LinearWeights load_linear_with_bias( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & prefix, + int64_t in_features, + int64_t out_features, + assets::TensorStorageType storage_type) { + return { + store.load_tensor(source, prefix + ".weight", storage_type, {out_features, in_features}), + store.load_f32_tensor(source, prefix + ".bias", {out_features}), + }; +} + +modules::LinearWeights load_linear_no_bias( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & prefix, + int64_t in_features, + int64_t out_features, + assets::TensorStorageType storage_type) { + return { + store.load_tensor(source, prefix + ".weight", storage_type, {out_features, in_features}), + std::nullopt, + }; +} + +modules::NormWeights load_layer_norm( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & prefix, + int64_t hidden_size) { + return { + store.load_f32_tensor(source, prefix + ".weight", {hidden_size}), + store.load_f32_tensor(source, prefix + ".bias", {hidden_size}), + }; +} + +void load_folded_depthwise_conv( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & prefix, + int64_t channels, + int64_t kernel_size, + assets::TensorStorageType storage_type, + core::TensorValue & out_weight, + core::TensorValue & out_bias) { + const auto raw_w = source.require_f32(prefix + ".depthwise_conv.weight", {channels, 1, kernel_size}); + const auto gamma = source.require_f32(prefix + ".norm.weight", {channels}); + const auto beta = source.require_f32(prefix + ".norm.bias", {channels}); + const auto mean = source.require_f32(prefix + ".norm.running_mean", {channels}); + const auto var = source.require_f32(prefix + ".norm.running_var", {channels}); + + std::vector folded_w(static_cast(channels * kernel_size), 0.0f); + std::vector folded_b(static_cast(channels), 0.0f); + + for (int64_t c = 0; c < channels; ++c) { + const float scale = gamma[static_cast(c)] / std::sqrt(var[static_cast(c)] + 1e-5f); + folded_b[static_cast(c)] = beta[static_cast(c)] - mean[static_cast(c)] * scale; + for (int64_t k = 0; k < kernel_size; ++k) { + folded_w[static_cast(c * kernel_size + k)] = + raw_w[static_cast(c * kernel_size + k)] * scale; + } + } + + out_weight = store.make_from_f32( + core::TensorShape::from_dims({channels, 1, kernel_size}), + storage_type == assets::TensorStorageType::Native ? assets::TensorStorageType::F32 : storage_type, + std::move(folded_w)); + out_bias = store.make_from_f32( + core::TensorShape::from_dims({channels}), + assets::TensorStorageType::F32, + std::move(folded_b)); +} + +Granite5LayerWeights load_layer_weights( + core::BackendWeightStore & store, + const assets::TensorSource & source, + int64_t layer_idx, + const Granite5EncoderConfig & config, + assets::TensorStorageType storage_type) { + Granite5LayerWeights w; + const std::string pfx = "encoder.layers." + std::to_string(layer_idx) + "."; + + w.is_subsample = (std::find(config.subsample_layers.begin(), config.subsample_layers.end(), layer_idx) != config.subsample_layers.end()); + + // FFN1 + w.ffn1_norm = load_layer_norm(store, source, pfx + "norm_feed_forward1", config.hidden_size); + w.ffn1_fc1 = load_linear_with_bias(store, source, pfx + "feed_forward1.linear1", config.hidden_size, config.intermediate_size, storage_type); + w.ffn1_fc2 = load_linear_with_bias(store, source, pfx + "feed_forward1.linear2", config.intermediate_size, config.hidden_size, storage_type); + + // Self-Attention + w.norm_self_att = load_layer_norm(store, source, pfx + "norm_self_att", config.hidden_size); + w.q_proj = store.load_tensor(source, pfx + "self_attn.q_proj.weight", storage_type, {config.hidden_size, config.hidden_size}); + w.k_proj = store.load_tensor(source, pfx + "self_attn.k_proj.weight", storage_type, {config.hidden_size, config.hidden_size}); + w.v_proj = store.load_tensor(source, pfx + "self_attn.v_proj.weight", storage_type, {config.hidden_size, config.hidden_size}); + w.o_proj = load_linear_with_bias(store, source, pfx + "self_attn.o_proj", config.hidden_size, config.hidden_size, storage_type); + + const int64_t max_pos_rows = 2 * config.max_position_embeddings + 1; + const auto raw_emb = source.require_f32(pfx + "self_attn.rel_pos_emb.weight", {max_pos_rows, config.head_dim}); + std::vector rel_128(static_cast(config.context_size * config.context_size * config.head_dim), 0.0f); + for (int64_t i = 0; i < config.context_size; ++i) { + for (int64_t j = 0; j < config.context_size; ++j) { + const int64_t dist = std::clamp(i - j, static_cast(-512), static_cast(512)) + 512; + const float * src_d = &raw_emb[static_cast(dist * config.head_dim)]; + float * dst_d = &rel_128[static_cast((i * config.context_size + j) * config.head_dim)]; + std::memcpy(dst_d, src_d, static_cast(config.head_dim) * sizeof(float)); + } + } + w.rel_pos_emb = store.make_from_f32( + core::TensorShape::from_dims({config.head_dim, config.context_size, config.context_size}), + storage_type == assets::TensorStorageType::Native ? assets::TensorStorageType::F32 : storage_type, + std::move(rel_128)); + + // Convolution + const int64_t conv_inner_dim = config.hidden_size * config.conv_expansion_factor; // 2048 + w.norm_conv = load_layer_norm(store, source, pfx + "norm_conv", config.hidden_size); + w.conv_pw1 = load_linear_with_bias(store, source, pfx + "conv.pointwise_lin1", config.hidden_size, conv_inner_dim * 2, storage_type); + load_folded_depthwise_conv(store, source, pfx + "conv", conv_inner_dim, config.conv_kernel_size, storage_type, w.conv_dw_weight, w.conv_dw_bias); + w.conv_pw2 = load_linear_with_bias(store, source, pfx + "conv.pointwise_lin2", conv_inner_dim, config.hidden_size, storage_type); + + // FFN2 + w.ffn2_norm = load_layer_norm(store, source, pfx + "norm_feed_forward2", config.hidden_size); + w.ffn2_fc1 = load_linear_with_bias(store, source, pfx + "feed_forward2.linear1", config.hidden_size, config.intermediate_size, storage_type); + w.ffn2_fc2 = load_linear_with_bias(store, source, pfx + "feed_forward2.linear2", config.intermediate_size, config.hidden_size, storage_type); + + // Norm Out + w.norm_out = load_layer_norm(store, source, pfx + "norm_out", config.hidden_size); + return w; +} + +Granite5EncoderWeights load_encoder_weights( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const Granite5ASRConfig & config, + assets::TensorStorageType storage_type) { + Granite5EncoderWeights weights; + weights.input_linear = load_linear_with_bias( + store, source, "encoder.input_linear", + config.encoder.input_features, config.encoder.hidden_size, storage_type); + + weights.layers.reserve(static_cast(config.encoder.num_layers)); + for (int64_t idx = 0; idx < config.encoder.num_layers; ++idx) { + weights.layers.push_back(load_layer_weights(store, source, idx, config.encoder, storage_type)); + } + + weights.out = load_linear_with_bias( + store, source, "encoder.out", + config.encoder.hidden_size, config.vocab_size, storage_type); + weights.out_mid = load_linear_with_bias( + store, source, "encoder.out_mid", + config.vocab_size, config.encoder.hidden_size, storage_type); + + store.upload(); + return weights; +} + +core::TensorValue build_block_attention( + core::ModuleBuildContext & ctx, + const core::TensorValue & x, + const Granite5LayerWeights & weights, + const Granite5EncoderConfig & config) { + const int64_t num_frames = x.shape.dims[1]; + const int64_t d_model = config.hidden_size; + const int64_t heads = config.num_attention_heads; + const int64_t d_head = config.head_dim; + const int64_t c = config.context_size; + const float scale = 1.0f / std::sqrt(static_cast(d_head)); + + auto q_all = modules::LinearModule({d_model, d_model, false}).build(ctx, x, {weights.q_proj, std::nullopt}); + auto k_all = modules::LinearModule({d_model, d_model, false}).build(ctx, x, {weights.k_proj, std::nullopt}); + auto v_all = modules::LinearModule({d_model, d_model, false}).build(ctx, x, {weights.v_proj, std::nullopt}); + + const int64_t nb_full = num_frames / c; + const int64_t nr = num_frames % c; + const int64_t num_blocks = nb_full + (nr > 0 ? 1 : 0); + + std::vector block_outputs; + block_outputs.reserve(static_cast(num_blocks)); + + for (int64_t b = 0; b < num_blocks; ++b) { + const int64_t blk = (b == nb_full) ? nr : c; + const int64_t t_offset = b * c; + + auto q_blk = modules::SliceModule({1, t_offset, blk}).build(ctx, q_all); + auto k_blk = modules::SliceModule({1, t_offset, blk}).build(ctx, k_all); + auto v_blk = modules::SliceModule({1, t_offset, blk}).build(ctx, v_all); + + auto q_3d = core::wrap_tensor( + ggml_reshape_3d(ctx.ggml, q_blk.tensor, d_head, heads, blk), + core::TensorShape::from_dims({blk, heads, d_head}), + GGML_TYPE_F32); + auto k_3d = core::wrap_tensor( + ggml_reshape_3d(ctx.ggml, k_blk.tensor, d_head, heads, blk), + core::TensorShape::from_dims({blk, heads, d_head}), + GGML_TYPE_F32); + auto v_3d = core::wrap_tensor( + ggml_reshape_3d(ctx.ggml, v_blk.tensor, d_head, heads, blk), + core::TensorShape::from_dims({blk, heads, d_head}), + GGML_TYPE_F32); + + auto q_perm = core::wrap_tensor( + ggml_permute(ctx.ggml, q_3d.tensor, 0, 2, 1, 3), + core::TensorShape::from_dims({heads, blk, d_head}), + GGML_TYPE_F32); + auto k_perm = core::wrap_tensor( + ggml_permute(ctx.ggml, k_3d.tensor, 0, 2, 1, 3), + core::TensorShape::from_dims({heads, blk, d_head}), + GGML_TYPE_F32); + + auto k_cont = core::wrap_tensor( + ggml_cont(ctx.ggml, k_perm.tensor), + k_perm.shape, + GGML_TYPE_F32); + auto scores = core::wrap_tensor( + ggml_mul_mat(ctx.ggml, k_cont.tensor, q_perm.tensor), + core::TensorShape::from_dims({heads, blk, blk}), + GGML_TYPE_F32); + + const size_t elem_size = ggml_type_size(weights.rel_pos_emb.type); + const size_t nb1 = static_cast(d_head) * elem_size; + const size_t nb2 = static_cast(c * d_head) * elem_size; + + auto rel_slice = core::wrap_tensor( + ggml_view_3d( + ctx.ggml, + weights.rel_pos_emb.tensor, + d_head, + blk, + blk, + nb1, + nb2, + 0), + core::TensorShape::from_dims({blk, blk, d_head}), + weights.rel_pos_emb.type); + auto rel_slice_cont = core::wrap_tensor( + ggml_cont(ctx.ggml, rel_slice.tensor), + rel_slice.shape, + weights.rel_pos_emb.type); + + auto pos_bias_c = core::wrap_tensor( + ggml_mul_mat(ctx.ggml, rel_slice_cont.tensor, q_3d.tensor), + core::TensorShape::from_dims({blk, heads, blk}), + GGML_TYPE_F32); + + auto pos_bias = core::wrap_tensor( + ggml_permute(ctx.ggml, pos_bias_c.tensor, 0, 2, 1, 3), + core::TensorShape::from_dims({heads, blk, blk}), + GGML_TYPE_F32); + + auto total_scores = core::wrap_tensor( + ggml_add(ctx.ggml, scores.tensor, pos_bias.tensor), + scores.shape, + GGML_TYPE_F32); + total_scores = core::wrap_tensor( + ggml_scale(ctx.ggml, total_scores.tensor, scale), + total_scores.shape, + GGML_TYPE_F32); + + auto attn_weights = core::wrap_tensor( + ggml_soft_max(ctx.ggml, total_scores.tensor), + total_scores.shape, + GGML_TYPE_F32); + + auto v_perm = core::wrap_tensor( + ggml_permute(ctx.ggml, v_3d.tensor, 1, 2, 0, 3), + core::TensorShape::from_dims({heads, d_head, blk}), + GGML_TYPE_F32); + auto v_cont = core::wrap_tensor( + ggml_cont(ctx.ggml, v_perm.tensor), + v_perm.shape, + GGML_TYPE_F32); + + auto out_head = core::wrap_tensor( + ggml_mul_mat(ctx.ggml, v_cont.tensor, attn_weights.tensor), + core::TensorShape::from_dims({heads, blk, d_head}), + GGML_TYPE_F32); + + auto out_perm = core::wrap_tensor( + ggml_permute(ctx.ggml, out_head.tensor, 0, 2, 1, 3), + core::TensorShape::from_dims({blk, heads, d_head}), + GGML_TYPE_F32); + auto out_cont = core::wrap_tensor( + ggml_cont(ctx.ggml, out_perm.tensor), + out_perm.shape, + GGML_TYPE_F32); + auto out_block = core::wrap_tensor( + ggml_reshape_3d(ctx.ggml, out_cont.tensor, d_model, blk, 1), + core::TensorShape::from_dims({1, blk, d_model}), + GGML_TYPE_F32); + block_outputs.push_back(out_block); + } + + core::TensorValue attn_out = block_outputs[0]; + for (size_t i = 1; i < block_outputs.size(); ++i) { + attn_out = modules::ConcatModule({1}).build(ctx, attn_out, block_outputs[i]); + } + + return modules::LinearModule({d_model, d_model, true}).build(ctx, attn_out, weights.o_proj); +} + +core::TensorValue build_conformer_block( + core::ModuleBuildContext & ctx, + const core::TensorValue & input_btc, + const Granite5LayerWeights & weights, + const Granite5EncoderConfig & config) { + const int64_t d_model = config.hidden_size; + const int64_t d_ffn = config.intermediate_size; + const int64_t d_conv = d_model * config.conv_expansion_factor; // 2048 + + // 1. FFN1: x + 0.5 * Linear2(SiLU(Linear1(LN(x)))) + auto h_norm1 = modules::LayerNormModule({d_model, kLayerNormEpsilon}).build(ctx, input_btc, weights.ffn1_norm); + auto h_ff1 = modules::LinearModule({d_model, d_ffn, true}).build(ctx, h_norm1, weights.ffn1_fc1); + h_ff1 = modules::SiluModule().build(ctx, h_ff1); + h_ff1 = modules::LinearModule({d_ffn, d_model, true}).build(ctx, h_ff1, weights.ffn1_fc2); + auto h_scaled1 = core::wrap_tensor( + ggml_scale(ctx.ggml, h_ff1.tensor, 0.5f), + h_ff1.shape, + GGML_TYPE_F32); + auto x1 = core::wrap_tensor( + ggml_add(ctx.ggml, input_btc.tensor, h_scaled1.tensor), + input_btc.shape, + GGML_TYPE_F32); + + // 2. Self-Attention: x1 + OutLinear(BlockAttn(LN(x1))) + auto h_norm_att = modules::LayerNormModule({d_model, kLayerNormEpsilon}).build(ctx, x1, weights.norm_self_att); + auto h_att = build_block_attention(ctx, h_norm_att, weights, config); + auto x2 = core::wrap_tensor( + ggml_add(ctx.ggml, x1.tensor, h_att.tensor), + x1.shape, + GGML_TYPE_F32); + + // 3. Conformer Convolution + auto h_norm_conv = modules::LayerNormModule({d_model, kLayerNormEpsilon}).build(ctx, x2, weights.norm_conv); + auto h_pw1 = modules::LinearModule({d_model, d_conv * 2, true}).build(ctx, h_norm_conv, weights.conv_pw1); + auto h_glu = modules::GLUModule().build(ctx, h_pw1); // [1, T, 2048] + + auto h_bct = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, h_glu); + + const int stride = weights.is_subsample ? 2 : 1; + auto conv_out = modules::DepthwiseConv1dModule({ + d_conv, + config.conv_kernel_size, + stride, + static_cast((config.conv_kernel_size - 1) / 2), + 1, + true + }).build(ctx, h_bct, {weights.conv_dw_weight, weights.conv_dw_bias}); + + auto conv_btc = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, conv_out); + conv_btc = modules::SiluModule().build(ctx, conv_btc); + auto conv_pw2 = modules::LinearModule({d_conv, d_model, true}).build(ctx, conv_btc, weights.conv_pw2); + + core::TensorValue x3; + if (weights.is_subsample) { + const int64_t t_half = x2.shape.dims[1] / 2; + auto x2_cont = core::wrap_tensor( + ggml_cont(ctx.ggml, x2.tensor), + x2.shape, + GGML_TYPE_F32); + auto x2_pooled = core::wrap_tensor( + ggml_pool_2d(ctx.ggml, x2_cont.tensor, GGML_OP_POOL_AVG, 1, 2, 1, 2, 0, 0), + core::TensorShape::from_dims({1, t_half, d_model}), + GGML_TYPE_F32); + auto conv_pw2_trimmed = (conv_pw2.shape.dims[1] == t_half) + ? conv_pw2 + : modules::SliceModule({1, 0, t_half}).build(ctx, conv_pw2); + x3 = core::wrap_tensor( + ggml_add(ctx.ggml, conv_pw2_trimmed.tensor, x2_pooled.tensor), + x2_pooled.shape, + GGML_TYPE_F32); + } else { + x3 = core::wrap_tensor( + ggml_add(ctx.ggml, x2.tensor, conv_pw2.tensor), + x2.shape, + GGML_TYPE_F32); + } + + // 4. FFN2: x3 + 0.5 * Linear2(SiLU(Linear1(LN(x3)))) + auto h_norm2 = modules::LayerNormModule({d_model, kLayerNormEpsilon}).build(ctx, x3, weights.ffn2_norm); + auto h_ff2 = modules::LinearModule({d_model, d_ffn, true}).build(ctx, h_norm2, weights.ffn2_fc1); + h_ff2 = modules::SiluModule().build(ctx, h_ff2); + h_ff2 = modules::LinearModule({d_ffn, d_model, true}).build(ctx, h_ff2, weights.ffn2_fc2); + auto h_scaled2 = core::wrap_tensor( + ggml_scale(ctx.ggml, h_ff2.tensor, 0.5f), + h_ff2.shape, + GGML_TYPE_F32); + auto x4 = core::wrap_tensor( + ggml_add(ctx.ggml, x3.tensor, h_scaled2.tensor), + x3.shape, + GGML_TYPE_F32); + + // 5. Output LayerNorm + return modules::LayerNormModule({d_model, kLayerNormEpsilon}).build(ctx, x4, weights.norm_out); +} + +} // namespace + +Granite5EncoderRuntime::Granite5EncoderRuntime( + std::shared_ptr assets, + engine::core::ExecutionContext & execution_context, + assets::TensorStorageType storage_type, + size_t graph_arena_bytes) + : assets_(std::move(assets)), + execution_context_(&execution_context), + weight_store_( + execution_context.backend(), + execution_context.backend_type(), + "Granite 5 ASR encoder weights", + 256ull * 1024ull * 1024ull), + graph_arena_bytes_(graph_arena_bytes) { + if (assets_ == nullptr) { + throw std::runtime_error("Granite 5 ASR encoder runtime requires assets"); + } + weights_ = load_encoder_weights(weight_store_, *assets_->source, assets_->config, storage_type); +} + +std::vector Granite5EncoderRuntime::transcribe_features( + const Granite5FrontendFeatures & features) { + if (features.frames <= 0 || features.values.empty()) { + return {}; + } + + const auto & config = assets_->config; + const int64_t num_frames = features.frames; + const int64_t feat_dim = features.feature_dim; + + ggml_init_params params{}; + params.mem_size = graph_arena_bytes_; + params.mem_buffer = nullptr; + params.no_alloc = true; + + ggml_context * ggml_ctx = ggml_init(params); + if (!ggml_ctx) { + throw std::runtime_error("Failed to initialize GGML context for Granite 5 ASR encoder"); + } + + ggml_gallocr * galloc = ggml_gallocr_new( + ggml_backend_get_default_buffer_type(execution_context_->backend())); + if (!galloc) { + ggml_free(ggml_ctx); + throw std::runtime_error("Failed to initialize GGML allocator for Granite 5 ASR encoder"); + } + + std::vector token_ids; + + try { + core::ModuleBuildContext ctx{ggml_ctx, "granite5asr_encoder", execution_context_->backend_type()}; + + auto in_tensor = core::wrap_tensor( + ggml_new_tensor_2d(ggml_ctx, GGML_TYPE_F32, feat_dim, num_frames), + core::TensorShape::from_dims({1, num_frames, feat_dim}), + GGML_TYPE_F32); + + auto h = modules::LinearModule({feat_dim, config.encoder.hidden_size, true}) + .build(ctx, in_tensor, weights_.input_linear); + + const int64_t mid_layer_idx = config.encoder.num_layers / 2; // 8 + for (int64_t idx = 0; idx < config.encoder.num_layers; ++idx) { + h = build_conformer_block(ctx, h, weights_.layers[static_cast(idx)], config.encoder); + + if (idx + 1 == mid_layer_idx) { + auto h_mid = modules::LinearModule({config.encoder.hidden_size, config.vocab_size, true}) + .build(ctx, h, weights_.out); + auto p_mid = core::wrap_tensor( + ggml_soft_max(ctx.ggml, h_mid.tensor), + h_mid.shape, + GGML_TYPE_F32); + auto h_inj = modules::LinearModule({config.vocab_size, config.encoder.hidden_size, true}) + .build(ctx, p_mid, weights_.out_mid); + h = core::wrap_tensor( + ggml_add(ctx.ggml, h.tensor, h_inj.tensor), + h.shape, + GGML_TYPE_F32); + } + } + + auto logits = modules::LinearModule({config.encoder.hidden_size, config.vocab_size, true}) + .build(ctx, h, weights_.out); + + ggml_cgraph * gf = ggml_new_graph_custom(ggml_ctx, kEncoderGraphNodes, false); + ggml_build_forward_expand(gf, logits.tensor); + + if (!ggml_gallocr_alloc_graph(galloc, gf)) { + throw std::runtime_error("Failed to allocate GGML graph for Granite 5 ASR encoder"); + } + + ggml_backend_tensor_set( + in_tensor.tensor, + features.values.data(), + 0, + features.values.size() * sizeof(float)); + + if (ggml_backend_graph_compute(execution_context_->backend(), gf) != GGML_STATUS_SUCCESS) { + throw std::runtime_error("Failed to compute GGML graph for Granite 5 ASR encoder"); + } + + const int64_t out_frames = logits.shape.dims[1]; + const int64_t vocab_size = config.vocab_size; + std::vector logits_data(static_cast(out_frames * vocab_size)); + ggml_backend_tensor_get( + logits.tensor, + logits_data.data(), + 0, + logits_data.size() * sizeof(float)); + + token_ids.reserve(static_cast(out_frames)); + for (int64_t t = 0; t < out_frames; ++t) { + const float * frame_logits = &logits_data[static_cast(t * vocab_size)]; + int32_t best_id = 0; + float max_val = frame_logits[0]; + for (int32_t v = 1; v < static_cast(vocab_size); ++v) { + if (frame_logits[v] > max_val) { + max_val = frame_logits[v]; + best_id = v; + } + } + token_ids.push_back(best_id); + } + } catch (...) { + ggml_gallocr_free(galloc); + ggml_free(ggml_ctx); + throw; + } + + ggml_gallocr_free(galloc); + ggml_free(ggml_ctx); + return token_ids; +} + +} // namespace engine::community_models::granite5asr diff --git a/src/community_models/granite5asr/frontend.cpp b/src/community_models/granite5asr/frontend.cpp new file mode 100644 index 000000000..c2b7513f0 --- /dev/null +++ b/src/community_models/granite5asr/frontend.cpp @@ -0,0 +1,229 @@ +#include "engine/community_models/granite5asr/frontend.h" + +#include "engine/framework/audio/conversion.h" +#include "engine/framework/audio/dsp.h" + +#include +#include +#include +#include + +namespace engine::community_models::granite5asr { +namespace { + +constexpr double kPi = 3.14159265358979323846; + +std::vector make_periodic_hann_window(int64_t win_length) { + std::vector window(static_cast(win_length), 0.0f); + if (win_length <= 1) { + if (!window.empty()) { + window[0] = 1.0f; + } + return window; + } + for (int64_t i = 0; i < win_length; ++i) { + window[static_cast(i)] = static_cast( + 0.5 - 0.5 * std::cos(2.0 * kPi * static_cast(i) / static_cast(win_length))); + } + return window; +} + +double hz_to_htk_mel(double hz) { + return 2595.0 * std::log10(1.0 + hz / 700.0); +} + +double htk_mel_to_hz(double mel) { + return 700.0 * (std::pow(10.0, mel / 2595.0) - 1.0); +} + +audio::AudioTensor build_htk_mel_filterbank(int64_t sample_rate, int64_t n_fft, int64_t n_mels) { + const int64_t num_bins = n_fft / 2 + 1; + const double f_min = 0.0; + const double f_max = static_cast(sample_rate) / 2.0; + + const double min_mel = hz_to_htk_mel(f_min); + const double max_mel = hz_to_htk_mel(f_max); + + std::vector mel_points(static_cast(n_mels + 2)); + for (size_t i = 0; i < mel_points.size(); ++i) { + mel_points[i] = min_mel + (max_mel - min_mel) * static_cast(i) / static_cast(n_mels + 1); + } + + std::vector fft_bins(static_cast(n_mels + 2)); + for (size_t i = 0; i < fft_bins.size(); ++i) { + const double hz = htk_mel_to_hz(mel_points[i]); + fft_bins[i] = static_cast(std::floor((static_cast(n_fft) + 1.0) * hz / static_cast(sample_rate))); + fft_bins[i] = std::clamp(fft_bins[i], static_cast(0), num_bins - 1); + } + + audio::AudioTensor fb; + fb.shape = {n_mels, num_bins}; + fb.values.assign(static_cast(n_mels * num_bins), 0.0f); + + for (int64_t m = 1; m <= n_mels; ++m) { + const int64_t f_left = fft_bins[static_cast(m - 1)]; + const int64_t f_center = fft_bins[static_cast(m)]; + const int64_t f_right = fft_bins[static_cast(m + 1)]; + + for (int64_t k = f_left; k < f_center; ++k) { + if (f_center > f_left) { + fb.values[static_cast((m - 1) * num_bins + k)] = + static_cast(k - f_left) / static_cast(f_center - f_left); + } + } + for (int64_t k = f_center; k < f_right; ++k) { + if (f_right > f_center) { + fb.values[static_cast((m - 1) * num_bins + k)] = + static_cast(f_right - k) / static_cast(f_right - f_center); + } + } + } + return fb; +} + +} // namespace + +Granite5Frontend::Granite5Frontend(std::shared_ptr assets) + : assets_(std::move(assets)) { + if (assets_ == nullptr) { + throw std::runtime_error("Granite 5 ASR frontend requires valid assets"); + } + const auto & config = assets_->config.frontend; + mel_filterbank_ = build_htk_mel_filterbank(config.sample_rate, config.n_fft, config.n_mels); + window_ = make_periodic_hann_window(config.win_length); +} + +std::vector Granite5Frontend::prepare_waveform(const runtime::AudioBuffer & audio) const { + if (audio.sample_rate <= 0 || audio.channels <= 0) { + throw std::runtime_error("Granite 5 ASR audio requires positive sample rate and channels"); + } + if (audio.samples.empty()) { + throw std::runtime_error("Granite 5 ASR audio input is empty"); + } + const auto & config = assets_->config.frontend; + return engine::audio::convert_interleaved_audio_to_mono_linear_resampled( + audio.samples, + audio.sample_rate, + audio.channels, + static_cast(config.sample_rate)); +} + +Granite5FrontendFeatures Granite5Frontend::extract(const runtime::AudioBuffer & audio) const { + return extract_waveform(prepare_waveform(audio)); +} + +Granite5FrontendFeatures Granite5Frontend::extract_waveform(const std::vector & waveform) const { + if (waveform.empty()) { + throw std::runtime_error("Granite 5 ASR waveform is empty"); + } + const auto & config = assets_->config.frontend; + const int64_t hop = config.hop_length; + const int64_t s = config.stack_factor; + const int64_t T = static_cast(waveform.size()); + const int64_t mel_frames = T / hop; + const int64_t n_frames = s * ((mel_frames + s - 1) / s); + const int64_t need_samples = (n_frames - 1) * hop + 1; + + std::vector padded_waveform = waveform; + if (static_cast(padded_waveform.size()) < need_samples) { + padded_waveform.resize(static_cast(need_samples), 0.0f); + } + + audio::STFTConfig stft_cfg; + stft_cfg.n_fft = config.n_fft; + stft_cfg.hop_length = config.hop_length; + stft_cfg.win_length = config.win_length; + stft_cfg.center = true; + stft_cfg.pad_mode = audio::STFTPadMode::Reflect; + + audio::STFT stft; + const auto mag_tensor = stft.compute_magnitude( + padded_waveform, + window_, + 1, + static_cast(padded_waveform.size()), + stft_cfg); + + const int64_t freq_bins = config.n_fft / 2 + 1; + const int64_t available_frames = mag_tensor.shape.size() >= 3 ? mag_tensor.shape[2] : (static_cast(mag_tensor.values.size()) / freq_bins); + const int64_t frames_to_use = std::min(n_frames, available_frames); + + const int64_t n_mels = config.n_mels; + std::vector logmel(static_cast(n_mels * n_frames), 0.0f); + + float max_logmel = -1e30f; + for (int64_t m = 0; m < n_mels; ++m) { + for (int64_t t = 0; t < frames_to_use; ++t) { + float energy = 0.0f; + for (int64_t f = 0; f < freq_bins; ++f) { + const float mag = mag_tensor.values[static_cast(f * available_frames + t)]; + const float power = mag * mag; + energy += mel_filterbank_.values[static_cast(m * freq_bins + f)] * power; + } + if (energy < 1e-10f) { + energy = 1e-10f; + } + const float log_val = std::log10(energy); + logmel[static_cast(m * n_frames + t)] = log_val; + if (log_val > max_logmel) { + max_logmel = log_val; + } + } + for (int64_t t = frames_to_use; t < n_frames; ++t) { + logmel[static_cast(m * n_frames + t)] = logmel[static_cast(m * n_frames + (frames_to_use > 0 ? frames_to_use - 1 : 0))]; + } + } + + const float floor_thresh = max_logmel - config.logmel_floor_db; + for (float & val : logmel) { + if (val < floor_thresh) { + val = floor_thresh; + } + val = (val / 4.0f) + 1.0f; + } + + std::vector deltas(static_cast(n_mels * n_frames), 0.0f); + if (config.deltas) { + for (int64_t m = 0; m < n_mels; ++m) { + const size_t offset = static_cast(m * n_frames); + if (n_frames == 1) { + deltas[offset] = 0.0f; + } else { + deltas[offset] = (logmel[offset + 1] - logmel[offset]) * 0.5f; + for (int64_t t = 1; t < n_frames - 1; ++t) { + deltas[offset + static_cast(t)] = + (logmel[offset + static_cast(t + 1)] - logmel[offset + static_cast(t - 1)]) * 0.5f; + } + deltas[offset + static_cast(n_frames - 1)] = + (logmel[offset + static_cast(n_frames - 1)] - logmel[offset + static_cast(n_frames - 2)]) * 0.5f; + } + } + } + + const int64_t out_frames = n_frames / s; + const int64_t feature_dim = n_mels * (config.deltas ? 2 : 1) * s; + + Granite5FrontendFeatures features; + features.frames = out_frames; + features.feature_dim = feature_dim; + features.values.assign(static_cast(out_frames * feature_dim), 0.0f); + + const int64_t channels_per_step = n_mels * (config.deltas ? 2 : 1); + for (int64_t t_out = 0; t_out < out_frames; ++t_out) { + float * dst_frame = &features.values[static_cast(t_out * feature_dim)]; + for (int64_t step = 0; step < s; ++step) { + const int64_t t_in = t_out * s + step; + float * dst_step = dst_frame + step * channels_per_step; + for (int64_t m = 0; m < n_mels; ++m) { + dst_step[m] = logmel[static_cast(m * n_frames + t_in)]; + if (config.deltas) { + dst_step[n_mels + m] = deltas[static_cast(m * n_frames + t_in)]; + } + } + } + } + + return features; +} + +} // namespace engine::community_models::granite5asr diff --git a/src/community_models/granite5asr/session.cpp b/src/community_models/granite5asr/session.cpp new file mode 100644 index 000000000..b014231c3 --- /dev/null +++ b/src/community_models/granite5asr/session.cpp @@ -0,0 +1,478 @@ +#include "engine/community_models/granite5asr/session.h" + +#include "engine/framework/audio/chunking.h" +#include "engine/framework/audio/conversion.h" +#include "engine/framework/debug/profiler.h" +#include "engine/framework/debug/trace.h" +#include "engine/framework/io/filesystem.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 "engine/framework/runtime/spec_backed_model.h" +#include "engine/models/silero_vad/session.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace engine::community_models::granite5asr { +namespace json = engine::io::json; +namespace { + +using Clock = std::chrono::steady_clock; +constexpr float kDefaultChunkSeconds = 30.0f; +constexpr size_t kDefaultGraphArenaBytes = 1024ull * 1024ull * 1024ull; + +std::filesystem::path default_spec_path() { + return engine::model_spec::default_spec_path("granite5asr"); +} + +std::shared_ptr require_assets( + std::shared_ptr assets) { + if (assets == nullptr) { + throw std::runtime_error("Granite 5 ASR session requires assets"); + } + return assets; +} + +assets::TensorStorageType parse_weight_storage( + const runtime::SessionOptions & options, + const std::string & family) { + const std::string key = family + ".weight_type"; + const auto it = options.options.find(key); + if (it == options.options.end()) { + return assets::TensorStorageType::Native; + } + const auto type = assets::parse_tensor_storage_type(it->second); + if (type == assets::TensorStorageType::Native || + type == assets::TensorStorageType::F32 || + type == assets::TensorStorageType::F16 || + type == assets::TensorStorageType::BF16 || + type == assets::TensorStorageType::Q8_0) { + return type; + } + throw std::runtime_error(key + " supports only native, f32, f16, bf16, and q8_0"); +} + +size_t parse_graph_arena_bytes( + const runtime::SessionOptions & options, + const std::string & family) { + const std::string key = family + ".encoder_graph_arena_mb"; + const auto it = options.options.find(key); + if (it == options.options.end()) { + return kDefaultGraphArenaBytes; + } + try { + const int64_t mb = std::stoll(it->second); + if (mb > 0) { + return static_cast(mb) * 1024ull * 1024ull; + } + } catch (...) {} + return kDefaultGraphArenaBytes; +} + +std::string parse_vad_model_path( + const runtime::SessionOptions & options, + const std::string & family) { + const std::string key = family + ".vad_model_path"; + const auto it = options.options.find(key); + if (it != options.options.end() && !it->second.empty()) { + return it->second; + } + return "assets/framework/models/silero_vad"; +} + +std::vector ctc_greedy_collapse( + const std::vector & raw_ids, + int64_t blank_id) { + std::vector collapsed; + int32_t prev = -1; + for (int32_t id : raw_ids) { + if (id != prev) { + if (id != static_cast(blank_id)) { + collapsed.push_back(id); + } + prev = id; + } + } + return collapsed; +} + +void append_chunk_transcript(std::string & merged, std::string chunk_text) { + chunk_text = engine::io::trim_ascii_whitespace(std::move(chunk_text)); + if (chunk_text.empty()) { + return; + } + if (!merged.empty()) { + const char last = merged.back(); + const char first = chunk_text.front(); + if (last != ' ' && first != ' ') { + merged.push_back(' '); + } + } + merged += chunk_text; +} + +class Granite5ASRLoader final : public runtime::IVoiceModelLoader { +public: + std::string family() const override { + return "granite5asr"; + } + + std::vector family_aliases() const override { + return {"granite_speech5_asr", "granite_speech", "granite_speech5_ctc"}; + } + + runtime::CapabilitySet advertised_capabilities() const override { + runtime::CapabilitySet out; + out.supported_tasks = { + {runtime::VoiceTaskKind::Asr, {runtime::RunMode::Offline, runtime::RunMode::Streaming}}, + }; + out.languages = {"en"}; + 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"); + const auto model_type = json::optional_string(config_root, "model_type", ""); + return model_type == "granite_speech5_ctc" || model_type == "granite5asr" || model_type == "granite_speech"; + } 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 = "470m-turboctc"; + inspection.metadata.description = "IBM Granite Speech 5.0 470M TurboCTC ASR model."; + inspection.capabilities.supported_tasks = { + {runtime::VoiceTaskKind::Asr, {runtime::RunMode::Offline, runtime::RunMode::Streaming}}, + }; + inspection.capabilities.languages = {"en"}; + inspection.capabilities.supports_timestamps = false; + 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_granite5asr_model(request); + } +}; + +} // namespace + +Granite5ASRSessionBase::Granite5ASRSessionBase( + 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_), + vad_model_path_(parse_vad_model_path(RuntimeSessionBase::options(), family_impl())) { + if (task_.task != runtime::VoiceTaskKind::Asr) { + throw std::runtime_error("Granite 5 ASR only supports VoiceTaskKind::Asr"); + } + const auto storage_type = parse_weight_storage(RuntimeSessionBase::options(), family_impl()); + const auto arena_bytes = parse_graph_arena_bytes(RuntimeSessionBase::options(), family_impl()); + encoder_ = std::unique_ptr( + new Granite5EncoderRuntime( + assets_, + execution_context(), + storage_type, + arena_bytes)); +} + +Granite5ASRSessionBase::~Granite5ASRSessionBase() = default; + +std::string Granite5ASRSessionBase::family_impl() const { + return "granite5asr"; +} + +runtime::VoiceTaskKind Granite5ASRSessionBase::task_kind_impl() const { + return task_.task; +} + +runtime::RunMode Granite5ASRSessionBase::run_mode_impl() const { + return task_.mode; +} + +runtime::IOfflineVoiceTaskSession & Granite5ASRSessionBase::vad_session() { + if (vad_session_ == nullptr) { + runtime::ModelLoadRequest load_request; + load_request.model_path = vad_model_path_; + vad_model_ = engine::models::silero_vad::load_silero_vad_model(load_request); + auto session = vad_model_->create_task_session( + runtime::TaskSpec{runtime::VoiceTaskKind::Vad, runtime::RunMode::Offline}, + runtime::SessionOptions{options().backend, {}}); + auto * offline = dynamic_cast(session.get()); + if (offline == nullptr) { + throw std::runtime_error("Granite 5 ASR internal VAD session does not support offline execution"); + } + session.release(); + vad_session_.reset(offline); + } + return *vad_session_; +} + +runtime::Transcript Granite5ASRSessionBase::transcribe_audio( + const runtime::AudioBuffer & audio, + const std::unordered_map & options) { + if (audio.samples.empty()) { + return {"", "en"}; + } + + auto transcribe_clip = [this](const runtime::AudioBuffer & clip) -> std::string { + if (clip.samples.empty()) { + return ""; + } + const auto features = frontend_.extract(clip); + const auto raw_tokens = encoder_->transcribe_features(features); + const auto collapsed = ctc_greedy_collapse(raw_tokens, assets_->config.blank_token_id); + if (collapsed.empty() || assets_->tokenizer == nullptr) { + return ""; + } + return assets_->tokenizer->decode_ids(collapsed); + }; + + const auto mode = engine::audio::parse_audio_chunk_mode(options); + if (mode == engine::audio::AudioChunkMode::None) { + return {engine::io::trim_ascii_whitespace(transcribe_clip(audio)), "en"}; + } + + const float duration_sec = static_cast(audio.samples.size()) / static_cast(std::max(1, audio.sample_rate * audio.channels)); + const float chunk_duration_sec = engine::audio::parse_audio_chunk_seconds_override(options).value_or(kDefaultChunkSeconds); + + if (duration_sec <= chunk_duration_sec) { + return {engine::io::trim_ascii_whitespace(transcribe_clip(audio)), "en"}; + } + + std::string full_text; + if ((mode == engine::audio::AudioChunkMode::Vad || mode == engine::audio::AudioChunkMode::Auto) && + engine::io::is_existing_file(vad_model_path_)) { + try { + const auto vad_options = engine::audio::VadAudioChunkOptions{ + static_cast(std::llround(static_cast(chunk_duration_sec) * static_cast(audio.sample_rate))), + static_cast(std::llround(0.5 * static_cast(audio.sample_rate))), + static_cast(std::llround(0.25 * static_cast(audio.sample_rate))), + }; + const auto spans = engine::audio::plan_vad_audio_chunks(audio, vad_session(), vad_options); + for (const auto & span : spans) { + const auto chunk_buf = engine::audio::slice_audio_buffer(audio, span); + append_chunk_transcript(full_text, transcribe_clip(chunk_buf)); + } + return {engine::io::trim_ascii_whitespace(std::move(full_text)), "en"}; + } catch (...) { + // Fallback to fixed chunking + } + } + + // Fixed chunking + const int64_t samples = static_cast(std::llround(static_cast(chunk_duration_sec) * static_cast(audio.sample_rate))); + const int64_t frames = static_cast(audio.samples.size() / std::max(1, audio.channels)); + const auto chunks = engine::audio::plan_audio_chunks( + frames, + {samples, samples, engine::audio::AudioChunkPadMode::Zero, engine::audio::AudioChunkTailAlignment::Start, 0}); + for (const auto & chunk : chunks) { + runtime::TimeSpan span{chunk.output_start_sample, chunk.output_start_sample + chunk.valid_samples}; + const auto chunk_buf = engine::audio::slice_audio_buffer(audio, span); + append_chunk_transcript(full_text, transcribe_clip(chunk_buf)); + } + + return {engine::io::trim_ascii_whitespace(std::move(full_text)), "en"}; +} + +// ----------------- Offline Session ----------------- + +Granite5ASROfflineSession::Granite5ASROfflineSession( + runtime::TaskSpec task, + runtime::SessionOptions options, + std::shared_ptr assets) + : Granite5ASRSessionBase(std::move(task), std::move(options), std::move(assets)) {} + +std::string Granite5ASROfflineSession::family() const { + return family_impl(); +} + +runtime::VoiceTaskKind Granite5ASROfflineSession::task_kind() const { + return task_kind_impl(); +} + +runtime::RunMode Granite5ASROfflineSession::run_mode() const { + return run_mode_impl(); +} + +void Granite5ASROfflineSession::prepare(const runtime::SessionPreparationRequest & request) { + mark_prepared(); +} + +runtime::TaskResult Granite5ASROfflineSession::run(const runtime::TaskRequest & request) { + require_prepared("Granite 5 ASR run()"); + if (!request.audio_input.has_value()) { + throw std::runtime_error("Granite 5 ASR run() requires audio_input"); + } + const auto wall_start = Clock::now(); + const auto transcript = transcribe_audio(*request.audio_input, request.options); + runtime::TaskResult result; + result.text_output = transcript; + engine::debug::timing_log_scalar("session.wall_ms", engine::debug::elapsed_ms(wall_start)); + return result; +} + +// ----------------- Streaming Session ----------------- + +Granite5ASRStreamingSession::Granite5ASRStreamingSession( + runtime::TaskSpec task, + runtime::SessionOptions options, + std::shared_ptr assets) + : Granite5ASRSessionBase(std::move(task), std::move(options), std::move(assets)) {} + +std::string Granite5ASRStreamingSession::family() const { + return family_impl(); +} + +runtime::VoiceTaskKind Granite5ASRStreamingSession::task_kind() const { + return task_kind_impl(); +} + +runtime::RunMode Granite5ASRStreamingSession::run_mode() const { + return run_mode_impl(); +} + +void Granite5ASRStreamingSession::prepare(const runtime::SessionPreparationRequest & request) { + mark_prepared(); +} + +runtime::StreamingPolicy Granite5ASRStreamingSession::streaming_policy() const { + runtime::StreamingPolicy policy; + policy.input = runtime::StreamingInputKind::AudioChunks; + policy.output = runtime::StreamingOutputKind::FinalResult; + policy.preferred_audio_chunk_samples = 512; + return policy; +} + +void Granite5ASRStreamingSession::start_stream(const runtime::TaskRequest & request) { + require_prepared("Granite 5 ASR start_stream()"); + streaming_request_ = request; + streaming_audio_ = runtime::AudioBuffer{}; + streaming_audio_.sample_rate = 16000; + streaming_audio_.channels = 1; +} + +void Granite5ASRStreamingSession::set_stream_event_sink(runtime::StreamEventCallback sink) { + stream_event_sink_ = std::move(sink); +} + +void Granite5ASRStreamingSession::reset() { + streaming_audio_.samples.clear(); +} + +runtime::StreamEvent Granite5ASRStreamingSession::process_audio_chunk(const runtime::AudioChunk & chunk) { + require_prepared("Granite 5 ASR process_audio_chunk()"); + if (!chunk.samples.empty()) { + const auto mono = engine::audio::convert_interleaved_audio_to_mono_linear_resampled( + chunk.samples, + chunk.sample_rate, + chunk.channels, + 16000); + streaming_audio_.samples.insert(streaming_audio_.samples.end(), mono.begin(), mono.end()); + } + return {}; +} + +runtime::TaskResult Granite5ASRStreamingSession::finish_stream() { + require_prepared("Granite 5 ASR finish_stream()"); + const auto transcript = transcribe_audio(streaming_audio_, streaming_request_.options); + runtime::TaskResult result; + result.text_output = transcript; + return result; +} + +runtime::TaskResult Granite5ASRStreamingSession::finalize() { + return finish_stream(); +} + +// ----------------- Loaded Model & Loader ----------------- + +Granite5ASRLoadedModel::Granite5ASRLoadedModel( + 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 & Granite5ASRLoadedModel::metadata() const noexcept { + return metadata_; +} + +const runtime::CapabilitySet & Granite5ASRLoadedModel::capabilities() const noexcept { + return capabilities_; +} + +std::unique_ptr Granite5ASRLoadedModel::create_task_session( + const runtime::TaskSpec & task, + const runtime::SessionOptions & options) const { + if (task.mode == runtime::RunMode::Streaming) { + return std::make_unique(task, options, assets_); + } + return std::make_unique(task, options, assets_); +} + +std::unique_ptr load_granite5asr_model( + const runtime::ModelLoadRequest & request) { + auto assets = load_granite5asr_assets(request.model_path); + + runtime::ModelMetadata metadata; + metadata.family = "granite5asr"; + metadata.variant = "470m-turboctc"; + metadata.description = "IBM Granite Speech 5.0 470M TurboCTC ASR model."; + + runtime::CapabilitySet capabilities; + capabilities.supported_tasks = { + {runtime::VoiceTaskKind::Asr, {runtime::RunMode::Offline, runtime::RunMode::Streaming}}, + }; + capabilities.languages = {"en"}; + capabilities.supports_timestamps = false; + + return std::make_unique( + std::move(metadata), + std::move(capabilities), + std::move(assets)); +} + +std::shared_ptr make_granite5asr_loader() { + return std::make_shared(); +} + +} // namespace engine::community_models::granite5asr diff --git a/src/framework/tokenizers/hf_tokenizer_json.cpp b/src/framework/tokenizers/hf_tokenizer_json.cpp index 9a837b1a5..317135750 100644 --- a/src/framework/tokenizers/hf_tokenizer_json.cpp +++ b/src/framework/tokenizers/hf_tokenizer_json.cpp @@ -5,6 +5,8 @@ #include #include +#include + namespace engine::tokenizers { namespace { @@ -20,11 +22,71 @@ std::string replace_all(std::string text, const std::string & needle, const std: return text; } +std::unordered_map build_unicode_to_bytes_map() { + std::unordered_map map; + std::vector bs; + std::vector cs; + for (int b = '!'; b <= '~'; ++b) { bs.push_back(b); cs.push_back(b); } + for (int b = 161; b <= 172; ++b) { bs.push_back(b); cs.push_back(b); } + for (int b = 174; b <= 255; ++b) { bs.push_back(b); cs.push_back(b); } + int n = 0; + for (int b = 0; b < 256; ++b) { + if (std::find(bs.begin(), bs.end(), b) == bs.end()) { + bs.push_back(b); + cs.push_back(256 + n); + ++n; + } + } + for (size_t i = 0; i < bs.size(); ++i) { + map[static_cast(cs[i])] = static_cast(bs[i]); + } + return map; +} + +std::string decode_byte_level(const std::string & text) { + static const auto map = build_unicode_to_bytes_map(); + std::string bytes; + bytes.reserve(text.size()); + for (size_t i = 0; i < text.size(); ) { + unsigned char c = static_cast(text[i]); + uint32_t cp = 0; + size_t len = 0; + if (c < 0x80) { + cp = c; + len = 1; + } else if ((c & 0xE0) == 0xC0 && i + 1 < text.size()) { + cp = ((c & 0x1F) << 6) | (static_cast(text[i + 1]) & 0x3F); + len = 2; + } else if ((c & 0xF0) == 0xE0 && i + 2 < text.size()) { + cp = ((c & 0x0F) << 12) | ((static_cast(text[i + 1]) & 0x3F) << 6) | (static_cast(text[i + 2]) & 0x3F); + len = 3; + } else if ((c & 0xF8) == 0xF0 && i + 3 < text.size()) { + cp = ((c & 0x07) << 18) | ((static_cast(text[i + 1]) & 0x3F) << 12) | ((static_cast(text[i + 2]) & 0x3F) << 6) | (static_cast(text[i + 3]) & 0x3F); + len = 4; + } else { + bytes.push_back(text[i]); + ++i; + continue; + } + i += len; + auto it = map.find(cp); + if (it != map.end()) { + bytes.push_back(static_cast(it->second)); + } else { + for (size_t k = i - len; k < i; ++k) { + bytes.push_back(text[k]); + } + } + } + return bytes; +} + std::string decode_with_optional_skip( const std::vector & id_to_token, const std::vector & ids, const std::string & metaspace_replacement, bool trim_leading_space, + bool byte_level, const int32_t * skip_token_id) { std::string text; for (const int32_t id : ids) { @@ -36,6 +98,9 @@ std::string decode_with_optional_skip( } text += id_to_token[static_cast(id)]; } + if (byte_level) { + text = decode_byte_level(text); + } if (!metaspace_replacement.empty()) { text = replace_all(text, metaspace_replacement, " "); } @@ -47,26 +112,51 @@ std::string decode_with_optional_skip( return text; } +bool has_byte_level(const engine::io::json::Value & root) { + if (const auto * decoder = root.find("decoder"); decoder != nullptr && decoder->is_object()) { + if (const auto * type = decoder->find("type"); type != nullptr && type->as_string() == "ByteLevel") { + return true; + } + } + if (const auto * pre = root.find("pre_tokenizer"); pre != nullptr && pre->is_object()) { + if (const auto * type = pre->find("type"); type != nullptr && type->as_string() == "ByteLevel") { + return true; + } + if (const auto * pretokenizers = pre->find("pretokenizers"); pretokenizers != nullptr && pretokenizers->is_array()) { + for (const auto & item : pretokenizers->as_array()) { + if (item.is_object()) { + if (const auto * type = item.find("type"); type != nullptr && type->as_string() == "ByteLevel") { + return true; + } + } + } + } + } + return false; +} + } // namespace HuggingFaceTokenizerJson::HuggingFaceTokenizerJson( std::vector id_to_token, std::string metaspace_replacement, - bool trim_leading_space) + bool trim_leading_space, + bool byte_level) : id_to_token_(std::move(id_to_token)), metaspace_replacement_(std::move(metaspace_replacement)), - trim_leading_space_(trim_leading_space) {} + trim_leading_space_(trim_leading_space), + byte_level_(byte_level) {} const std::vector & HuggingFaceTokenizerJson::id_to_token() const noexcept { return id_to_token_; } std::string HuggingFaceTokenizerJson::decode_ids(const std::vector & ids) const { - return decode_with_optional_skip(id_to_token_, ids, metaspace_replacement_, trim_leading_space_, nullptr); + return decode_with_optional_skip(id_to_token_, ids, metaspace_replacement_, trim_leading_space_, byte_level_, nullptr); } std::string HuggingFaceTokenizerJson::decode_ids(const std::vector & ids, int32_t skip_token_id) const { - return decode_with_optional_skip(id_to_token_, ids, metaspace_replacement_, trim_leading_space_, &skip_token_id); + return decode_with_optional_skip(id_to_token_, ids, metaspace_replacement_, trim_leading_space_, byte_level_, &skip_token_id); } std::shared_ptr load_huggingface_tokenizer_json( @@ -114,10 +204,13 @@ std::shared_ptr load_huggingface_tokenizer_json( trim_leading_space = true; } + const bool byte_level = has_byte_level(root); + return std::make_shared( std::move(id_to_token), std::move(metaspace_replacement), - trim_leading_space); + trim_leading_space, + byte_level); } } // namespace engine::tokenizers diff --git a/tests/granite5asr/granite5asr_warm_bench.cpp b/tests/granite5asr/granite5asr_warm_bench.cpp new file mode 100644 index 000000000..cafee9f54 --- /dev/null +++ b/tests/granite5asr/granite5asr_warm_bench.cpp @@ -0,0 +1,18 @@ +#include "../core/audio_task_warm_bench.h" + +int main(int argc, char ** argv) { + try { + engine::tools::AudioTaskBenchConfig config; + config.family = "granite5asr"; + config.default_model = "granite5asr"; + config.task = engine::runtime::VoiceTaskKind::Asr; + config.output_kind = engine::tools::AudioTaskOutputKind::Asr; + return engine::tools::run_audio_task_warm_bench( + argc, + argv, + config); + } catch (const std::exception & ex) { + std::cerr << "granite5asr_warm_bench failed: " << ex.what() << "\n"; + return 1; + } +} diff --git a/tests/granite5asr/test_granite5asr_golden_transcription.cpp b/tests/granite5asr/test_granite5asr_golden_transcription.cpp new file mode 100644 index 000000000..5128a605a --- /dev/null +++ b/tests/granite5asr/test_granite5asr_golden_transcription.cpp @@ -0,0 +1,136 @@ +#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 = "concord returned to its place amidst the tents"; + +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("granite5asr").string()); + const std::filesystem::path audio_path = arg_value( + argc, argv, "--audio", repo_path("assets/asr_validation/librispeech/librispeech_test_clean_6930-75918-0000.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_granite5asr_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 = "granite5asr"; + 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["granite5asr.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: Granite 5 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"; + + if (actual != expected) { + std::cerr << "FAIL: transcript mismatch!\n"; + return kExitFail; + } + + std::cout << "PASS: Granite 5 ASR golden transcription verified successfully.\n"; + return kExitPass; + } catch (const std::exception & ex) { + std::cerr << "FAIL: exception: " << ex.what() << "\n"; + return kExitFail; + } +} diff --git a/tools/community_models/convert_granite5asr.py b/tools/community_models/convert_granite5asr.py new file mode 100644 index 000000000..f8a902b6a --- /dev/null +++ b/tools/community_models/convert_granite5asr.py @@ -0,0 +1,116 @@ +#!/usr/bin/env python3 +"""Convert IBM Granite Speech 5.0 470M TurboCTC safetensors checkpoints to audio.cpp GGUF packages. + +Produces a self-contained GGUF with weights from model.safetensors and copies +tokenizer.json and config.json alongside it, giving a complete model directory +that audiocpp_cli / audiocpp_server load with --family granite5asr. + +Examples: + python tools/community_models/convert_granite5asr.py \ + --checkpoint granite5asr \ + --converter build/windows-cpu-release/bin/audiocpp_gguf.exe \ + --type q8_0 \ + --output models/granite5asr-gguf/granite5asr-q8_0.gguf +""" + +import argparse +import shutil +import subprocess +import sys +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[2] +SPEC = REPO_ROOT / "model_specs" / "granite5asr.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}") + ckpt = candidates[-1] + + output.parent.mkdir(parents=True, exist_ok=True) + + command = [ + str(converter), + "--input", + str(ckpt), + "--root", + str(checkpoint), + "--family", + "granite5asr", + "--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 ["tokenizer.json", "config.json", "preprocessor_config.json"]: + src_asset = checkpoint / asset_name + if src_asset.exists(): + dst_asset = output.parent / asset_name + 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 / "granite5asr", + help="checkpoint directory containing model.safetensors and tokenizer.json", + ) + parser.add_argument( + "--converter", + type=Path, + required=True, + help="path to the audiocpp_gguf binary", + ) + parser.add_argument( + "--output", + type=Path, + default=Path("gguf-out/granite-speech-5.0-470m-turboctc-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() diff --git a/webui/configs/models_catalog.json b/webui/configs/models_catalog.json index 45e2e2103..2a3417055 100644 --- a/webui/configs/models_catalog.json +++ b/webui/configs/models_catalog.json @@ -105,6 +105,9 @@ { "id": "kroko-asr", "display_name": "Kroko Community ASR (asr, GGUF Q8)", "display_name_en": "Kroko Community ASR (asr, GGUF Q8)", "family": "kroko_asr", "path": "models/Kroko-ASR-GGUF", "task": "asr", "mode": "offline", "download_id": "kroko_asr_community_q8_0", "min_vram_gb": 4, "input_hint": "**Kroko Community ASR**：GGUF Q8 包；离线转写，支持时间戳。", "input_hint_en": "**Kroko Community ASR**: GGUF Q8 package for offline transcription with timestamps." }, + { "id": "granite5asr", "display_name": "Granite Speech 5.0 470M TurboCTC (asr)", "display_name_en": "Granite Speech 5.0 470M TurboCTC (asr)", "family": "granite5asr", "path": "granite5asr", "task": "asr", "mode": "offline", "download_id": "granite5asr_q8_0", "min_vram_gb": 4, + "input_hint": "**Granite Speech 5.0 TurboCTC**：IBM 470M 英语 ASR；超快 Conformer CTC 转写；支持长音频自动分段与流式模式。", + "input_hint_en": "**Granite Speech 5.0 TurboCTC**: IBM 470M English ASR with ultra-fast Conformer CTC architecture, supporting long-form audio segmentation and streaming mode." }, { "id": "sense-asr", "display_name": "SenseVoice-Small (asr, 流式, 社区)", "display_name_en": "SenseVoice-Small (asr + streaming, community)", "family": "sense_asr", "path": "models/SenseVoice-Small-GGUF", "task": "asr", "mode": "offline", "download_id": "sensevoice_small_q8", "min_vram_gb": 4, "input_hint": "**SenseVoice-Small**（社区模型）：多语种 ASR，事件/情感/语言标签，ITN 可开关；离线与流式模式。", "input_hint_en": "**SenseVoice-Small** (community): multilingual ASR with event/emotion/language tags, optional ITN; offline and streaming modes." },