AUSTIN2526/FireRedTrainAlign

5

stars

5

commits

Python

primary language

Aug 9, 2026

updated

README

πŸ”₯ FireRedASR-HF

A complete rewrite of FireRedTeam/FireRedASR-AED-L following the standard Hugging Face architecture:

PretrainedConfig / PreTrainedModel / FeatureExtractor / Tokenizer / Processor / Seq2SeqTrainer

The model itself β€” Conformer encoder + Transformer decoder + beam search β€” remains fully consistent with the original implementation. Only the external API has been redesigned to support Hugging Face interfaces such as from_pretrained, save_pretrained, generate, and Trainer.


Directory Structure

FireRedASR-HF/
β”œβ”€β”€ fireredasr/
β”‚   β”œβ”€β”€ configuration_fireredasr.py        # FireRedAsrConfig(PretrainedConfig)
β”‚   β”œβ”€β”€ modeling_fireredasr.py             # FireRedAsrForConditionalGeneration(PreTrainedModel, GenerationMixin)
β”‚   β”œβ”€β”€ modeling_conformer_encoder.py      # Conformer encoder (official checkpoint-compatible weight names)
β”‚   β”œβ”€β”€ modeling_transformer_decoder.py    # Transformer decoder + batch beam search
β”‚   β”œβ”€β”€ feature_extraction_fireredasr.py   # FireRedAsrFeatureExtractor(SequenceFeatureExtractor)
β”‚   β”œβ”€β”€ tokenization_fireredasr.py         # FireRedAsrTokenizer(PreTrainedTokenizer)
β”‚   β”œβ”€β”€ processing_fireredasr.py           # FireRedAsrProcessor(ProcessorMixin)
β”‚   └── audio_utils.py                     # Audio loading / mono conversion / resampling
β”œβ”€β”€ scripts/
β”‚   β”œβ”€β”€ convert_original_checkpoint.py     # Official model.pth.tar β†’ HF repository
β”‚   β”œβ”€β”€ run_asr.py                         # Batch inference (optional WhisperX alignment)
β”‚   └── run_finetune.py                    # Seq2SeqTrainer fine-tuning
β”œβ”€β”€ audio/                                 # Example audio files
└── requirements.txt

Installation

conda create -n firered python=3.10
conda activate firered
pip install -r requirements.txt

If kaldi_native_fbank cannot be installed, torchaudio is automatically used as the fallback implementation for fbank extraction.

whisperx is only required when word-level timestamps are needed.


1. Convert the Official Checkpoint

First, download the official model directory from Hugging Face. It should contain:

  • model.pth.tar
  • cmvn.ark
  • dict.txt
  • train_bpe1000.model

Then convert it to the standard Hugging Face format:

python scripts/convert_original_checkpoint.py \
    --original_dir FireRedASR-AED-L \
    --output_dir   FireRedASR-AED-L-hf

The resulting directory will look like:

FireRedASR-AED-L-hf/
β”œβ”€β”€ config.json
β”œβ”€β”€ generation_config.json
β”œβ”€β”€ model.safetensors
β”œβ”€β”€ preprocessor_config.json        # Contains CMVN statistics
β”œβ”€β”€ vocab.txt
β”œβ”€β”€ bpe.model
β”œβ”€β”€ tokenizer_config.json
└── *.py                            # Model source files for trust_remote_code

The converted model can then be pushed to the Hugging Face Hub using push_to_hub() or loaded with trust_remote_code=True.

If you do not want to include the Python source code in the output directory, add:

--no_code

2. Inference

Python API

import torch

from fireredasr import (
    FireRedAsrForConditionalGeneration,
    FireRedAsrProcessor,
    load_audio,
)

processor = FireRedAsrProcessor.from_pretrained(
    "./FireRedASR-AED-L-hf"
)

model = (
    FireRedAsrForConditionalGeneration
    .from_pretrained("./FireRedASR-AED-L-hf")
    .eval()
    .cuda()
)

wav, sr = load_audio("audio/90061.wav", 16000)

inputs = processor(
    audio=wav,
    sampling_rate=sr,
    return_tensors="pt",
).to("cuda")

ids = model.generate(
    **inputs,
    num_beams=20,
    length_penalty=0.1,
)

print(
    processor.batch_decode(
        ids,
        skip_special_tokens=True,
    )[0]
)

The model can also be loaded through the Hugging Face Auto class. Importing fireredasr automatically registers the model:

from transformers import AutoModelForSpeechSeq2Seq

model = AutoModelForSpeechSeq2Seq.from_pretrained(
    "./FireRedASR-AED-L-hf"
)

Command Line

Transcription only

python scripts/run_asr.py \
    --model ./FireRedASR-AED-L-hf \
    --audio_dir audio \
    --output_dir asr_result

Transcription with WhisperX word-level alignment

python scripts/run_asr.py \
    --model ./FireRedASR-AED-L-hf \
    --audio_dir audio \
    --align

The output format is consistent with the original project.


3. Fine-Tuning

The dataset format is the same as the original project:

audio/<utt_id>.wav
labels.txt

where each line in labels.txt follows:

<utt_id>\t<transcription>

Example:

python scripts/run_finetune.py \
    --model ./FireRedASR-AED-L-hf \
    --audio_dir audio \
    --label_file labels.txt \
    --output_dir ./firered-finetuned \
    --num_train_epochs 10 \
    --learning_rate 1e-5 \
    --per_device_train_batch_size 2 \
    --gradient_accumulation_steps 8 \
    --do_eval \
    --predict_with_generate \
    --bf16

Common Options

ArgumentDescription
--freeze_encoderFreeze the Conformer encoder and train only the decoder
--difficulty_strategylength_based / unk_based / combined, following the difficulty-weighted sampling strategy from the original project
--predict_with_generateRun actual beam search during evaluation and calculate CER
--ditherAdd noise to fbank features during training (default: 1.0; inference always uses 0.0)

Other Seq2SeqTrainingArguments features can be used directly, including:

  • DeepSpeed
  • resume_from_checkpoint
  • report_to="wandb"
  • Gradient accumulation
  • Mixed precision
  • Distributed training

Mapping to the Original Project

OriginalFireRedASR-HF
FireRedAsr(model_name, weight)FireRedAsrForConditionalGeneration.from_pretrained(path)
HFASRFeatExtractor(dir)FireRedAsrFeatureExtractor.from_pretrained(path)
ChineseCharEnglishSpmTokenizer(dir)FireRedAsrTokenizer.from_pretrained(path)
Manually construct [sos] + ids + [eos]Pass labels to the model; tokens are shifted internally using shift_tokens_right
wrapper.forward(feat, lengths, tgt)model(input_features=..., attention_mask=..., labels=...) β†’ Seq2SeqLMOutput
model.transcribe(feat, lengths, durs, beam_size=...)model.generate(input_features=..., num_beams=...)
Custom training loop + torch.save(state_dict)Seq2SeqTrainer + save_pretrained
cmvn.ark + kaldiioCMVN statistics embedded directly in preprocessor_config.json

Implementation Notes

Custom generate()

generate() is implemented specifically for FireRedASR.

The incremental cache used by the FireRedASR decoder is different from the standard Hugging Face KV-cache mechanism. Therefore, the original, validated batch beam search implementation is preserved.

Only the external parameter names are adapted to Hugging Face conventions:

num_beams
max_new_tokens
length_penalty
num_return_sequences

FireRedASR-specific parameters are also preserved:

softmax_smoothing
eos_penalty

The returned sequence format is compatible with Hugging Face:

<sos> ... <eos> <pad> <pad> ...

This allows the output to be passed directly to:

processor.batch_decode(...)

and also enables:

Seq2SeqTrainer(
    predict_with_generate=True
)

Weight Names Remain Unchanged

The original checkpoint weight names are preserved:

encoder.*
decoder.*

The conversion process does not rename the model parameters.

This maintains a 1:1 correspondence with the official FireRedASR checkpoint.


Shared Embeddings and Output Projection

The decoder embedding and output projection are tied.

The implementation uses:

_tied_weights_keys

to correctly handle the shared parameters.

This ensures that save_pretrained() can save the model in safetensors format without triggering shared-tensor errors.


Compatibility

The implementation has been tested with Transformers 5.x.

Fallbacks are also provided for important API differences between Transformers 4.x and 5.x, including:

  • processing_class
  • _tied_weights_keys type differences
  • Processor serialization behavior

Important Notes

For stable decoding, it is recommended to keep individual audio segments below:

60 seconds

Longer segments may increase the likelihood of repeated phrases during decoding.

Maximum Audio Length

Audio longer than approximately:

200 seconds

may exceed the positional encoding range.

The effective limit is determined by:

max_source_positions

Audio Preprocessing

Input audio is automatically:

  1. Converted to mono
  2. Resampled to 16 kHz
  3. Converted into fbank features
  4. Normalized using the embedded CMVN statistics

No external cmvn.ark or kaldiio file is required after checkpoint conversion.


Summary

FireRedASR-HF keeps the original FireRedASR-AED-L architecture and decoding behavior while exposing it through the standard Hugging Face ecosystem.

The main goal is to make the original implementation compatible with familiar Hugging Face workflows:

model = FireRedAsrForConditionalGeneration.from_pretrained(...)
processor = FireRedAsrProcessor.from_pretrained(...)
output = model.generate(...)
model.save_pretrained(...)

and:

trainer = Seq2SeqTrainer(...)
trainer.train()

This makes FireRedASR easier to integrate into existing Hugging Face projects, fine-tuning pipelines, distributed training workflows, and Hugging Face Hub repositories without changing the underlying ASR model architecture.

Contributors

AUSTIN2526

5 commits

AUSTIN2526/FireRedTrainAlign

5

stars

5

commits

Python

primary language

Aug 9, 2026

updated

README

πŸ”₯ FireRedASR-HF

A complete rewrite of FireRedTeam/FireRedASR-AED-L following the standard Hugging Face architecture:

PretrainedConfig / PreTrainedModel / FeatureExtractor / Tokenizer / Processor / Seq2SeqTrainer

The model itself β€” Conformer encoder + Transformer decoder + beam search β€” remains fully consistent with the original implementation. Only the external API has been redesigned to support Hugging Face interfaces such as from_pretrained, save_pretrained, generate, and Trainer.


Directory Structure

FireRedASR-HF/
β”œβ”€β”€ fireredasr/
β”‚   β”œβ”€β”€ configuration_fireredasr.py        # FireRedAsrConfig(PretrainedConfig)
β”‚   β”œβ”€β”€ modeling_fireredasr.py             # FireRedAsrForConditionalGeneration(PreTrainedModel, GenerationMixin)
β”‚   β”œβ”€β”€ modeling_conformer_encoder.py      # Conformer encoder (official checkpoint-compatible weight names)
β”‚   β”œβ”€β”€ modeling_transformer_decoder.py    # Transformer decoder + batch beam search
β”‚   β”œβ”€β”€ feature_extraction_fireredasr.py   # FireRedAsrFeatureExtractor(SequenceFeatureExtractor)
β”‚   β”œβ”€β”€ tokenization_fireredasr.py         # FireRedAsrTokenizer(PreTrainedTokenizer)
β”‚   β”œβ”€β”€ processing_fireredasr.py           # FireRedAsrProcessor(ProcessorMixin)
β”‚   └── audio_utils.py                     # Audio loading / mono conversion / resampling
β”œβ”€β”€ scripts/
β”‚   β”œβ”€β”€ convert_original_checkpoint.py     # Official model.pth.tar β†’ HF repository
β”‚   β”œβ”€β”€ run_asr.py                         # Batch inference (optional WhisperX alignment)
β”‚   └── run_finetune.py                    # Seq2SeqTrainer fine-tuning
β”œβ”€β”€ audio/                                 # Example audio files
└── requirements.txt

Installation

conda create -n firered python=3.10
conda activate firered
pip install -r requirements.txt

If kaldi_native_fbank cannot be installed, torchaudio is automatically used as the fallback implementation for fbank extraction.

whisperx is only required when word-level timestamps are needed.


1. Convert the Official Checkpoint

First, download the official model directory from Hugging Face. It should contain:

  • model.pth.tar
  • cmvn.ark
  • dict.txt
  • train_bpe1000.model

Then convert it to the standard Hugging Face format:

python scripts/convert_original_checkpoint.py \
    --original_dir FireRedASR-AED-L \
    --output_dir   FireRedASR-AED-L-hf

The resulting directory will look like:

FireRedASR-AED-L-hf/
β”œβ”€β”€ config.json
β”œβ”€β”€ generation_config.json
β”œβ”€β”€ model.safetensors
β”œβ”€β”€ preprocessor_config.json        # Contains CMVN statistics
β”œβ”€β”€ vocab.txt
β”œβ”€β”€ bpe.model
β”œβ”€β”€ tokenizer_config.json
└── *.py                            # Model source files for trust_remote_code

The converted model can then be pushed to the Hugging Face Hub using push_to_hub() or loaded with trust_remote_code=True.

If you do not want to include the Python source code in the output directory, add:

--no_code

2. Inference

Python API

import torch

from fireredasr import (
    FireRedAsrForConditionalGeneration,
    FireRedAsrProcessor,
    load_audio,
)

processor = FireRedAsrProcessor.from_pretrained(
    "./FireRedASR-AED-L-hf"
)

model = (
    FireRedAsrForConditionalGeneration
    .from_pretrained("./FireRedASR-AED-L-hf")
    .eval()
    .cuda()
)

wav, sr = load_audio("audio/90061.wav", 16000)

inputs = processor(
    audio=wav,
    sampling_rate=sr,
    return_tensors="pt",
).to("cuda")

ids = model.generate(
    **inputs,
    num_beams=20,
    length_penalty=0.1,
)

print(
    processor.batch_decode(
        ids,
        skip_special_tokens=True,
    )[0]
)

The model can also be loaded through the Hugging Face Auto class. Importing fireredasr automatically registers the model:

from transformers import AutoModelForSpeechSeq2Seq

model = AutoModelForSpeechSeq2Seq.from_pretrained(
    "./FireRedASR-AED-L-hf"
)

Command Line

Transcription only

python scripts/run_asr.py \
    --model ./FireRedASR-AED-L-hf \
    --audio_dir audio \
    --output_dir asr_result

Transcription with WhisperX word-level alignment

python scripts/run_asr.py \
    --model ./FireRedASR-AED-L-hf \
    --audio_dir audio \
    --align

The output format is consistent with the original project.


3. Fine-Tuning

The dataset format is the same as the original project:

audio/<utt_id>.wav
labels.txt

where each line in labels.txt follows:

<utt_id>\t<transcription>

Example:

python scripts/run_finetune.py \
    --model ./FireRedASR-AED-L-hf \
    --audio_dir audio \
    --label_file labels.txt \
    --output_dir ./firered-finetuned \
    --num_train_epochs 10 \
    --learning_rate 1e-5 \
    --per_device_train_batch_size 2 \
    --gradient_accumulation_steps 8 \
    --do_eval \
    --predict_with_generate \
    --bf16

Common Options

ArgumentDescription
--freeze_encoderFreeze the Conformer encoder and train only the decoder
--difficulty_strategylength_based / unk_based / combined, following the difficulty-weighted sampling strategy from the original project
--predict_with_generateRun actual beam search during evaluation and calculate CER
--ditherAdd noise to fbank features during training (default: 1.0; inference always uses 0.0)

Other Seq2SeqTrainingArguments features can be used directly, including:

  • DeepSpeed
  • resume_from_checkpoint
  • report_to="wandb"
  • Gradient accumulation
  • Mixed precision
  • Distributed training

Mapping to the Original Project

OriginalFireRedASR-HF
FireRedAsr(model_name, weight)FireRedAsrForConditionalGeneration.from_pretrained(path)
HFASRFeatExtractor(dir)FireRedAsrFeatureExtractor.from_pretrained(path)
ChineseCharEnglishSpmTokenizer(dir)FireRedAsrTokenizer.from_pretrained(path)
Manually construct [sos] + ids + [eos]Pass labels to the model; tokens are shifted internally using shift_tokens_right
wrapper.forward(feat, lengths, tgt)model(input_features=..., attention_mask=..., labels=...) β†’ Seq2SeqLMOutput
model.transcribe(feat, lengths, durs, beam_size=...)model.generate(input_features=..., num_beams=...)
Custom training loop + torch.save(state_dict)Seq2SeqTrainer + save_pretrained
cmvn.ark + kaldiioCMVN statistics embedded directly in preprocessor_config.json

Implementation Notes

Custom generate()

generate() is implemented specifically for FireRedASR.

The incremental cache used by the FireRedASR decoder is different from the standard Hugging Face KV-cache mechanism. Therefore, the original, validated batch beam search implementation is preserved.

Only the external parameter names are adapted to Hugging Face conventions:

num_beams
max_new_tokens
length_penalty
num_return_sequences

FireRedASR-specific parameters are also preserved:

softmax_smoothing
eos_penalty

The returned sequence format is compatible with Hugging Face:

<sos> ... <eos> <pad> <pad> ...

This allows the output to be passed directly to:

processor.batch_decode(...)

and also enables:

Seq2SeqTrainer(
    predict_with_generate=True
)

Weight Names Remain Unchanged

The original checkpoint weight names are preserved:

encoder.*
decoder.*

The conversion process does not rename the model parameters.

This maintains a 1:1 correspondence with the official FireRedASR checkpoint.


Shared Embeddings and Output Projection

The decoder embedding and output projection are tied.

The implementation uses:

_tied_weights_keys

to correctly handle the shared parameters.

This ensures that save_pretrained() can save the model in safetensors format without triggering shared-tensor errors.


Compatibility

The implementation has been tested with Transformers 5.x.

Fallbacks are also provided for important API differences between Transformers 4.x and 5.x, including:

  • processing_class
  • _tied_weights_keys type differences
  • Processor serialization behavior

Important Notes

For stable decoding, it is recommended to keep individual audio segments below:

60 seconds

Longer segments may increase the likelihood of repeated phrases during decoding.

Maximum Audio Length

Audio longer than approximately:

200 seconds

may exceed the positional encoding range.

The effective limit is determined by:

max_source_positions

Audio Preprocessing

Input audio is automatically:

  1. Converted to mono
  2. Resampled to 16 kHz
  3. Converted into fbank features
  4. Normalized using the embedded CMVN statistics

No external cmvn.ark or kaldiio file is required after checkpoint conversion.


Summary

FireRedASR-HF keeps the original FireRedASR-AED-L architecture and decoding behavior while exposing it through the standard Hugging Face ecosystem.

The main goal is to make the original implementation compatible with familiar Hugging Face workflows:

model = FireRedAsrForConditionalGeneration.from_pretrained(...)
processor = FireRedAsrProcessor.from_pretrained(...)
output = model.generate(...)
model.save_pretrained(...)

and:

trainer = Seq2SeqTrainer(...)
trainer.train()

This makes FireRedASR easier to integrate into existing Hugging Face projects, fine-tuning pipelines, distributed training workflows, and Hugging Face Hub repositories without changing the underlying ASR model architecture.

Contributors

AUSTIN2526

5 commits

Languages

Python

100.0%