diff --git a/docs/models/arkasr.md b/docs/models/arkasr.md new file mode 100644 index 000000000..998630465 --- /dev/null +++ b/docs/models/arkasr.md @@ -0,0 +1,63 @@ +# ARK-ASR candidate port + +This candidate reuses the ARK implementation and converter published by +harshav in [ARK-ASR-3B-GGUF](https://huggingface.co/harshav/ARK-ASR-3B-GGUF). +The source archive `transcribe.cpp-arkasr-b838a2b.tar.gz` has SHA-256 +`776c2a7a690dc21dc629aa79d03003a0a9f9a463f1e0a4f3a4cefc6ee59dbe97`. +The runtime additions are MIT licensed; original weights are Apache-2.0. +The original model and inference implementation are maintained by +[AutoArk](https://github.com/AutoArk/open-audio-opd). + +The source was adapted to the current device-selection API and gate/up packing +status contract. Optional Q/K/V projection biases were added to the shared +causal decoder, preserving null defaults for other families. The converter +now records variant identity, supported languages and automatic language +detection, accepts a pinned source revision, and checks variant dimensions. + +## Variants + +`ark-asr-0.6b` uses a 896-dimension Qwen2 decoder with 24 layers, 14 attention +heads and 2 KV heads. `ark-asr-3b` uses the existing 2048-dimension decoder. +Both have a Whisper-style encoder with partial interleaved RoPE and a merge-4 +MLP adapter. The 0.6B name describes the decoder, not the full checkpoint. + +## Conversion + +```sh +uv run scripts/convert-arkasr-to-gguf.py \ + --input Edge0/ARK-ASR-0.6B \ + --revision 45776b56d58cdfb2e2eb632f7e110f38684633e0 \ + --variant ark-asr-0.6b --outtype q8_0 \ + --output ark-asr-0.6b-Q8_0.gguf +``` + +The reused converter's Q8 path retains the tied embedding/output matrix and +convolutions in F16, and norms/biases/frontend tensors in F32. This is a +candidate conversion path; moving quantization to the standard C++ policy is +part of upstream acceptance work. + +## Runtime contract + +16 kHz mono PCM; offline automatic-language transcription; cancellation; +no timestamps, translation, streaming or recognition hints. A language option +does not steer this port. Long audio is processed in windows of at most 30 +seconds, with low-energy boundaries, overlap and transcript stitching. + +```sh +build/bin/transcribe-cli --backend metal -m ark-asr-0.6b-Q8_0.gguf samples/jfk.wav +build/bin/transcribe-cli --backend cpu -m ark-asr-0.6b-Q8_0.gguf samples/jfk.wav +uv run scripts/verify-arkasr-reference.py --model /path/to/pinned/checkpoint \ + --audio samples/jfk.wav --device mps +``` + +## Acceptance status + +This is a draft candidate for +[issue 190](https://github.com/handy-computer/transcribe.cpp/issues/190). +It is not a canonical published model family. The existing default test suite, +real-model transcription comparisons and original-checkpoint oracle provide +initial evidence. The intake signoff, standard golden tensor manifest/dumper, +tolerance profile, family fixture/structural/ABI smokes, full numerical gates, +WER dataset evaluation and canonical artifact publication remain required +before upstream acceptance/publication. Transcript comparisons must not be +reported as human-ground-truth WER. diff --git a/scripts/convert-arkasr-to-gguf.py b/scripts/convert-arkasr-to-gguf.py new file mode 100644 index 000000000..4bf7ca9e1 --- /dev/null +++ b/scripts/convert-arkasr-to-gguf.py @@ -0,0 +1,369 @@ +#!/usr/bin/env python3 +# /// script +# requires-python = ">=3.12,<3.15" +# dependencies = [ +# "gguf", +# "huggingface-hub", +# "numpy", +# "safetensors", +# "torch", +# "transformers", +# ] +# /// +"""Convert ARK-ASR 0.6B or 3B to GGUF format for transcribe.cpp. + +Original model: ARK-ASR / AutoArk, by Yu Lin, Yiming Wang, Runyuan Cai, +and Xiaodong Zeng. https://github.com/AutoArk/open-audio-opd +Reused GGUF converter: harshav, published with ARK-ASR-3B-GGUF. +https://huggingface.co/harshav/ARK-ASR-3B-GGUF +Native engine: The transcribe.cpp authors (MIT). +0.6B conversion/metadata adaptation: maxffarrell, with Codex assistance. +This converter is MIT licensed; original model weights are Apache-2.0. + +Architecture (see PLAN.md §ARK): + * Audio encoder: Whisper-large-v3 (128 mel, d_model 1280, 32 layers, 20 heads) + with *partial interleaved RoPE* (rot_dim 32 of head_dim 64, theta 10000, + Q+K only); k_proj has no bias; the encoder's own final LayerNorm is dropped + (nn.Identity). + * Adapter: LayerNorm(1280) -> merge 4 consecutive frames (->5120) -> + Linear 5120->4096 -> GELU -> Linear 4096->2048. + * Decoder: stock Qwen2.5-3B (2048 hid, 36 L, GQA 16/2, SwiGLU, RMSNorm, + rope_theta 1e6, vocab 151936, tied embeddings). + +Audio placeholder token <|audio|> (151663) embeddings are replaced by the first +N = ((mel_frames+1)//2)//4 adapter frames. + +Streams tensors one-at-a-time via safe_open (BF16->F16). ~9 GB peak RAM. + +Usage: + uv run --python 3.12 scripts/convert-arkasr-to-gguf.py --input AutoArk-AI/ARK-ASR-3B \ + --output arkasr-3b-f16.gguf --outtype f16 +""" + +import argparse +import json +import sys +from pathlib import Path + +import numpy as np +import torch + +try: + from gguf import GGUFWriter, GGMLQuantizationType +except ImportError: + print("Error: gguf package not found. Install with: pip install gguf") + sys.exit(1) + +try: + from safetensors import safe_open +except ImportError: + print("Error: safetensors package not found. Install with: pip install safetensors") + sys.exit(1) + +try: + from huggingface_hub import snapshot_download +except ImportError: + snapshot_download = None + + +def load_model_dir(model_id: str, revision: str | None = None) -> Path: + p = Path(model_id) + if p.is_dir(): + return p + if snapshot_download is None: + print("Error: huggingface_hub not installed and --input is not a local dir") + sys.exit(1) + print(f"Downloading model from HuggingFace: {model_id}") + return Path(snapshot_download( + model_id, revision=revision, + allow_patterns=["*.safetensors", "*.json", "merges.txt", "*.jinja"], + )) + + +def map_tensor_name(name: str): + """Map an HF tensor name to a GGUF name, or None to skip it.""" + # --- decoder (Qwen2.5-3B) --- + if name == "model.embed_tokens.weight": + return "dec.embed.weight" + if name == "model.norm.weight": + return "dec.norm.weight" + if name == "lm_head.weight": + return None # tied to dec.embed.weight; reused at runtime + if name.startswith("model.layers."): + rest = name[len("model.layers."):] + idx, _, sub = rest.partition(".") + m = { + "self_attn.q_proj.weight": "attn.q.weight", + "self_attn.q_proj.bias": "attn.q.bias", + "self_attn.k_proj.weight": "attn.k.weight", + "self_attn.k_proj.bias": "attn.k.bias", + "self_attn.v_proj.weight": "attn.v.weight", + "self_attn.v_proj.bias": "attn.v.bias", + "self_attn.o_proj.weight": "attn.o.weight", + "input_layernorm.weight": "attn_norm.weight", + "post_attention_layernorm.weight": "ffn_norm.weight", + "mlp.gate_proj.weight": "ffn.gate.weight", + "mlp.up_proj.weight": "ffn.up.weight", + "mlp.down_proj.weight": "ffn.down.weight", + }.get(sub) + return f"dec.blk.{idx}.{m}" if m else None + + # --- adapter --- + if name == "audio_encoder.layer_norm.weight": + return "adapter.ln.weight" + if name == "audio_encoder.layer_norm.bias": + return "adapter.ln.bias" + if name == "audio_encoder.adapting.0.weight": + return "adapter.fc1.weight" + if name == "audio_encoder.adapting.0.bias": + return "adapter.fc1.bias" + if name == "audio_encoder.adapting.2.weight": + return "adapter.fc2.weight" + if name == "audio_encoder.adapting.2.bias": + return "adapter.fc2.bias" + + # --- whisper encoder --- + if name.startswith("audio_encoder.whisper."): + rest = name[len("audio_encoder.whisper."):] + if rest == "conv1.weight": + return "enc.conv1.weight" + if rest == "conv1.bias": + return "enc.conv1.bias" + if rest == "conv2.weight": + return "enc.conv2.weight" + if rest == "conv2.bias": + return "enc.conv2.bias" + if rest == "embed_positions.weight": + return None # unused (use_rope=True) + if rest in ("layer_norm.weight", "layer_norm.bias"): + return None # encoder final LN replaced by nn.Identity + if rest.startswith("layers."): + r2 = rest[len("layers."):] + idx, _, sub = r2.partition(".") + m = { + "self_attn.q_proj.weight": "attn.q.weight", + "self_attn.q_proj.bias": "attn.q.bias", + "self_attn.k_proj.weight": "attn.k.weight", # no bias + "self_attn.v_proj.weight": "attn.v.weight", + "self_attn.v_proj.bias": "attn.v.bias", + "self_attn.out_proj.weight": "attn.o.weight", + "self_attn.out_proj.bias": "attn.o.bias", + "self_attn_layer_norm.weight": "attn_ln.weight", + "self_attn_layer_norm.bias": "attn_ln.bias", + "fc1.weight": "fc1.weight", + "fc1.bias": "fc1.bias", + "fc2.weight": "fc2.weight", + "fc2.bias": "fc2.bias", + "final_layer_norm.weight": "ffn_ln.weight", + "final_layer_norm.bias": "ffn_ln.bias", + }.get(sub) + return f"enc.blk.{idx}.{m}" if m else None + return None + return None + + +def build_mel_and_window(n_fft: int, n_mels: int, sr: int): + """Whisper feature-extractor mel filterbank (n_freqs, n_mels) + hann window. + + Mirrors transformers.WhisperFeatureExtractor: slaney mel scale + slaney + norm, 0..sr/2 Hz, and torch.hann_window(periodic=True).""" + from transformers.audio_utils import mel_filter_bank + n_freqs = 1 + n_fft // 2 + fb = mel_filter_bank( + num_frequency_bins=n_freqs, + num_mel_filters=n_mels, + min_frequency=0.0, + max_frequency=float(sr) / 2.0, + sampling_rate=sr, + norm="slaney", + mel_scale="slaney", + ) # (n_freqs, n_mels) float32 + fb = np.ascontiguousarray(fb, dtype=np.float32) + # periodic hann window of length n_fft + win = np.hanning(n_fft + 1)[:-1].astype(np.float32) + return fb, win + + +def main(): + ap = argparse.ArgumentParser(description="Convert ARK-ASR-3B to GGUF") + ap.add_argument("--input", required=True, help="HF model ID or local dir") + ap.add_argument("--output", required=True, help="Output GGUF path") + ap.add_argument("--outtype", default="f16", choices=["f16", "f32", "q8_0"], + help="dtype for 2D weights (default f16). q8_0 quantizes 2D " + "decoder/encoder/adapter weights; keeps the tied embedding, " + "conv kernels (F16) and norms/filterbank (F32).") + ap.add_argument("--variant", required=True, choices=["ark-asr-0.6b", "ark-asr-3b"]) + ap.add_argument("--revision", help="Pinned Hugging Face source revision") + args = ap.parse_args() + + mdir = load_model_dir(args.input, args.revision) + with open(mdir / "config.json", encoding="utf-8") as f: + cfg = json.load(f) + wc = cfg["whisper_config"] + expected_hidden = {"ark-asr-0.6b": 896, "ark-asr-3b": 2048}[args.variant] + if cfg["model_type"] != "arkasr" or cfg["hidden_size"] != expected_hidden: + raise ValueError("Checkpoint does not match the selected ARK variant") + if not cfg.get("tie_word_embeddings") or not cfg.get("use_rope"): + raise ValueError("ARK runtime requires tied embeddings and RoPE") + + w = GGUFWriter(args.output, "arkasr") + w.add_name(args.variant.upper()) + source_repo = "Edge0/ARK-ASR-0.6B" if args.variant == "ark-asr-0.6b" else "Edge0/ARK-ASR-3B" + w.add_string("general.author", "ARK-ASR / AutoArk: Yu Lin, Yiming Wang, Runyuan Cai, Xiaodong Zeng") + w.add_string("general.license", "apache-2.0") + w.add_string("general.license.link", "https://www.apache.org/licenses/LICENSE-2.0") + w.add_string("general.source.url", "https://huggingface.co/" + source_repo) + w.add_string("general.source.huggingface.repository", source_repo) + w.add_string("arkasr.attribution.original_project", "https://github.com/AutoArk/open-audio-opd") + w.add_string("arkasr.attribution.paper", "https://arxiv.org/abs/2605.28139") + w.add_string("arkasr.attribution.reused_converter_author", "harshav") + w.add_string("arkasr.attribution.reused_converter_url", "https://huggingface.co/harshav/ARK-ASR-3B-GGUF") + w.add_string("arkasr.attribution.native_engine", "The transcribe.cpp authors; https://github.com/handy-computer/transcribe.cpp; MIT") + w.add_string("arkasr.attribution.community_conversion", "maxffarrell: 0.6B adaptation, quantization, validation and packaging; Codex assistance") + w.add_string("arkasr.attribution.notice", "Community format conversion only; no model retraining; no original-author endorsement. See accompanying NOTICE and ATTRIBUTION.md.") + if args.revision: + w.add_string("general.source.huggingface.revision", args.revision) + w.add_string("stt.variant", args.variant) + w.add_bool("stt.capability.lang_detect", True) + w.add_array("general.languages", ["zh", "en", "de", "ja", "fr", "ko", "es", "pl", "it", "ro", "hu", "cs", "nl", "fi", "hr", "sk", "sl", "et", "lt"] if args.variant == "ark-asr-0.6b" else ["zh", "en", "de", "ja", "fr", "ko", "es"]) + w.add_description("ARK-ASR (Whisper-RoPE encoder + Qwen2 decoder)") + + # ---- decoder (Qwen2.5-3B) hparams ---- + hid = int(cfg["hidden_size"]) + n_heads = int(cfg["num_attention_heads"]) + w.add_uint32("arkasr.llm.hidden_size", hid) + w.add_uint32("arkasr.llm.num_layers", int(cfg["num_hidden_layers"])) + w.add_uint32("arkasr.llm.num_heads", n_heads) + w.add_uint32("arkasr.llm.num_kv_heads", int(cfg["num_key_value_heads"])) + w.add_uint32("arkasr.llm.head_dim", hid // n_heads) + w.add_uint32("arkasr.llm.intermediate_size", int(cfg["intermediate_size"])) + w.add_uint32("arkasr.llm.vocab_size", int(cfg["vocab_size"])) + w.add_uint32("arkasr.llm.max_position_embeddings", int(cfg.get("max_position_embeddings", 32768))) + w.add_float32("arkasr.llm.rope_theta", float(cfg.get("rope_theta", 1000000.0))) + w.add_float32("arkasr.llm.rms_norm_eps", float(cfg.get("rms_norm_eps", 1e-6))) + + # ---- whisper encoder hparams ---- + d_model = int(wc["d_model"]) + e_heads = int(wc["encoder_attention_heads"]) + e_head_dim = d_model // e_heads + w.add_uint32("arkasr.whisper.d_model", d_model) + w.add_uint32("arkasr.whisper.num_layers", int(wc["encoder_layers"])) + w.add_uint32("arkasr.whisper.num_heads", e_heads) + w.add_uint32("arkasr.whisper.head_dim", e_head_dim) + w.add_uint32("arkasr.whisper.ffn_dim", int(wc["encoder_ffn_dim"])) + w.add_uint32("arkasr.whisper.num_mel_bins", int(wc["num_mel_bins"])) + w.add_uint32("arkasr.whisper.max_source_positions", int(wc["max_source_positions"])) + # RotaryEmbedding(dim = head_dim // 2), base 10000 (rope_ratio 1) + w.add_uint32("arkasr.whisper.rot_dim", e_head_dim // 2) + w.add_float32("arkasr.whisper.rope_theta", 10000.0) + w.add_float32("arkasr.whisper.ln_eps", 1e-5) + + # ---- adapter / audio / tokenizer hparams ---- + w.add_uint32("arkasr.adapter.merge_factor", int(cfg.get("merge_factor", 4))) + w.add_uint32("arkasr.audio_token_id", int(cfg["audio_token_id"])) + w.add_uint32("arkasr.bos_token_id", int(cfg.get("bos_token_id", 151643))) + w.add_uint32("arkasr.eos_token_id", int(cfg.get("eos_token_id", 151645))) + w.add_uint32("arkasr.pad_token_id", int(cfg.get("pad_token_id", 151643))) + w.add_uint32("arkasr.n_fft", 400) + w.add_uint32("arkasr.hop_length", 160) + w.add_uint32("arkasr.sample_rate", 16000) + + # ---- mel filterbank + window ---- + fb, win = build_mel_and_window(400, int(wc["num_mel_bins"]), 16000) + w.add_tensor("enc.mel_filters", fb, raw_dtype=GGMLQuantizationType.F32) + w.add_tensor("enc.mel_window", win, raw_dtype=GGMLQuantizationType.F32) + + # ---- tokenizer: GPT-2 byte-level BPE (vocab + merges + specials) ---- + vocab_size = int(cfg["vocab_size"]) + tokens = [f"[PAD{i}]" for i in range(vocab_size)] + vj = mdir / "vocab.json" + if vj.exists(): + with open(vj, encoding="utf-8") as f: + vmap = json.load(f) + for tok, idx in vmap.items(): + if 0 <= idx < vocab_size: + tokens[idx] = tok + aj = mdir / "added_tokens.json" + if aj.exists(): + with open(aj, encoding="utf-8") as f: + for tok, idx in json.load(f).items(): + if 0 <= idx < vocab_size: + tokens[idx] = tok + w.add_tokenizer_model("gpt2") + w.add_token_list(tokens) + mt = mdir / "merges.txt" + if mt.exists(): + merges = [] + with open(mt, encoding="utf-8") as f: + for line in f: + line = line.rstrip("\n") + if line and not line.startswith("#"): + merges.append(line) + if merges: + w.add_token_merges(merges) + else: + print("WARNING: merges.txt missing — decode-only is fine, encode falls back to per-byte") + + # ---- stream tensors ---- + index_path = mdir / "model.safetensors.index.json" + if index_path.exists(): + with open(index_path, encoding="utf-8") as f: + weight_map = json.load(f)["weight_map"] + else: + # single-file checkpoint + only = next(mdir.glob("*.safetensors")) + weight_map = {} + with safe_open(str(only), framework="pt") as h: + for k in h.keys(): + weight_map[k] = only.name + + handles = {} + def get_handle(fname): + if fname not in handles: + handles[fname] = safe_open(str(mdir / fname), framework="pt") + return handles[fname] + + # q8_0: keep these at F16/F32 (sampling-critical / conv / filterbank). + def keep_unquantized(name): + return (name == "dec.embed.weight" or name.endswith("conv1.weight") or name.endswith("conv2.weight")) + + try: + from gguf import quants as _gguf_quants + except ImportError: + _gguf_quants = None + + n_written = 0 + n_quant = 0 + for hf_name in sorted(weight_map.keys()): + gguf_name = map_tensor_name(hf_name) + if gguf_name is None: + continue + t = get_handle(weight_map[hf_name]).get_tensor(hf_name) + if t.ndim <= 1: + data = np.ascontiguousarray(t.to(torch.float32).numpy()) + w.add_tensor(gguf_name, data, raw_dtype=GGMLQuantizationType.F32) + elif args.outtype == "q8_0" and t.ndim == 2 and not keep_unquantized(gguf_name) \ + and (t.shape[-1] % 32 == 0) and _gguf_quants is not None: + arr = np.ascontiguousarray(t.to(torch.float32).numpy()) + qd = _gguf_quants.quantize(arr, GGMLQuantizationType.Q8_0) + w.add_tensor(gguf_name, qd, raw_dtype=GGMLQuantizationType.Q8_0) + n_quant += 1 + else: + dt = torch.float32 if args.outtype == "f32" else torch.float16 + qt = GGMLQuantizationType.F32 if args.outtype == "f32" else GGMLQuantizationType.F16 + data = np.ascontiguousarray(t.to(dt).numpy()) + w.add_tensor(gguf_name, data, raw_dtype=qt) + n_written += 1 + if n_written % 50 == 0: + print(f" ... {n_written} tensors ({n_quant} q8_0)") + + print(f"Writing GGUF ({n_written} tensors) -> {args.output}") + w.write_header_to_file() + w.write_kv_data_to_file() + w.write_tensors_to_file() + w.close() + print("Done.") + + +if __name__ == "__main__": + main() diff --git a/scripts/verify-arkasr-reference.py b/scripts/verify-arkasr-reference.py new file mode 100644 index 000000000..a2d752fc1 --- /dev/null +++ b/scripts/verify-arkasr-reference.py @@ -0,0 +1,84 @@ +# /// script +# requires-python = ">=3.12,<3.15" +# dependencies = [ +# "librosa", +# "soundfile", +# "torch", +# "transformers", +# ] +# /// +"""Run the official ARK-ASR checkpoint as an independent parity oracle.""" + +import argparse + +import torch +from transformers import AutoModelForCausalLM, AutoProcessor, AutoTokenizer + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--model", required=True) + parser.add_argument("--audio", required=True) + parser.add_argument("--device", choices=("auto", "cpu", "mps"), default="auto") + args = parser.parse_args() + + device = args.device + if device == "auto": + device = "mps" if torch.backends.mps.is_available() else "cpu" + dtype = torch.bfloat16 if device == "mps" else torch.float32 + + processor = AutoProcessor.from_pretrained(args.model, trust_remote_code=True) + tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True) + model = AutoModelForCausalLM.from_pretrained( + args.model, + trust_remote_code=True, + dtype=dtype, + attn_implementation="sdpa", + ).to(device) + model.eval() + + conversation = [{ + "role": "user", + "content": [ + {"type": "audio", "path": args.audio}, + {"type": "text", "text": "Please transcribe this audio."}, + ], + }] + inputs = processor.apply_chat_template( + conversation, + add_generation_prompt=True, + return_tensors="pt", + sampling_rate=16000, + audio_padding="longest", + text_kwargs={"padding": "longest"}, + audio_max_length=30 * 16000, + ).to(device) + if "audios" in inputs: + inputs["audios"] = inputs["audios"].to(dtype=dtype) + + eos_ids = tokenizer.eos_token_id + keep_ids = {eos_ids} if isinstance(eos_ids, int) else set(eos_ids or []) + bad_ids = set(tokenizer.all_special_ids) - keep_ids + bad_ids.update( + token_id + for token, token_id in tokenizer.get_added_vocab().items() + if token.startswith("<") and token.endswith(">") and token_id not in keep_ids + ) + + with torch.inference_mode(): + outputs = model.generate( + **inputs, + do_sample=False, + max_new_tokens=256, + pad_token_id=tokenizer.pad_token_id, + eos_token_id=tokenizer.eos_token_id, + bad_words_ids=[[token_id] for token_id in sorted(bad_ids)], + ) + print(tokenizer.batch_decode( + outputs[:, inputs.input_ids.shape[1]:], + skip_special_tokens=True, + )[0]) + + +if __name__ == "__main__": + main() diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 5b0931f71..706db0bdd 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -55,6 +55,11 @@ add_library(transcribe arch/qwen3_asr/weights.cpp arch/qwen3_asr/encoder.cpp arch/qwen3_asr/decoder.cpp + arch/arkasr/model.cpp + arch/arkasr/capabilities.cpp + arch/arkasr/weights.cpp + arch/arkasr/encoder.cpp + arch/arkasr/decoder.cpp arch/moss/model.cpp arch/moss/capabilities.cpp arch/moss/diarize.cpp diff --git a/src/arch/arkasr/arkasr.h b/src/arch/arkasr/arkasr.h new file mode 100644 index 000000000..eedd26ede --- /dev/null +++ b/src/arch/arkasr/arkasr.h @@ -0,0 +1,56 @@ +#pragma once + +#include "causal_lm/causal_lm.h" +#include "transcribe-backend.h" +#include "transcribe-mel.h" +#include "transcribe-model.h" +#include "transcribe-session.h" +#include "transcribe-tokenizer.h" +#include "weights.h" + +#include +#include + +namespace transcribe::arkasr { + +void apply_family_invariants(transcribe_model & model); + +struct SpecialTokens { + int32_t user = -1; + int32_t begin_audio = -1; + int32_t audio = -1; + int32_t end_audio = -1; + int32_t assistant = -1; + int32_t eos = -1; +}; + +struct ArkAsrModel final : public transcribe_model { + Tokenizer tok; + ArkAsrHParams hparams; + ArkAsrWeights weights; + SpecialTokens special; + std::vector instruction_tokens; + ggml_context * ctx_meta = nullptr; + + BackendPlan plan; + ggml_backend_buffer_t backend_buffer = nullptr; + causal_lm::PackedGateUpHandles packed_gate_up; + std::optional mel; + + ~ArkAsrModel() override; + + const Tokenizer * tokenizer() const override { return &tok; } +}; + +struct ArkAsrSession final : public transcribe_session { + ggml_context * compute_ctx = nullptr; + ggml_backend_sched_t sched = nullptr; + causal_lm::KvCache kv_cache; + std::vector mel_buf; + std::vector audio_buf; + bool decoder_use_flash = true; + + ~ArkAsrSession() override; +}; + +} // namespace transcribe::arkasr diff --git a/src/arch/arkasr/capabilities.cpp b/src/arch/arkasr/capabilities.cpp new file mode 100644 index 000000000..adfae1c61 --- /dev/null +++ b/src/arch/arkasr/capabilities.cpp @@ -0,0 +1,26 @@ +// arch/arkasr/capabilities.cpp - family invariants. + +#include "arkasr.h" + +namespace transcribe::arkasr { + +void apply_family_invariants(transcribe_model & model) { + transcribe_capabilities & caps = model.caps; + + caps.native_sample_rate = 16000; + + // Transcript-only chat response; no timestamps (the sibling + // Qwen3-ForcedAligner would be its own family). + caps.max_timestamp_kind = TRANSCRIBE_TIMESTAMPS_NONE; + + // The fixed transcription prompt lets the model infer language, but this + // port does not yet consume a caller-supplied language hint. Translation + // is not advertised. + caps.supports_translate = false; + + // Cancellation is wired at the per-run level. No PNC/ITN toggle; the + // Whisper-specific features do not apply here. + transcribe::set_feature(&model, TRANSCRIBE_FEATURE_CANCELLATION, true); +} + +} // namespace transcribe::arkasr diff --git a/src/arch/arkasr/decoder.cpp b/src/arch/arkasr/decoder.cpp new file mode 100644 index 000000000..2ec910a5c --- /dev/null +++ b/src/arch/arkasr/decoder.cpp @@ -0,0 +1,146 @@ +#include "decoder.h" + +#include "ggml.h" +#include "transcribe-log.h" + +namespace transcribe::arkasr { + +namespace { + +causal_lm::BlockView block_view(const ArkAsrDecBlock & b) { + causal_lm::BlockView v{}; + v.norm_attn_w = b.attn_norm_w; + v.norm_ffn_w = b.ffn_norm_w; + v.attn_q_w = b.q_w; + v.attn_k_w = b.k_w; + v.attn_v_w = b.v_w; + v.attn_o_w = b.o_w; + v.attn_q_b = b.q_b; + v.attn_k_b = b.k_b; + v.attn_v_b = b.v_b; + v.ffn_gate_up_w = b.ffn_gate_up_w; + v.ffn_down_w = b.ffn_down_w; + return v; +} + +causal_lm::BlockParams block_params(const ArkAsrHParams & hp) { + causal_lm::BlockParams p{}; + p.n_heads = hp.dec_n_heads; + p.n_kv_heads = hp.dec_n_kv_heads; + p.head_dim = hp.dec_head_dim; + p.max_position = hp.dec_max_position; + p.rms_eps = hp.dec_rms_eps; + p.rope_theta = hp.dec_rope_theta; + return p; +} + +ggml_tensor * final_logits(ggml_context * ctx, ggml_tensor * x, const ArkAsrWeights & w, const ArkAsrHParams & hp) { + x = ggml_mul(ctx, ggml_rms_norm(ctx, x, hp.dec_rms_eps), w.dec_norm_w); + x = ggml_mul_mat(ctx, w.dec_embed_w, x); + return ggml_reshape_1d(ctx, x, hp.dec_vocab_size); +} + +} // namespace + +PrefillBuild build_prefill_graph(ggml_context * ctx, + const ArkAsrWeights & weights, + const ArkAsrHParams & hp, + transcribe::causal_lm::KvCache & kv_cache, + int prompt_len, + int audio_len, + int audio_start, + bool use_flash) { + PrefillBuild out{}; + if (ctx == nullptr || prompt_len <= 0 || audio_len <= 0 || audio_start < 0 || + audio_start + audio_len > prompt_len || prompt_len > kv_cache.n_ctx) { + log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "arkasr: invalid prefill dimensions"); + return out; + } + + out.input_ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, prompt_len); + out.audio = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hp.dec_hidden, audio_len); + out.positions = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, prompt_len); + out.mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, prompt_len, prompt_len); + for (auto * t : { out.input_ids, out.audio, out.positions, out.mask }) { + ggml_set_input(t); + } + ggml_set_name(out.input_ids, "ark.input_ids"); + ggml_set_name(out.audio, "ark.audio_embeddings"); + ggml_set_name(out.positions, "ark.positions"); + ggml_set_name(out.mask, "ark.mask"); + + out.graph = ggml_new_graph_custom(ctx, 16384, false); + if (out.graph == nullptr) { + return {}; + } + + ggml_tensor * token_embeddings = ggml_get_rows(ctx, weights.dec_embed_w, out.input_ids); + const size_t elem = ggml_element_size(token_embeddings); + ggml_tensor * x = out.audio; + if (audio_start > 0) { + auto * prefix = ggml_view_2d(ctx, token_embeddings, hp.dec_hidden, audio_start, elem * hp.dec_hidden, 0); + x = ggml_concat(ctx, ggml_cont(ctx, prefix), x, 1); + } + const int suffix_len = prompt_len - audio_start - audio_len; + if (suffix_len > 0) { + auto * suffix = ggml_view_2d(ctx, token_embeddings, hp.dec_hidden, suffix_len, elem * hp.dec_hidden, + elem * hp.dec_hidden * static_cast(audio_start + audio_len)); + x = ggml_concat(ctx, x, ggml_cont(ctx, suffix), 1); + } + + const auto params = block_params(hp); + for (int layer = 0; layer < hp.dec_n_layers; ++layer) { + causal_lm::BlockOpts opts{}; + opts.use_flash = use_flash; + opts.slice_last_before_ffn = layer == hp.dec_n_layers - 1; + x = causal_lm::block_prefill(ctx, out.graph, x, block_view(weights.dec_blocks[layer]), params, kv_cache, layer, + prompt_len, out.mask, out.positions, opts); + } + + out.logits = final_logits(ctx, x, weights, hp); + ggml_set_name(out.logits, "ark.prefill_logits"); + ggml_set_output(out.logits); + ggml_build_forward_expand(out.graph, out.logits); + return out; +} + +StepBuild build_step_graph(ggml_context * ctx, + const ArkAsrWeights & weights, + const ArkAsrHParams & hp, + transcribe::causal_lm::KvCache & kv_cache, + int max_kv, + bool use_flash) { + StepBuild out{}; + if (ctx == nullptr || max_kv <= 0 || max_kv > kv_cache.n_ctx) { + return out; + } + + out.input_id = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1); + out.position = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1); + out.kv_index = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, 1); + out.mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, max_kv, 1); + out.logit_mask = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hp.dec_vocab_size); + for (auto * t : { out.input_id, out.position, out.kv_index, out.mask, out.logit_mask }) { + ggml_set_input(t); + } + + out.graph = ggml_new_graph_custom(ctx, 8192, false); + if (out.graph == nullptr) { + return {}; + } + + ggml_tensor * x = ggml_get_rows(ctx, weights.dec_embed_w, out.input_id); + const auto params = block_params(hp); + for (int layer = 0; layer < hp.dec_n_layers; ++layer) { + x = causal_lm::block_step(ctx, out.graph, x, block_view(weights.dec_blocks[layer]), params, kv_cache, layer, + max_kv, out.mask, out.position, out.kv_index, use_flash); + } + auto * logits = final_logits(ctx, x, weights, hp); + logits = ggml_add(ctx, logits, out.logit_mask); + out.token = ggml_argmax(ctx, logits); + ggml_set_output(out.token); + ggml_build_forward_expand(out.graph, out.token); + return out; +} + +} // namespace transcribe::arkasr diff --git a/src/arch/arkasr/decoder.h b/src/arch/arkasr/decoder.h new file mode 100644 index 000000000..0a998b7e8 --- /dev/null +++ b/src/arch/arkasr/decoder.h @@ -0,0 +1,47 @@ +#pragma once + +#include "causal_lm/causal_lm.h" +#include "weights.h" + +struct ggml_context; +struct ggml_cgraph; +struct ggml_tensor; + +namespace transcribe::arkasr { + +struct PrefillBuild { + ggml_tensor * input_ids = nullptr; + ggml_tensor * audio = nullptr; + ggml_tensor * positions = nullptr; + ggml_tensor * mask = nullptr; + ggml_tensor * logits = nullptr; + ggml_cgraph * graph = nullptr; +}; + +PrefillBuild build_prefill_graph(ggml_context * ctx, + const ArkAsrWeights & weights, + const ArkAsrHParams & hp, + transcribe::causal_lm::KvCache & kv_cache, + int prompt_len, + int audio_len, + int audio_start, + bool use_flash); + +struct StepBuild { + ggml_tensor * input_id = nullptr; + ggml_tensor * position = nullptr; + ggml_tensor * kv_index = nullptr; + ggml_tensor * mask = nullptr; + ggml_tensor * logit_mask = nullptr; + ggml_tensor * token = nullptr; + ggml_cgraph * graph = nullptr; +}; + +StepBuild build_step_graph(ggml_context * ctx, + const ArkAsrWeights & weights, + const ArkAsrHParams & hp, + transcribe::causal_lm::KvCache & kv_cache, + int max_kv, + bool use_flash); + +} // namespace transcribe::arkasr diff --git a/src/arch/arkasr/encoder.cpp b/src/arch/arkasr/encoder.cpp new file mode 100644 index 000000000..cf7b6af2d --- /dev/null +++ b/src/arch/arkasr/encoder.cpp @@ -0,0 +1,119 @@ +#include "encoder.h" + +#include "ggml.h" +#include "transcribe-log.h" + +#include + +namespace transcribe::arkasr { + +namespace { + +ggml_tensor * layer_norm(ggml_context * ctx, ggml_tensor * x, ggml_tensor * weight, ggml_tensor * bias, float eps) { + return ggml_add(ctx, ggml_mul(ctx, ggml_norm(ctx, x, eps), weight), bias); +} + +ggml_tensor * conv_bias(ggml_context * ctx, ggml_tensor * bias) { + return ggml_reshape_2d(ctx, bias, 1, bias->ne[0]); +} + +ggml_tensor * encoder_attention(ggml_context * ctx, + ggml_tensor * x, + const ArkAsrEncBlock & b, + const ArkAsrHParams & hp, + ggml_tensor * positions) { + const int head_dim = hp.enc_head_dim; + const int heads = hp.enc_n_heads; + const int tokens = static_cast(x->ne[1]); + + auto * q = ggml_add(ctx, ggml_mul_mat(ctx, b.q_w, x), b.q_b); + auto * k = ggml_mul_mat(ctx, b.k_w, x); + auto * v = ggml_add(ctx, ggml_mul_mat(ctx, b.v_w, x), b.v_b); + + q = ggml_reshape_3d(ctx, q, head_dim, heads, tokens); + k = ggml_reshape_3d(ctx, k, head_dim, heads, tokens); + v = ggml_reshape_3d(ctx, v, head_dim, heads, tokens); + + q = ggml_rope_ext(ctx, q, positions, nullptr, hp.enc_rot_dim, GGML_ROPE_TYPE_NORMAL, 0, hp.enc_rope_theta, 1.0f, + 0.0f, 1.0f, 0.0f, 0.0f); + k = ggml_rope_ext(ctx, k, positions, nullptr, hp.enc_rot_dim, GGML_ROPE_TYPE_NORMAL, 0, hp.enc_rope_theta, 1.0f, + 0.0f, 1.0f, 0.0f, 0.0f); + + q = ggml_cont(ctx, ggml_permute(ctx, q, 0, 2, 1, 3)); + k = ggml_cont(ctx, ggml_permute(ctx, k, 0, 2, 1, 3)); + v = ggml_cont(ctx, ggml_permute(ctx, v, 0, 2, 1, 3)); + + auto * out = ggml_flash_attn_ext(ctx, q, k, v, nullptr, 1.0f / std::sqrt(static_cast(head_dim)), 0.0f, 0.0f); + out = ggml_reshape_2d(ctx, out, hp.enc_d_model, tokens); + return ggml_add(ctx, ggml_mul_mat(ctx, b.o_w, out), b.o_b); +} + +} // namespace + +EncoderBuild build_encoder_graph(ggml_context * ctx, + const ArkAsrWeights & w, + const ArkAsrHParams & hp, + int n_mel_frames) { + EncoderBuild out{}; + if (ctx == nullptr || n_mel_frames <= 0) { + return out; + } + + out.graph = ggml_new_graph_custom(ctx, 16384, false); + if (out.graph == nullptr) { + return {}; + } + + out.mel = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_mel_frames, hp.enc_num_mels); + ggml_set_name(out.mel, "ark.mel"); + ggml_set_input(out.mel); + + auto * x = ggml_conv_1d(ctx, w.conv1_w, out.mel, 1, 1, 1); + x = ggml_gelu_erf(ctx, ggml_add(ctx, x, conv_bias(ctx, w.conv1_b))); + x = ggml_conv_1d(ctx, w.conv2_w, x, 2, 1, 1); + x = ggml_gelu_erf(ctx, ggml_add(ctx, x, conv_bias(ctx, w.conv2_b))); + const int enc_tokens = static_cast(x->ne[0]); + x = ggml_cont(ctx, ggml_transpose(ctx, x)); + + out.positions = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, enc_tokens); + ggml_set_name(out.positions, "ark.encoder_positions"); + ggml_set_input(out.positions); + + for (int layer = 0; layer < hp.enc_n_layers; ++layer) { + const auto & b = w.enc_blocks[layer]; + auto * residual = x; + auto * h = layer_norm(ctx, x, b.attn_ln_w, b.attn_ln_b, hp.enc_ln_eps); + x = ggml_add(ctx, residual, encoder_attention(ctx, h, b, hp, out.positions)); + + residual = x; + h = layer_norm(ctx, x, b.ffn_ln_w, b.ffn_ln_b, hp.enc_ln_eps); + h = ggml_add(ctx, ggml_mul_mat(ctx, b.fc1_w, h), b.fc1_b); + h = ggml_gelu_erf(ctx, h); + h = ggml_add(ctx, ggml_mul_mat(ctx, b.fc2_w, h), b.fc2_b); + x = ggml_add(ctx, residual, h); + } + + x = layer_norm(ctx, x, w.adapter_ln_w, w.adapter_ln_b, hp.enc_ln_eps); + const int usable = (enc_tokens / hp.merge_factor) * hp.merge_factor; + if (usable <= 0) { + log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "arkasr: audio is too short after encoder subsampling"); + return {}; + } + if (usable != enc_tokens) { + x = ggml_cont(ctx, ggml_view_2d(ctx, x, hp.enc_d_model, usable, x->nb[1], 0)); + } + + out.n_audio = usable / hp.merge_factor; + x = ggml_reshape_2d(ctx, x, hp.enc_d_model * hp.merge_factor, out.n_audio); + x = ggml_add(ctx, ggml_mul_mat(ctx, w.adapter_fc1_w, x), w.adapter_fc1_b); + x = ggml_gelu_erf(ctx, x); + x = ggml_add(ctx, ggml_mul_mat(ctx, w.adapter_fc2_w, x), w.adapter_fc2_b); + + out.output = x; + ggml_set_name(out.output, "ark.audio_embeddings"); + ggml_set_output(out.output); + ggml_build_forward_expand(out.graph, out.output); + return out; +} + +} // namespace transcribe::arkasr diff --git a/src/arch/arkasr/encoder.h b/src/arch/arkasr/encoder.h new file mode 100644 index 000000000..d4da822de --- /dev/null +++ b/src/arch/arkasr/encoder.h @@ -0,0 +1,24 @@ +#pragma once + +#include "weights.h" + +struct ggml_context; +struct ggml_cgraph; +struct ggml_tensor; + +namespace transcribe::arkasr { + +struct EncoderBuild { + ggml_tensor * mel = nullptr; + ggml_tensor * positions = nullptr; + ggml_tensor * output = nullptr; + ggml_cgraph * graph = nullptr; + int n_audio = 0; +}; + +EncoderBuild build_encoder_graph(ggml_context * ctx, + const ArkAsrWeights & weights, + const ArkAsrHParams & hp, + int n_mel_frames); + +} // namespace transcribe::arkasr diff --git a/src/arch/arkasr/model.cpp b/src/arch/arkasr/model.cpp new file mode 100644 index 000000000..650d67ec4 --- /dev/null +++ b/src/arch/arkasr/model.cpp @@ -0,0 +1,672 @@ +#include "arkasr.h" +#include "decoder.h" +#include "encoder.h" +#include "ggml-backend.h" +#include "ggml.h" +#include "gguf.h" +#include "transcribe-arch.h" +#include "transcribe-batch-util.h" +#include "transcribe-flash-policy.h" +#include "transcribe-load-common.h" +#include "transcribe-loader.h" +#include "transcribe-log.h" +#include "transcribe-meta.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace transcribe::arkasr { + +extern const Arch arch; + +static_assert(std::is_base_of_v); +static_assert(std::is_base_of_v); + +ArkAsrSession::~ArkAsrSession() { + kv_cache.free(); + if (sched != nullptr) { + safe_sched_free(sched); + } + if (compute_ctx != nullptr) { + ggml_free(compute_ctx); + } +} + +ArkAsrModel::~ArkAsrModel() { + if (ctx_meta != nullptr) { + ggml_free(ctx_meta); + } + if (backend_buffer != nullptr) { + safe_buffer_free(backend_buffer); + } + packed_gate_up.free(); + for (auto it = plan.scheduler_list.rbegin(); it != plan.scheduler_list.rend(); ++it) { + safe_backend_free(*it); + } +} + +namespace { + +constexpr int kFirstSpecialToken = 151643; +constexpr int kMaxAudioSamples = 30 * 16000; +constexpr int kMinAudioSamples = 20 * 16000; +constexpr int kOverlapSamples = 2 * 16000; + +int context_ceiling(const ArkAsrSession & session, const ArkAsrHParams & hp) { + return session.n_ctx > 0 ? std::min(session.n_ctx, hp.dec_max_position) : hp.dec_max_position; +} + +bool reset_compute_context(ArkAsrSession & session) { + if (session.sched != nullptr) { + ggml_backend_sched_reset(session.sched); + } + if (session.compute_ctx != nullptr) { + ggml_free(session.compute_ctx); + } + ggml_init_params params{}; + params.mem_size = 32 * 1024 * 1024; + params.no_alloc = true; + session.compute_ctx = ggml_init(params); + return session.compute_ctx != nullptr; +} + +transcribe_status resolve_special_tokens(ArkAsrModel & model) { + struct Entry { + const char * text; + int32_t * id; + }; + + const Entry entries[] = { + { "<|user|>", &model.special.user }, + { "<|begin_of_audio|>", &model.special.begin_audio }, + { "<|end_of_audio|>", &model.special.end_audio }, + { "<|assistant|>", &model.special.assistant }, + }; + for (const auto & entry : entries) { + *entry.id = model.tok.find(entry.text); + if (*entry.id < 0) { + log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "arkasr: tokenizer missing %s", entry.text); + return TRANSCRIBE_ERR_GGUF; + } + } + model.special.audio = model.hparams.audio_token_id; + model.special.eos = model.hparams.eos_token_id; + const int audio_vocab_id = model.tok.find("<|audio|>"); + if (audio_vocab_id >= 0 && audio_vocab_id != model.special.audio) { + log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "arkasr: audio token metadata does not match tokenizer"); + return TRANSCRIBE_ERR_GGUF; + } + return TRANSCRIBE_OK; +} + +void build_prompt(const ArkAsrModel & model, int audio_tokens, std::vector & ids, int & audio_start) { + ids.clear(); + ids.reserve(static_cast(audio_tokens) + model.instruction_tokens.size() + 4); + ids.push_back(model.special.user); + ids.push_back(model.special.begin_audio); + audio_start = static_cast(ids.size()); + ids.insert(ids.end(), audio_tokens, model.special.audio); + ids.push_back(model.special.end_audio); + ids.insert(ids.end(), model.instruction_tokens.begin(), model.instruction_tokens.end()); + ids.push_back(model.special.assistant); +} + +std::vector causal_mask(int tokens) { + std::vector mask(static_cast(tokens) * tokens); + for (int q = 0; q < tokens; ++q) { + for (int k = 0; k < tokens; ++k) { + const float value = k <= q ? 0.0f : -std::numeric_limits::infinity(); + mask[static_cast(q) * tokens + k] = ggml_fp32_to_fp16(value); + } + } + return mask; +} + +std::vector step_mask(int width, int position) { + std::vector mask(width); + for (int i = 0; i < width; ++i) { + const float value = i <= position ? 0.0f : -std::numeric_limits::infinity(); + mask[i] = ggml_fp32_to_fp16(value); + } + return mask; +} + +int argmax_initial_visible(const std::vector & logits) { + int best = -1; + const int visible = std::min(kFirstSpecialToken, static_cast(logits.size())); + for (int i = 0; i < visible; ++i) { + if (best < 0 || logits[i] > logits[best]) { + best = i; + } + } + return best; +} + +void normalize_transcript(std::string & text) { + const auto trim_left = [&] { + const auto pos = text.find_first_not_of(" \t\r\n"); + text.erase(0, pos == std::string::npos ? text.size() : pos); + }; + trim_left(); + if (!text.empty() && text.front() == '.' && + (text.size() == 1 || text[1] == ' ' || text[1] == '\t' || text[1] == '\n')) { + text.erase(0, 1); + trim_left(); + } +} + +struct WindowResult { + std::string raw_text; + std::string text; + bool truncated = false; +}; + +struct WordSpan { + size_t begin = 0; + size_t end = 0; + std::string key; +}; + +std::vector word_spans(const std::string & text) { + std::vector spans; + for (size_t i = 0; i < text.size();) { + while (i < text.size() && std::isspace(static_cast(text[i]))) { + ++i; + } + const size_t begin = i; + while (i < text.size() && !std::isspace(static_cast(text[i]))) { + ++i; + } + if (begin == i) { + continue; + } + WordSpan span{ begin, i, {} }; + for (size_t j = begin; j < i; ++j) { + const unsigned char ch = static_cast(text[j]); + if (ch >= 0x80 || std::isalnum(ch)) { + span.key.push_back(ch < 0x80 ? static_cast(std::tolower(ch)) : text[j]); + } + } + spans.push_back(std::move(span)); + } + return spans; +} + +std::string append_with_word_overlap(std::string & full, const std::string & segment) { + if (segment.empty()) { + return {}; + } + if (full.empty()) { + full = segment; + return segment; + } + + const auto previous = word_spans(full); + const auto current = word_spans(segment); + const size_t limit = std::min({ 16, previous.size(), current.size() }); + size_t overlap = 0; + for (size_t count = 1; count <= limit; ++count) { + bool equal = true; + for (size_t i = 0; i < count; ++i) { + if (previous[previous.size() - count + i].key != current[i].key) { + equal = false; + break; + } + } + if (equal) { + overlap = count; + } + } + + size_t begin = overlap == 0 ? 0 : current[overlap - 1].end; + while (begin < segment.size() && std::isspace(static_cast(segment[begin]))) { + ++begin; + } + std::string appended = segment.substr(begin); + if (!appended.empty()) { + if (!std::isspace(static_cast(full.back()))) { + full.push_back(' '); + } + full += appended; + } + return appended; +} + +int choose_window_samples(const float * pcm, int remaining) { + if (remaining <= kMaxAudioSamples) { + return remaining; + } + + constexpr int kFrameSamples = 80 * 16000 / 1000; + constexpr int kHopSamples = 20 * 16000 / 1000; + double total_energy = 0.0; + for (int i = 0; i < kMaxAudioSamples; ++i) { + total_energy += static_cast(pcm[i]) * pcm[i]; + } + + double frame_energy = 0.0; + for (int i = kMinAudioSamples; i < kMinAudioSamples + kFrameSamples; ++i) { + frame_energy += static_cast(pcm[i]) * pcm[i]; + } + double best_energy = frame_energy; + int best_start = kMinAudioSamples; + for (int start = kMinAudioSamples + kHopSamples; start + kFrameSamples <= kMaxAudioSamples; start += kHopSamples) { + for (int i = start - kHopSamples; i < start; ++i) { + frame_energy -= static_cast(pcm[i]) * pcm[i]; + } + for (int i = start + kFrameSamples - kHopSamples; i < start + kFrameSamples; ++i) { + frame_energy += static_cast(pcm[i]) * pcm[i]; + } + if (frame_energy < best_energy) { + best_energy = frame_energy; + best_start = start; + } + } + + const double window_mean_square = total_energy / kMaxAudioSamples; + const double best_mean_square = best_energy / kFrameSamples; + if (window_mean_square > 0.0 && best_mean_square <= window_mean_square * 0.01) { + return best_start + kFrameSamples / 2; + } + return kMaxAudioSamples; +} + +transcribe_status load(Loader & loader, const transcribe_model_load_params * params, transcribe_model ** out_model) { + const int64_t start = ggml_time_us(); + auto model = std::make_unique(); + model->arch = &arch; + model->variant = loader.variant().empty() ? "ark-asr-3b" : loader.variant(); + apply_family_invariants(*model); + + if (const auto st = read_capability_kv(loader.gguf(), model->caps); st != TRANSCRIBE_OK) { + return st; + } + if (const auto st = read_languages_kv(loader.gguf(), *model); st != TRANSCRIBE_OK) { + return st; + } + if (const auto st = model->tok.load(loader.gguf()); st != TRANSCRIBE_OK) { + return st; + } + if (const auto st = read_arkasr_hparams(loader.gguf(), model->hparams); st != TRANSCRIBE_OK) { + return st; + } + if (model->tok.n_tokens() != model->hparams.dec_vocab_size) { + log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "arkasr: tokenizer and decoder vocabulary sizes differ"); + return TRANSCRIBE_ERR_GGUF; + } + if (const auto st = resolve_special_tokens(*model); st != TRANSCRIBE_OK) { + return st; + } + if (const auto st = model->tok.encode("Please transcribe this audio.", model->instruction_tokens); + st != TRANSCRIBE_OK) { + return st; + } + + model->caps.max_audio_ms = 0; + model->limits = {}; + + MelConfig mel_cfg{}; + mel_cfg.sample_rate = model->hparams.sample_rate; + mel_cfg.num_mels = model->hparams.enc_num_mels; + mel_cfg.n_fft = model->hparams.n_fft; + mel_cfg.win_length = model->hparams.n_fft; + mel_cfg.hop_length = model->hparams.hop_length; + mel_cfg.pre_emphasis = 0.0f; + mel_cfg.f_min = 0.0f; + mel_cfg.f_max = static_cast(model->hparams.sample_rate) / 2.0f; + mel_cfg.pad_mode = "reflect"; + mel_cfg.window_type = "hann_periodic"; + mel_cfg.normalize = "per_utterance"; + { + using R = load_common::ReadF32Result; + std::vector freq_major; + const size_t n_freq = static_cast(mel_cfg.n_fft / 2 + 1); + const size_t n_mels = static_cast(mel_cfg.num_mels); + const auto fb = load_common::read_f32_tensor_checked(loader.gguf(), loader.path(), "enc.mel_filters", + n_freq * n_mels, "arkasr", freq_major); + const auto win = load_common::read_f32_tensor_checked(loader.gguf(), loader.path(), "enc.mel_window", + mel_cfg.win_length, "arkasr", mel_cfg.window); + if (fb != R::Ok || win != R::Ok) { + return TRANSCRIBE_ERR_GGUF; + } + mel_cfg.filterbank.resize(freq_major.size()); + for (size_t f = 0; f < n_freq; ++f) { + for (size_t m = 0; m < n_mels; ++m) { + mel_cfg.filterbank[m * n_freq + f] = freq_major[f * n_mels + m]; + } + } + } + model->mel.emplace(mel_cfg); + + gguf_init_params init{}; + init.no_alloc = true; + init.ctx = &model->ctx_meta; + gguf_context * gguf = gguf_init_from_file(loader.path().c_str(), init); + if (gguf == nullptr) { + return TRANSCRIBE_ERR_GGUF; + } + if (const auto st = build_arkasr_weights(model->ctx_meta, model->hparams, model->weights); st != TRANSCRIBE_OK) { + gguf_free(gguf); + return st; + } + + const auto backend = params != nullptr ? params->backend : TRANSCRIBE_BACKEND_AUTO; + const auto device = params != nullptr ? params->device : nullptr; + if (const auto st = load_common::init_backends(backend, device, "arkasr", model->plan); st != TRANSCRIBE_OK) { + gguf_free(gguf); + return st; + } + model->backend = ggml_backend_name(model->plan.primary); + model->primary_backend = model->plan.primary; + model->backend_buffer = ggml_backend_alloc_ctx_tensors(model->ctx_meta, model->plan.primary); + if (model->backend_buffer == nullptr) { + gguf_free(gguf); + return TRANSCRIBE_ERR_OOM; + } + ggml_backend_buffer_set_usage(model->backend_buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + if (const auto st = load_common::stream_tensor_data(loader.path(), gguf, model->ctx_meta, "arkasr"); + st != TRANSCRIBE_OK) { + gguf_free(gguf); + return st; + } + gguf_free(gguf); + + std::vector entries; + entries.reserve(model->weights.dec_blocks.size()); + for (auto & block : model->weights.dec_blocks) { + entries.push_back({ block.ffn_gate_w, block.ffn_up_w, &block.ffn_gate_up_w }); + } + const auto pack_status = + causal_lm::pack_gate_up(model->plan.primary, model->hparams.dec_hidden, model->hparams.dec_intermediate, + entries, model->packed_gate_up, "arkasr"); + if (pack_status != TRANSCRIBE_OK) { + return pack_status; + } + + model->t_load_us = ggml_time_us() - start; + *out_model = model.release(); + return TRANSCRIBE_OK; +} + +transcribe_status init_context(transcribe_model * model, + const transcribe_session_params * params, + transcribe_session ** out_session) { + if (model == nullptr || model->arch != &arch || out_session == nullptr) { + return TRANSCRIBE_ERR_INVALID_ARG; + } + auto session = std::make_unique(); + session->model = model; + session->n_threads = params != nullptr ? params->n_threads : 0; + session->kv_type = params != nullptr ? params->kv_type : TRANSCRIBE_KV_TYPE_AUTO; + session->n_ctx = transcribe_session_params_n_ctx(params); + bool encoder_flash = true; + flash::apply_env_overrides(encoder_flash, session->decoder_use_flash); + *out_session = session.release(); + return TRANSCRIBE_OK; +} + +transcribe_status run_window(ArkAsrSession & session, const float * pcm, int n_samples, WindowResult & result) { + if (pcm == nullptr || n_samples <= 0 || n_samples > kMaxAudioSamples) { + return TRANSCRIBE_ERR_INVALID_ARG; + } + auto & model = *static_cast(session.model); + if (!model.mel.has_value() || model.plan.scheduler_list.empty()) { + return TRANSCRIBE_ERR_INVALID_ARG; + } + if (session.poll_abort()) { + return TRANSCRIBE_ERR_ABORTED; + } + + const int64_t mel_start = ggml_time_us(); + int n_mels = 0; + int n_frames = 0; + if (const auto st = model.mel->compute(pcm, static_cast(n_samples), session.mel_buf, n_mels, n_frames, + session.n_threads); + st != TRANSCRIBE_OK) { + return st; + } + session.t_mel_us += ggml_time_us() - mel_start; + + if (session.sched == nullptr) { + session.sched = ggml_backend_sched_new(model.plan.scheduler_list.data(), nullptr, + static_cast(model.plan.scheduler_list.size()), 16384, false, true); + if (session.sched == nullptr) { + return TRANSCRIBE_ERR_OOM; + } + } + transcribe::configure_sched_n_threads(session.sched, session.n_threads); + + if (!reset_compute_context(session)) { + return TRANSCRIBE_ERR_OOM; + } + const int64_t enc_start = ggml_time_us(); + auto encoder = build_encoder_graph(session.compute_ctx, model.weights, model.hparams, n_frames); + if (encoder.graph == nullptr || encoder.output == nullptr) { + return TRANSCRIBE_ERR_GGUF; + } + ggml_backend_sched_reset(session.sched); + if (!ggml_backend_sched_alloc_graph(session.sched, encoder.graph)) { + return TRANSCRIBE_ERR_OOM; + } + ggml_backend_tensor_set(encoder.mel, session.mel_buf.data(), 0, session.mel_buf.size() * sizeof(float)); + std::vector enc_positions(static_cast(encoder.positions->ne[0])); + for (int i = 0; i < static_cast(enc_positions.size()); ++i) { + enc_positions[i] = i; + } + ggml_backend_tensor_set(encoder.positions, enc_positions.data(), 0, enc_positions.size() * sizeof(int32_t)); + if (ggml_backend_sched_graph_compute(session.sched, encoder.graph) != GGML_STATUS_SUCCESS) { + return TRANSCRIBE_ERR_GGUF; + } + const int audio_tokens = encoder.n_audio; + session.audio_buf.resize(static_cast(model.hparams.dec_hidden) * audio_tokens); + ggml_backend_tensor_get(encoder.output, session.audio_buf.data(), 0, session.audio_buf.size() * sizeof(float)); + session.t_encode_us += ggml_time_us() - enc_start; + + const int64_t decode_start = ggml_time_us(); + std::vector prompt; + int audio_start = 0; + build_prompt(model, audio_tokens, prompt, audio_start); + const int prompt_len = static_cast(prompt.size()); + const int max_new = std::clamp(n_samples / 1280 + 64, 256, 4096); + const int ceiling = context_ceiling(session, model.hparams); + if (prompt_len + max_new > ceiling) { + return TRANSCRIBE_ERR_INPUT_TOO_LONG; + } + + int kv_width = 1; + while (kv_width < prompt_len + max_new) { + kv_width *= 2; + } + if (session.kv_cache.ctx == nullptr || session.kv_cache.n_ctx < kv_width) { + session.kv_cache.free(); + const ggml_type kv_type = session.kv_type == TRANSCRIBE_KV_TYPE_F32 ? GGML_TYPE_F32 : GGML_TYPE_F16; + if (!causal_lm::kv_init(session.kv_cache, model.plan.primary, kv_width, model.hparams.dec_n_kv_heads, + model.hparams.dec_head_dim, model.hparams.dec_n_layers, kv_type)) { + return TRANSCRIBE_ERR_OOM; + } + } else { + ggml_backend_buffer_clear(session.kv_cache.buffer, 0); + } + session.kv_cache.n = 0; + session.kv_cache.head = 0; + + if (!reset_compute_context(session)) { + return TRANSCRIBE_ERR_OOM; + } + auto prefill = build_prefill_graph(session.compute_ctx, model.weights, model.hparams, session.kv_cache, prompt_len, + audio_tokens, audio_start, session.decoder_use_flash); + if (prefill.graph == nullptr || prefill.logits == nullptr) { + return TRANSCRIBE_ERR_GGUF; + } + ggml_backend_sched_reset(session.sched); + if (!ggml_backend_sched_alloc_graph(session.sched, prefill.graph)) { + return TRANSCRIBE_ERR_OOM; + } + ggml_backend_tensor_set(prefill.input_ids, prompt.data(), 0, prompt.size() * sizeof(int32_t)); + ggml_backend_tensor_set(prefill.audio, session.audio_buf.data(), 0, session.audio_buf.size() * sizeof(float)); + std::vector positions(prompt.size()); + for (int i = 0; i < prompt_len; ++i) { + positions[i] = i; + } + const auto mask = causal_mask(prompt_len); + ggml_backend_tensor_set(prefill.positions, positions.data(), 0, positions.size() * sizeof(int32_t)); + ggml_backend_tensor_set(prefill.mask, mask.data(), 0, mask.size() * sizeof(ggml_fp16_t)); + if (ggml_backend_sched_graph_compute(session.sched, prefill.graph) != GGML_STATUS_SUCCESS) { + return TRANSCRIBE_ERR_GGUF; + } + std::vector logits(model.hparams.dec_vocab_size); + ggml_backend_tensor_get(prefill.logits, logits.data(), 0, logits.size() * sizeof(float)); + + std::vector generated; + generated.reserve(max_new); + int next = argmax_initial_visible(logits); + if (next < 0) { + return TRANSCRIBE_ERR_GGUF; + } + generated.push_back(next); + + std::vector logit_mask(static_cast(model.hparams.dec_vocab_size), 0.0f); + for (int id = kFirstSpecialToken; id < model.hparams.dec_vocab_size; ++id) { + logit_mask[id] = -std::numeric_limits::infinity(); + } + if (model.special.eos >= 0 && model.special.eos < model.hparams.dec_vocab_size) { + logit_mask[model.special.eos] = 0.0f; + } + bool reached_eos = false; + for (int count = 1; count < max_new; ++count) { + if (next == model.special.eos) { + reached_eos = true; + break; + } + if (session.poll_abort()) { + return TRANSCRIBE_ERR_ABORTED; + } + if (!reset_compute_context(session)) { + return TRANSCRIBE_ERR_OOM; + } + auto step = build_step_graph(session.compute_ctx, model.weights, model.hparams, session.kv_cache, kv_width, + session.decoder_use_flash); + if (step.graph == nullptr || step.token == nullptr) { + return TRANSCRIBE_ERR_GGUF; + } + ggml_backend_sched_reset(session.sched); + if (!ggml_backend_sched_alloc_graph(session.sched, step.graph)) { + return TRANSCRIBE_ERR_OOM; + } + ggml_backend_tensor_set(step.logit_mask, logit_mask.data(), 0, logit_mask.size() * sizeof(float)); + const int32_t input = next; + const int32_t position = prompt_len + count - 1; + const int64_t kv_index = position; + const auto smask = step_mask(kv_width, position); + ggml_backend_tensor_set(step.input_id, &input, 0, sizeof(input)); + ggml_backend_tensor_set(step.position, &position, 0, sizeof(position)); + ggml_backend_tensor_set(step.kv_index, &kv_index, 0, sizeof(kv_index)); + ggml_backend_tensor_set(step.mask, smask.data(), 0, smask.size() * sizeof(ggml_fp16_t)); + if (ggml_backend_sched_graph_compute(session.sched, step.graph) != GGML_STATUS_SUCCESS) { + return TRANSCRIBE_ERR_GGUF; + } + ggml_backend_tensor_get(step.token, &next, 0, sizeof(next)); + generated.push_back(next); + } + if (!generated.empty() && generated.back() == model.special.eos) { + generated.pop_back(); + reached_eos = true; + } + result.truncated = !reached_eos; + session.t_decode_us += ggml_time_us() - decode_start; + + result.raw_text = model.tok.decode(generated.data(), static_cast(generated.size())); + if (std::getenv("TRANSCRIBE_ARKASR_DEBUG_GEN") != nullptr) { + std::fprintf(stderr, "arkasr tokens:"); + for (const int32_t id : generated) { + std::fprintf(stderr, " %d", id); + } + std::fprintf(stderr, "\n"); + } + std::vector visible; + visible.reserve(generated.size()); + for (const int32_t id : generated) { + if (id >= 0 && id < kFirstSpecialToken) { + visible.push_back(id); + } + } + result.text = model.tok.decode(visible.data(), static_cast(visible.size())); + normalize_transcript(result.text); + return result.truncated ? TRANSCRIBE_ERR_OUTPUT_TRUNCATED : TRANSCRIBE_OK; +} + +transcribe_status run(transcribe_session * base, const float * pcm, int n_samples, const transcribe_run_params *) { + if (base == nullptr || pcm == nullptr || n_samples <= 0) { + return TRANSCRIBE_ERR_INVALID_ARG; + } + + auto & session = *static_cast(base); + auto & model = *static_cast(session.model); + if (!model.mel.has_value() || model.plan.scheduler_list.empty()) { + return TRANSCRIBE_ERR_INVALID_ARG; + } + + session.t_mel_us = 0; + session.t_encode_us = 0; + session.t_decode_us = 0; + session.was_truncated = false; + session.raw_text.clear(); + session.full_text.clear(); + session.segments.clear(); + + for (int offset = 0; offset < n_samples;) { + if (session.poll_abort()) { + return TRANSCRIBE_ERR_ABORTED; + } + const int window_samples = choose_window_samples(pcm + offset, n_samples - offset); + WindowResult window; + const auto st = run_window(session, pcm + offset, window_samples, window); + if (st != TRANSCRIBE_OK && st != TRANSCRIBE_ERR_OUTPUT_TRUNCATED) { + return st; + } + session.was_truncated = session.was_truncated || window.truncated; + if (!window.raw_text.empty()) { + if (!session.raw_text.empty()) { + session.raw_text.push_back(' '); + } + session.raw_text += window.raw_text; + } + + const std::string appended = append_with_word_overlap(session.full_text, window.text); + if (!appended.empty()) { + transcribe_session::SegmentEntry segment{}; + segment.text = appended; + segment.t0_ms = static_cast(offset) * 1000 / model.hparams.sample_rate; + segment.t1_ms = static_cast(offset + window_samples) * 1000 / model.hparams.sample_rate; + session.segments.push_back(std::move(segment)); + } + + if (offset + window_samples >= n_samples) { + break; + } + offset += window_samples - kOverlapSamples; + } + + session.result_kind = TRANSCRIBE_TIMESTAMPS_SEGMENT; + session.has_result = true; + return session.was_truncated ? TRANSCRIBE_ERR_OUTPUT_TRUNCATED : TRANSCRIBE_OK; +} + +} // namespace + +extern const Arch arch = { + "arkasr", load, init_context, run, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, +}; + +} // namespace transcribe::arkasr diff --git a/src/arch/arkasr/weights.cpp b/src/arch/arkasr/weights.cpp new file mode 100644 index 000000000..195e0a1b7 --- /dev/null +++ b/src/arch/arkasr/weights.cpp @@ -0,0 +1,172 @@ +#include "weights.h" + +#include "ggml.h" +#include "gguf.h" +#include "transcribe-log.h" +#include "transcribe-meta.h" +#include "transcribe-weights-util.h" + +namespace transcribe::arkasr { + +namespace { +constexpr const char * kTag = "arkasr"; + +#define READ_U32(key, field) \ + do { \ + if (const auto st = read_required_u32_kv(gguf, (key), kTag, hp.field); st != TRANSCRIBE_OK) \ + return st; \ + } while (0) +#define READ_F32(key, field) \ + do { \ + if (const auto st = read_required_f32_kv(gguf, (key), kTag, hp.field); st != TRANSCRIBE_OK) \ + return st; \ + } while (0) + +#define GET_F32(slot, name, ...) \ + do { \ + auto * t = transcribe::weights::find_tensor(ctx, (name), { GGML_TYPE_F32 }, { __VA_ARGS__ }, kTag); \ + if (t == nullptr) \ + return TRANSCRIBE_ERR_GGUF; \ + (slot) = t; \ + } while (0) +#define GET_CONV(slot, name, ...) \ + do { \ + auto * t = \ + transcribe::weights::find_tensor(ctx, (name), { TRANSCRIBE_QUANT_CONV_TYPES }, { __VA_ARGS__ }, kTag); \ + if (t == nullptr) \ + return TRANSCRIBE_ERR_GGUF; \ + (slot) = t; \ + } while (0) +#define GET_LIN(slot, name, ...) \ + do { \ + auto * t = \ + transcribe::weights::find_tensor(ctx, (name), { TRANSCRIBE_QUANT_LINEAR_TYPES }, { __VA_ARGS__ }, kTag); \ + if (t == nullptr) \ + return TRANSCRIBE_ERR_GGUF; \ + (slot) = t; \ + } while (0) +} // namespace + +transcribe_status read_arkasr_hparams(const gguf_context * gguf, ArkAsrHParams & hp) { + if (gguf == nullptr) { + return TRANSCRIBE_ERR_INVALID_ARG; + } + + READ_U32("arkasr.llm.hidden_size", dec_hidden); + READ_U32("arkasr.llm.num_layers", dec_n_layers); + READ_U32("arkasr.llm.num_heads", dec_n_heads); + READ_U32("arkasr.llm.num_kv_heads", dec_n_kv_heads); + READ_U32("arkasr.llm.head_dim", dec_head_dim); + READ_U32("arkasr.llm.intermediate_size", dec_intermediate); + READ_U32("arkasr.llm.vocab_size", dec_vocab_size); + READ_U32("arkasr.llm.max_position_embeddings", dec_max_position); + READ_F32("arkasr.llm.rope_theta", dec_rope_theta); + READ_F32("arkasr.llm.rms_norm_eps", dec_rms_eps); + + READ_U32("arkasr.whisper.d_model", enc_d_model); + READ_U32("arkasr.whisper.num_layers", enc_n_layers); + READ_U32("arkasr.whisper.num_heads", enc_n_heads); + READ_U32("arkasr.whisper.head_dim", enc_head_dim); + READ_U32("arkasr.whisper.ffn_dim", enc_ffn_dim); + READ_U32("arkasr.whisper.num_mel_bins", enc_num_mels); + READ_U32("arkasr.whisper.max_source_positions", enc_max_position); + READ_U32("arkasr.whisper.rot_dim", enc_rot_dim); + READ_F32("arkasr.whisper.rope_theta", enc_rope_theta); + READ_F32("arkasr.whisper.ln_eps", enc_ln_eps); + + READ_U32("arkasr.adapter.merge_factor", merge_factor); + READ_U32("arkasr.audio_token_id", audio_token_id); + READ_U32("arkasr.bos_token_id", bos_token_id); + READ_U32("arkasr.eos_token_id", eos_token_id); + READ_U32("arkasr.pad_token_id", pad_token_id); + READ_U32("arkasr.n_fft", n_fft); + READ_U32("arkasr.hop_length", hop_length); + READ_U32("arkasr.sample_rate", sample_rate); + + if (hp.dec_hidden <= 0 || hp.dec_n_layers <= 0 || hp.dec_n_heads <= 0 || hp.dec_n_kv_heads <= 0 || + hp.dec_head_dim <= 0 || hp.dec_intermediate <= 0 || hp.dec_vocab_size <= 0 || hp.dec_max_position <= 0 || + hp.dec_hidden != hp.dec_n_heads * hp.dec_head_dim || hp.dec_n_heads % hp.dec_n_kv_heads != 0) { + log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "arkasr: invalid Qwen2.5 decoder dimensions"); + return TRANSCRIBE_ERR_GGUF; + } + if (hp.enc_d_model <= 0 || hp.enc_n_layers <= 0 || hp.enc_n_heads <= 0 || hp.enc_head_dim <= 0 || + hp.enc_ffn_dim <= 0 || hp.enc_num_mels <= 0 || hp.enc_d_model != hp.enc_n_heads * hp.enc_head_dim || + hp.enc_rot_dim <= 0 || hp.enc_rot_dim > hp.enc_head_dim || (hp.enc_rot_dim % 2) != 0) { + log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "arkasr: invalid Whisper encoder dimensions"); + return TRANSCRIBE_ERR_GGUF; + } + if (hp.merge_factor <= 0 || hp.sample_rate != 16000 || hp.n_fft <= 0 || hp.hop_length <= 0 || + hp.audio_token_id < 0 || hp.eos_token_id < 0) { + log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "arkasr: invalid adapter, frontend, or token metadata"); + return TRANSCRIBE_ERR_GGUF; + } + return TRANSCRIBE_OK; +} + +transcribe_status build_arkasr_weights(ggml_context * ctx, const ArkAsrHParams & hp, ArkAsrWeights & w) { + if (ctx == nullptr) { + return TRANSCRIBE_ERR_INVALID_ARG; + } + + const int64_t ed = hp.enc_d_model; + const int64_t ef = hp.enc_ffn_dim; + const int64_t dh = hp.dec_hidden; + const int64_t di = hp.dec_intermediate; + const int64_t qd = static_cast(hp.dec_n_heads) * hp.dec_head_dim; + const int64_t kd = static_cast(hp.dec_n_kv_heads) * hp.dec_head_dim; + const int64_t ah = static_cast(hp.dec_hidden) * 2; + + GET_CONV(w.conv1_w, "enc.conv1.weight", 3, hp.enc_num_mels, ed); + GET_F32(w.conv1_b, "enc.conv1.bias", ed); + GET_CONV(w.conv2_w, "enc.conv2.weight", 3, ed, ed); + GET_F32(w.conv2_b, "enc.conv2.bias", ed); + + w.enc_blocks.assign(hp.enc_n_layers, ArkAsrEncBlock{}); + for (int i = 0; i < hp.enc_n_layers; ++i) { + auto & b = w.enc_blocks[i]; + GET_F32(b.attn_ln_w, transcribe::weights::lname("enc.blk.%d.attn_ln.weight", i), ed); + GET_F32(b.attn_ln_b, transcribe::weights::lname("enc.blk.%d.attn_ln.bias", i), ed); + GET_LIN(b.q_w, transcribe::weights::lname("enc.blk.%d.attn.q.weight", i), ed, ed); + GET_F32(b.q_b, transcribe::weights::lname("enc.blk.%d.attn.q.bias", i), ed); + GET_LIN(b.k_w, transcribe::weights::lname("enc.blk.%d.attn.k.weight", i), ed, ed); + GET_LIN(b.v_w, transcribe::weights::lname("enc.blk.%d.attn.v.weight", i), ed, ed); + GET_F32(b.v_b, transcribe::weights::lname("enc.blk.%d.attn.v.bias", i), ed); + GET_LIN(b.o_w, transcribe::weights::lname("enc.blk.%d.attn.o.weight", i), ed, ed); + GET_F32(b.o_b, transcribe::weights::lname("enc.blk.%d.attn.o.bias", i), ed); + GET_F32(b.ffn_ln_w, transcribe::weights::lname("enc.blk.%d.ffn_ln.weight", i), ed); + GET_F32(b.ffn_ln_b, transcribe::weights::lname("enc.blk.%d.ffn_ln.bias", i), ed); + GET_LIN(b.fc1_w, transcribe::weights::lname("enc.blk.%d.fc1.weight", i), ed, ef); + GET_F32(b.fc1_b, transcribe::weights::lname("enc.blk.%d.fc1.bias", i), ef); + GET_LIN(b.fc2_w, transcribe::weights::lname("enc.blk.%d.fc2.weight", i), ef, ed); + GET_F32(b.fc2_b, transcribe::weights::lname("enc.blk.%d.fc2.bias", i), ed); + } + + GET_F32(w.adapter_ln_w, "adapter.ln.weight", ed); + GET_F32(w.adapter_ln_b, "adapter.ln.bias", ed); + GET_LIN(w.adapter_fc1_w, "adapter.fc1.weight", ed * hp.merge_factor, ah); + GET_F32(w.adapter_fc1_b, "adapter.fc1.bias", ah); + GET_LIN(w.adapter_fc2_w, "adapter.fc2.weight", ah, dh); + GET_F32(w.adapter_fc2_b, "adapter.fc2.bias", dh); + + GET_LIN(w.dec_embed_w, "dec.embed.weight", dh, hp.dec_vocab_size); + GET_F32(w.dec_norm_w, "dec.norm.weight", dh); + w.dec_blocks.assign(hp.dec_n_layers, ArkAsrDecBlock{}); + for (int i = 0; i < hp.dec_n_layers; ++i) { + auto & b = w.dec_blocks[i]; + GET_F32(b.attn_norm_w, transcribe::weights::lname("dec.blk.%d.attn_norm.weight", i), dh); + GET_LIN(b.q_w, transcribe::weights::lname("dec.blk.%d.attn.q.weight", i), dh, qd); + GET_F32(b.q_b, transcribe::weights::lname("dec.blk.%d.attn.q.bias", i), qd); + GET_LIN(b.k_w, transcribe::weights::lname("dec.blk.%d.attn.k.weight", i), dh, kd); + GET_F32(b.k_b, transcribe::weights::lname("dec.blk.%d.attn.k.bias", i), kd); + GET_LIN(b.v_w, transcribe::weights::lname("dec.blk.%d.attn.v.weight", i), dh, kd); + GET_F32(b.v_b, transcribe::weights::lname("dec.blk.%d.attn.v.bias", i), kd); + GET_LIN(b.o_w, transcribe::weights::lname("dec.blk.%d.attn.o.weight", i), qd, dh); + GET_F32(b.ffn_norm_w, transcribe::weights::lname("dec.blk.%d.ffn_norm.weight", i), dh); + GET_LIN(b.ffn_gate_w, transcribe::weights::lname("dec.blk.%d.ffn.gate.weight", i), dh, di); + GET_LIN(b.ffn_up_w, transcribe::weights::lname("dec.blk.%d.ffn.up.weight", i), dh, di); + GET_LIN(b.ffn_down_w, transcribe::weights::lname("dec.blk.%d.ffn.down.weight", i), di, dh); + } + return TRANSCRIBE_OK; +} + +} // namespace transcribe::arkasr diff --git a/src/arch/arkasr/weights.h b/src/arch/arkasr/weights.h new file mode 100644 index 000000000..147719cd4 --- /dev/null +++ b/src/arch/arkasr/weights.h @@ -0,0 +1,103 @@ +#pragma once + +#include "transcribe.h" + +#include +#include + +struct gguf_context; +struct ggml_context; +struct ggml_tensor; + +namespace transcribe::arkasr { + +struct ArkAsrHParams { + int32_t dec_hidden = 0; + int32_t dec_n_layers = 0; + int32_t dec_n_heads = 0; + int32_t dec_n_kv_heads = 0; + int32_t dec_head_dim = 0; + int32_t dec_intermediate = 0; + int32_t dec_vocab_size = 0; + int32_t dec_max_position = 0; + float dec_rope_theta = 0.0f; + float dec_rms_eps = 0.0f; + + int32_t enc_d_model = 0; + int32_t enc_n_layers = 0; + int32_t enc_n_heads = 0; + int32_t enc_head_dim = 0; + int32_t enc_ffn_dim = 0; + int32_t enc_num_mels = 0; + int32_t enc_max_position = 0; + int32_t enc_rot_dim = 0; + float enc_rope_theta = 0.0f; + float enc_ln_eps = 0.0f; + + int32_t merge_factor = 0; + int32_t audio_token_id = -1; + int32_t bos_token_id = -1; + int32_t eos_token_id = -1; + int32_t pad_token_id = -1; + int32_t n_fft = 0; + int32_t hop_length = 0; + int32_t sample_rate = 0; +}; + +struct ArkAsrEncBlock { + ggml_tensor * attn_ln_w = nullptr; + ggml_tensor * attn_ln_b = nullptr; + ggml_tensor * q_w = nullptr; + ggml_tensor * q_b = nullptr; + ggml_tensor * k_w = nullptr; + ggml_tensor * v_w = nullptr; + ggml_tensor * v_b = nullptr; + ggml_tensor * o_w = nullptr; + ggml_tensor * o_b = nullptr; + ggml_tensor * ffn_ln_w = nullptr; + ggml_tensor * ffn_ln_b = nullptr; + ggml_tensor * fc1_w = nullptr; + ggml_tensor * fc1_b = nullptr; + ggml_tensor * fc2_w = nullptr; + ggml_tensor * fc2_b = nullptr; +}; + +struct ArkAsrDecBlock { + ggml_tensor * attn_norm_w = nullptr; + ggml_tensor * q_w = nullptr; + ggml_tensor * q_b = nullptr; + ggml_tensor * k_w = nullptr; + ggml_tensor * k_b = nullptr; + ggml_tensor * v_w = nullptr; + ggml_tensor * v_b = nullptr; + ggml_tensor * o_w = nullptr; + ggml_tensor * ffn_norm_w = nullptr; + ggml_tensor * ffn_gate_w = nullptr; + ggml_tensor * ffn_up_w = nullptr; + ggml_tensor * ffn_gate_up_w = nullptr; + ggml_tensor * ffn_down_w = nullptr; +}; + +struct ArkAsrWeights { + ggml_tensor * conv1_w = nullptr; + ggml_tensor * conv1_b = nullptr; + ggml_tensor * conv2_w = nullptr; + ggml_tensor * conv2_b = nullptr; + std::vector enc_blocks; + + ggml_tensor * adapter_ln_w = nullptr; + ggml_tensor * adapter_ln_b = nullptr; + ggml_tensor * adapter_fc1_w = nullptr; + ggml_tensor * adapter_fc1_b = nullptr; + ggml_tensor * adapter_fc2_w = nullptr; + ggml_tensor * adapter_fc2_b = nullptr; + + ggml_tensor * dec_embed_w = nullptr; + ggml_tensor * dec_norm_w = nullptr; + std::vector dec_blocks; +}; + +transcribe_status read_arkasr_hparams(const gguf_context * gguf, ArkAsrHParams & hp); +transcribe_status build_arkasr_weights(ggml_context * ctx, const ArkAsrHParams & hp, ArkAsrWeights & weights); + +} // namespace transcribe::arkasr diff --git a/src/causal_lm/causal_lm.cpp b/src/causal_lm/causal_lm.cpp index 4db8c0730..85fe3bca6 100644 --- a/src/causal_lm/causal_lm.cpp +++ b/src/causal_lm/causal_lm.cpp @@ -272,8 +272,17 @@ ggml_tensor * block_prefill(ggml_context * ctx, // Q/K/V projections (bias-free on Qwen3). Packing into one mul_mat // consistently regresses on Metal; left separate. ggml_tensor * Q = mul_mat_f32acc(ctx, view.attn_q_w, x_norm); + if (view.attn_q_b != nullptr) { + Q = ggml_add(ctx, Q, view.attn_q_b); + } ggml_tensor * K = mul_mat_f32acc(ctx, view.attn_k_w, x_norm); + if (view.attn_k_b != nullptr) { + K = ggml_add(ctx, K, view.attn_k_b); + } ggml_tensor * V = mul_mat_f32acc(ctx, view.attn_v_w, x_norm); + if (view.attn_v_b != nullptr) { + V = ggml_add(ctx, V, view.attn_v_b); + } Q = ggml_reshape_4d(ctx, Q, head_dim, n_heads, T_seq, 1); K = ggml_reshape_4d(ctx, K, head_dim, n_kv_heads, T_seq, 1); @@ -409,8 +418,17 @@ ggml_tensor * block_step(ggml_context * ctx, ggml_tensor * x_norm = rms_norm(ctx, x, view.norm_attn_w, rms_eps); ggml_tensor * Q = mul_mat_f32acc(ctx, view.attn_q_w, x_norm); + if (view.attn_q_b != nullptr) { + Q = ggml_add(ctx, Q, view.attn_q_b); + } ggml_tensor * K = mul_mat_f32acc(ctx, view.attn_k_w, x_norm); + if (view.attn_k_b != nullptr) { + K = ggml_add(ctx, K, view.attn_k_b); + } ggml_tensor * V = mul_mat_f32acc(ctx, view.attn_v_w, x_norm); + if (view.attn_v_b != nullptr) { + V = ggml_add(ctx, V, view.attn_v_b); + } Q = ggml_reshape_4d(ctx, Q, head_dim, n_heads, 1, 1); K = ggml_reshape_4d(ctx, K, head_dim, n_kv_heads, 1, 1); @@ -526,8 +544,17 @@ ggml_tensor * block_step_n(ggml_context * ctx, ggml_tensor * x_norm = rms_norm(ctx, x, view.norm_attn_w, rms_eps); ggml_tensor * Q = mul_mat_f32acc(ctx, view.attn_q_w, x_norm); + if (view.attn_q_b != nullptr) { + Q = ggml_add(ctx, Q, view.attn_q_b); + } ggml_tensor * K = mul_mat_f32acc(ctx, view.attn_k_w, x_norm); + if (view.attn_k_b != nullptr) { + K = ggml_add(ctx, K, view.attn_k_b); + } ggml_tensor * V = mul_mat_f32acc(ctx, view.attn_v_w, x_norm); + if (view.attn_v_b != nullptr) { + V = ggml_add(ctx, V, view.attn_v_b); + } Q = ggml_reshape_4d(ctx, Q, head_dim, n_heads, T_seq, 1); K = ggml_reshape_4d(ctx, K, head_dim, n_kv_heads, T_seq, 1); @@ -640,8 +667,17 @@ ggml_tensor * block_step_batched(ggml_context * ctx, ggml_tensor * x_norm = rms_norm(ctx, x, view.norm_attn_w, rms_eps); ggml_tensor * Q = mul_mat_f32acc(ctx, view.attn_q_w, x_norm); // [q_dim, B] + if (view.attn_q_b != nullptr) { + Q = ggml_add(ctx, Q, view.attn_q_b); + } ggml_tensor * K = mul_mat_f32acc(ctx, view.attn_k_w, x_norm); // [kv_dim, B] + if (view.attn_k_b != nullptr) { + K = ggml_add(ctx, K, view.attn_k_b); + } ggml_tensor * V = mul_mat_f32acc(ctx, view.attn_v_w, x_norm); // [kv_dim, B] + if (view.attn_v_b != nullptr) { + V = ggml_add(ctx, V, view.attn_v_b); + } // Batch on ne[2] (the per-token/position axis) so RoPE applies each // utterance's own position. T == 1 per utterance, so ne[2] == B. @@ -753,8 +789,17 @@ ggml_tensor * block_prefill_batched(ggml_context * ctx, ggml_tensor * x_norm = rms_norm(ctx, x, view.norm_attn_w, rms_eps); ggml_tensor * Q = mul_mat_f32acc(ctx, view.attn_q_w, x_norm); // [q_dim, T, B] + if (view.attn_q_b != nullptr) { + Q = ggml_add(ctx, Q, view.attn_q_b); + } ggml_tensor * K = mul_mat_f32acc(ctx, view.attn_k_w, x_norm); // [kv_dim, T, B] + if (view.attn_k_b != nullptr) { + K = ggml_add(ctx, K, view.attn_k_b); + } ggml_tensor * V = mul_mat_f32acc(ctx, view.attn_v_w, x_norm); // [kv_dim, T, B] + if (view.attn_v_b != nullptr) { + V = ggml_add(ctx, V, view.attn_v_b); + } Q = ggml_reshape_4d(ctx, Q, head_dim, n_heads, T, B); K = ggml_reshape_4d(ctx, K, head_dim, n_kv_heads, T, B); diff --git a/src/causal_lm/causal_lm.h b/src/causal_lm/causal_lm.h index 1992f86bf..68ad5bbc8 100644 --- a/src/causal_lm/causal_lm.h +++ b/src/causal_lm/causal_lm.h @@ -36,6 +36,10 @@ struct BlockView { ggml_tensor * attn_k_w = nullptr; // [hidden, n_kv_heads * head_dim] ggml_tensor * attn_v_w = nullptr; // [hidden, n_kv_heads * head_dim] ggml_tensor * attn_o_w = nullptr; // [n_heads * head_dim, hidden] + // Optional Qwen2 projection biases; null for existing callers. + ggml_tensor * attn_q_b = nullptr; + ggml_tensor * attn_k_b = nullptr; + ggml_tensor * attn_v_b = nullptr; // Per-head Q/K RMSNorm (Qwen3). Null on Llama-style decoders (e.g. // Voxtral's Ministral backbone); helpers skip the norm when null. ggml_tensor * attn_q_norm = nullptr; // [head_dim] per-head Q-norm, or null diff --git a/src/transcribe-arch.cpp b/src/transcribe-arch.cpp index fb768cbac..6cfb04ee8 100644 --- a/src/transcribe-arch.cpp +++ b/src/transcribe-arch.cpp @@ -27,6 +27,10 @@ namespace qwen3_asr { extern const Arch arch; } +namespace arkasr { +extern const Arch arch; +} + namespace moss { extern const Arch arch; } @@ -93,10 +97,26 @@ const Arch * find_arch(const char * name) { } static const Arch * const k_archs[] = { - ¶keet::arch, &cohere::arch, &canary::arch, &qwen3_asr::arch, &voxtral::arch, - &voxtral_realtime::arch, &canary_qwen::arch, &whisper::arch, &moonshine::arch, &moonshine_streaming::arch, - &sensevoice::arch, &funasr_nano::arch, &gigaam::arch, &granite::arch, &granite_nar::arch, - &medasr::arch, &moss::arch, &sortformer::arch, &granite5_ctc::arch, + ¶keet::arch, + &cohere::arch, + &canary::arch, + &qwen3_asr::arch, + &arkasr::arch, + &voxtral::arch, + &voxtral_realtime::arch, + &canary_qwen::arch, + &whisper::arch, + &moonshine::arch, + &moonshine_streaming::arch, + &sensevoice::arch, + &funasr_nano::arch, + &gigaam::arch, + &granite::arch, + &granite_nar::arch, + &medasr::arch, + &moss::arch, + &sortformer::arch, + &granite5_ctc::arch, }; constexpr size_t k_n = sizeof(k_archs) / sizeof(k_archs[0]);