ZihanLiummyycc/drax-dbank-full-decoder

Drax full-decoder adaptation for DementiaBank

Adapted parameters accompanying Diffusion and Flow Matching ASR for Elderly Speech. Code and experiment recipes.

Contents and scope

This package contains 579,883,674 updated parameter elements in 303 tensors, approximately 2.16 GiB. It is the complete adapted Drax decoder, including its acoustic projection and conditioning modules, not a small LoRA adapter. These are final parameter values to replace the corresponding native parameters, not additive deltas.

The frozen audio encoder, unadapted base checkpoint, optimizer, raw predictions, training logs and DementiaBank data are not included. The deterministic rotary buffer is supplied by the native base and was checked against the original checkpoint. This is not a standalone Transformers model or a standard PEFT adapter.

Native loading

Use the Drax environment described in the accompanying code repository. Resolve bases at the immutable revisions in base_assets.json; that file also records their hashes. Review the loading scripts before executing them.

git clone https://github.com/ZihanLiummyycc/diffusion-asr-elderly-speech.git
cd diffusion-asr-elderly-speech
git checkout 29709f075399be63ed6d1e536a0a0ca3bba0a072

from pathlib import Path
import sys
import torch
from huggingface_hub import snapshot_download
from transformers import WhisperProcessor

code_root = Path.cwd()  # cloned repository above
component_dir = Path(snapshot_download("ZihanLiummyycc/drax-dbank-full-decoder"))
base_dir = snapshot_download("aiola/drax-v1", revision="96dda0e1d86c9b4c7f06a377410283ede1b671ab")
processor_dir = snapshot_download(
    "openai/whisper-large-v3", revision="06f233fe06e710322aca913c1bc4249a0d71fce1",
    allow_patterns=["*.json", "*.txt"],
)
sys.path.insert(0, str(code_root / "benchmarks/repos/drax"))
sys.path.insert(0, str(component_dir))
from drax.model.drax_model import Drax
from load_drax_component import load_drax_component

model = Drax.from_pretrained(base_dir, device="cuda", torch_dtype=torch.float32).eval()
load_drax_component(model, component_dir)
processor = WhisperProcessor.from_pretrained(processor_dir, local_files_only=True)

The base download includes encoder weights needed locally. They are not redistributed here. Use the native Drax solver and the accompanying benchmarks/tools/dbank_drax/evaluate_dbank.py recipe for inference. Decoder parameters are replaced by the component; no training is performed by the loader.

Evaluation and verification

The original evaluation reported 34.85% WER on 928 DementiaBank utterances. This is a historical aggregate associated with the original experiment, not a new benchmark run during publication. Consult the paper and experiment recipes for data partitions, decoding configuration and the limitations of exploratory scope comparisons.

Verification passed for local file hashes, tensor names, shapes, finite values, strict native parameter loading, and original-versus-exported generation. The native check used an NVIDIA A40, FP32, two synthetic one-second inputs (silence and chirp), four discrete-flow steps, and the same random seeds for both loading paths. Before loading the exported state, every selected decoder parameter was cleared. Token outputs matched exactly. Base assets and the used native source files were checked against immutable public versions.

These checks do not measure WER, RTF, robustness or privacy leakage. The convenience wrapper has local unit tests; a fresh installation from this public repository has not been evaluated on the full clinical corpus. See verification.json and checksums.json.

Intended use and limitations

Research on adapting ASR to elderly speech. Not a diagnostic model and not validated for clinical decisions. Transcription errors and dataset bias remain possible. No raw clinical records are released, but this does not establish that the weights are free of memorization. Obtain any required data access separately and process only audio you are authorized to use.

Attribution and terms

Based on aiOla's Drax: Speech Recognition with Discrete Flow Matching, by Aviv Navon and colleagues, paper. Its pinned model card declares CC BY 4.0. The native Drax code separately carries CC BY-NC 4.0. These are different license scopes, not interchangeable labels. Retained notices and the authors' contribution policy are in LICENSE.md and licenses/. This is a modified research release, not an official aiOla model or an endorsement by upstream authors.