RASMUS/Finnish-ASR-Canary-v2

1# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7#     http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14#15 16"""17# This script evaluates CTC and Transducer (RNNT) models (only Hybrid Transducer-CTC in case of Transducer) in context biasing mode18# by applying CTC-based Word Spotter (paper link) 19 20# Config Help21 22To discover all arguments of the script, please run :23python eval_greedy_decoding_with_context_biasing.py --help24python eval_greedy_decoding_with_context_biasing.py --cfg job25 26# USAGE27 28python eval_greedy_decoding_with_context_biasing.py \29            nemo_model_file=<path to the .nemo file of the model> \30            input_manifest=<path to the evaluation JSON manifest file \31            preds_output_folder=<folder to store the predictions> \32            decoder_type=<type of model decoder [ctc or rnnt]> \33            acoustic_batch_size=<batch size to calculate log probabilities> \34            apply_context_biasing=<True or False to apply context biasing> \35            context_file=<path to the context biasing file with key words/phrases> \36            beam_threshold=[<list of the beam thresholds, separated with commas>] \37            context_score=[<list of the context scores, separated with commas>] \38            ctc_ali_token_weight=[<list of the ctc alignment token weights, separated with commas>] \39            ...40 41# Description of context biasing graph:42Context biasing file contains words/phrases with their spellings43(one word/phrase per line, spellings are separated from word/phrase by underscore symbol):44WORD1_SPELLING145WORD2_SPELLING1_SPELLING246...47nvidia_nvidia48gpu_gpu_g p u49nvlink_nvlink_nv link50...51alternative spellings help to improve the recognition accuracy of abbreviations and complicated words,52which are often recognized as separate words (gpu -> g p u, nvlink -> nv link, tensorrt -> tensor rt, and so on).53 54 55# Grid Search for Hyper parameters56 57For grid search, you can provide a list of arguments as follows -58 59            beam_threshold=[4.0,5.0,6.0,....] \60            context_score=[1.0,1.5,...,4.0,4.5] \61            ctc_ali_token_weight=[0.1,0.2,...,0.7,0.8] \62 63"""64 65 66import contextlib67import json68import os69import tempfile70from dataclasses import dataclass, field, is_dataclass71from pathlib import Path72from typing import Dict, Optional73 74import editdistance75import numpy as np76import torch77from omegaconf import MISSING, OmegaConf78from sklearn.model_selection import ParameterGrid79from tqdm.auto import tqdm80 81import nemo.collections.asr as nemo_asr82from nemo.collections.asr.models import EncDecCTCModelBPE, EncDecHybridRNNTCTCModel83from nemo.collections.asr.parts import context_biasing84from nemo.core.config import hydra_runner85from nemo.utils import logging86 87 88@dataclass89class EvalContextBiasingConfig:90    """91    Evaluate CTC and Transducer (RNNT) ASR models in greedy decoding with context biasing.92    """93 94    # # The path of the '.nemo' file of the ASR model or the name of a pretrained model (ngc / huggingface)95    nemo_model_file: str = MISSING96 97    # File paths98    input_manifest: str = MISSING  # The manifest file of the evaluation set99    preds_output_folder: str = MISSING  # The folder where the predictions are stored100 101    # Parameters for inference102    acoustic_batch_size: int = 128  # The batch size to calculate log probabilities103    beam_batch_size: int = 128  # The batch size to be used for beam search decoding104    device: str = "cuda"  # The device to load the model onto to calculate log probabilities105    use_amp: bool = False  # Whether to use AMP if available to calculate log probabilities106    num_workers: int = 1  # Number of workers for DataLoader107    decoder_type: Optional[str] = None  # [ctc, rnnt] decoder type for asr model108 109    # Context-Biasing params110    apply_context_biasing: bool = False  # True in case of context biasing111    context_file: str = MISSING  # text file with context biasing words and their spellings112    spelling_separator: str = "_"  # separator between word and its spellings in context biasing file113    beam_threshold: list[float] = field(default_factory=lambda: [5.0])  # beam pruning threshold for ctc-ws decoding114    context_score: list[float] = field(default_factory=lambda: [3.0])  # per token weight for context biasing words115    ctc_ali_token_weight: list[float] = field(116        default_factory=lambda: [0.6]117    )  # weight of CTC tokens to prevent false accept errors118    print_cb_stats: bool = False  # print context biasing stats (mostly for debugging)119 120    # Auxiliary parameters121    sort_logits: bool = True  # do logits sorting before decoding - it reduces computation on puddings122    softmax_temperature: float = 1.00123    preserve_alignments: bool = False124 125 126def decoding_step(127    asr_model: nemo_asr.models.ASRModel,128    cfg: EvalContextBiasingConfig,129    encoder_outputs: list[torch.Tensor],130    ctc_logprobs: list[np.ndarray],131    target_transcripts: list[str],132    audio_file_paths: list[str],133    durations: list[str],134    preds_output_manifest: str,135    beam_batch_size: int = 128,136    progress_bar: bool = True,137    context_graph: context_biasing.ContextGraphCTC = None,138    blank_idx: int = 0,139    hp: Optional[Dict] = None,140) -> tuple[float, float]:141 142    # run CTC-based Word Spotter:143    if cfg.apply_context_biasing:144        ws_results = {}145        for idx, logits in tqdm(146            enumerate(ctc_logprobs), desc=f"Eval CTC-based Word Spotter...", ncols=120, total=len(ctc_logprobs)147        ):148            ws_results[audio_file_paths[idx]] = context_biasing.run_word_spotter(149                logits,150                context_graph,151                asr_model,152                blank_idx=blank_idx,153                beam_threshold=hp['beam_threshold'],154                cb_weight=hp['context_score'],155                ctc_ali_token_weight=hp['ctc_ali_token_weight'],156            )157 158    level = logging.getEffectiveLevel()159    logging.setLevel(logging.CRITICAL)160    # reset config161    asr_model.change_decoding_strategy(None)162 163    # preserve alignment:164    asr_model.cfg.decoding.preserve_alignments = cfg.preserve_alignments165 166    # update model's decoding strategy config167    if isinstance(asr_model, EncDecCTCModelBPE):168        # in case of ctc169        asr_model.cfg.decoding.strategy = "greedy"170    else:171        # in case of rnnt172        asr_model.cfg.decoding.strategy = "greedy_batch"173        # fast greedy batch decoding:174        asr_model.cfg.decoding.greedy.loop_labels = True175 176    # update model's decoding strategy177    asr_model.change_decoding_strategy(asr_model.cfg.decoding)178    logging.setLevel(level)179 180    wer_dist_first = cer_dist_first = 0181    words_count = chars_count = sample_idx = 0182 183    out_manifest = open(preds_output_manifest, 'w', encoding='utf_8', newline='\n')184 185    # ctc part for both EncDecCTCModelBPE and EncDecHybridRNNTCTCModel186    if cfg.decoder_type == "ctc":187        for batch_idx, probs in enumerate(ctc_logprobs):188            preds = np.argmax(probs, axis=1)189            if cfg.apply_context_biasing and ws_results[audio_file_paths[batch_idx]]:190                # make new text by mearging alignment with ctc-ws predictions:191                if cfg.print_cb_stats:192                    logging.info("\n" + "********" * 10)193                    logging.info(f"File name: {audio_file_paths[batch_idx]}")194                pred_text, raw_text = context_biasing.merge_alignment_with_ws_hyps(195                    preds,196                    asr_model,197                    ws_results[audio_file_paths[batch_idx]],198                    decoder_type="ctc",199                    blank_idx=blank_idx,200                    print_stats=cfg.print_cb_stats,201                )202                if cfg.print_cb_stats:203                    logging.info(f"raw text: {raw_text}")204                    logging.info(f"hyp text: {pred_text}")205                    logging.info(f"ref text: {target_transcripts[batch_idx]}")206            else:207                preds_tensor = torch.tensor(preds, device='cpu').unsqueeze(0)208                if isinstance(asr_model, EncDecHybridRNNTCTCModel):209                    hyp = asr_model.ctc_decoding.ctc_decoder_predictions_tensor(preds_tensor)[0]210                else:211                    hyp = asr_model.wer.decoding.ctc_decoder_predictions_tensor(preds_tensor)[0]212                pred_text = hyp.text213            pred_split_w = pred_text.split()214            target_split_w = target_transcripts[batch_idx].split()215            pred_split_c = list(pred_text)216            target_split_c = list(target_transcripts[batch_idx])217 218            wer_dist = editdistance.eval(target_split_w, pred_split_w)219            cer_dist = editdistance.eval(target_split_c, pred_split_c)220 221            wer_dist_first += wer_dist222            cer_dist_first += cer_dist223            words_count += len(target_split_w)224            chars_count += len(target_split_c)225 226            if preds_output_manifest:227                item = {228                    'audio_filepath': audio_file_paths[batch_idx],229                    'duration': durations[batch_idx],230                    'text': target_transcripts[batch_idx],231                    'pred_text': pred_text,232                    'wer': f"{wer_dist/len(target_split_w):.4f}",233                }234                print(json.dumps(item), file=out_manifest)235        out_manifest.close()236 237        return wer_dist_first / words_count, cer_dist_first / chars_count238 239    # rnnt part for EncDecHybridRNNTCTCModel240    else:241        if progress_bar:242            description = "Greedy_batch decoding.."243            it = tqdm(range(int(np.ceil(len(encoder_outputs) / beam_batch_size))), desc=description, ncols=120)244        else:245            it = range(int(np.ceil(len(encoder_outputs) / beam_batch_size)))246        for batch_idx in it:247            probs_batch = encoder_outputs[batch_idx * beam_batch_size : (batch_idx + 1) * beam_batch_size]248            probs_lens = torch.tensor([prob.shape[-1] for prob in probs_batch])249            with torch.no_grad():250                packed_batch = torch.zeros(len(probs_batch), probs_batch[0].shape[0], max(probs_lens), device='cpu')251 252                for prob_index in range(len(probs_batch)):253                    packed_batch[prob_index, :, : probs_lens[prob_index]] = torch.tensor(254                        probs_batch[prob_index].unsqueeze(0), device=packed_batch.device, dtype=packed_batch.dtype255                    )256                best_hyp_batch = asr_model.decoding.rnnt_decoder_predictions_tensor(257                    packed_batch,258                    probs_lens,259                    return_hypotheses=True,260                )261            beams_batch = [[x] for x in best_hyp_batch]262 263            for beams_idx, beams in enumerate(beams_batch):264                target = target_transcripts[sample_idx + beams_idx]265                target_split_w = target.split()266                target_split_c = list(target)267                words_count += len(target_split_w)268                chars_count += len(target_split_c)269                for candidate_idx, candidate in enumerate(beams):270                    if cfg.apply_context_biasing and ws_results[audio_file_paths[sample_idx + beams_idx]]:271                        # make new text by mearging alignment with ctc-ws predictions:272                        if cfg.print_cb_stats:273                            logging.info("\n" + "********" * 10)274                            logging.info(f"File name: {audio_file_paths[batch_idx]}")275                        pred_text, raw_text = context_biasing.merge_alignment_with_ws_hyps(276                            candidate,277                            asr_model,278                            ws_results[audio_file_paths[sample_idx + beams_idx]],279                            decoder_type="rnnt",280                            blank_idx=blank_idx,281                            print_stats=cfg.print_cb_stats,282                        )283                        if cfg.print_cb_stats:284                            logging.info(f"raw text: {raw_text}")285                            logging.info(f"hyp text: {pred_text}")286                            logging.info(f"ref text: {target_transcripts[sample_idx + beams_idx]}")287                    else:288                        pred_text = candidate.text289 290                    pred_split_w = pred_text.split()291                    wer_dist = editdistance.eval(target_split_w, pred_split_w)292                    pred_split_c = list(pred_text)293                    cer_dist = editdistance.eval(target_split_c, pred_split_c)294 295                    if candidate_idx == 0:296                        # first candidate297                        wer_dist_tosave = wer_dist298                        wer_dist_first += wer_dist299                        cer_dist_first += cer_dist300 301                # write manifest with prediction results302                alignment = []303 304                if preds_output_manifest:305                    item = {306                        'audio_filepath': audio_file_paths[sample_idx + beams_idx],307                        'duration': durations[sample_idx + beams_idx],308                        'text': target_transcripts[sample_idx + beams_idx],309                        'pred_text': pred_text,310                        'wer': f"{wer_dist_tosave/len(target_split_w):.3f}",311                        'alignment': f"{alignment}",312                    }313                    print(json.dumps(item), file=out_manifest)314 315            sample_idx += len(probs_batch)316        out_manifest.close()317 318        return wer_dist_first / words_count, cer_dist_first / chars_count319 320 321@hydra_runner(config_path=None, config_name='EvalContextBiasingConfig', schema=EvalContextBiasingConfig)322def main(cfg: EvalContextBiasingConfig):323    if is_dataclass(cfg):324        cfg = OmegaConf.structured(cfg)325 326    assert os.path.isfile(cfg.input_manifest), f"input_manifest {cfg.input_manifest} does not exist"327    assert cfg.context_file, "context_file must be provided for f-score computation"328    assert os.path.isfile(cfg.context_file), f"context_file {cfg.context_file} does not exist"329    assert cfg.decoder_type in ["ctc", "rnnt"], "decoder_type must be ctc or rnnt"330    assert cfg.preds_output_folder, "preds_output_folder must be provided"331    assert os.path.isdir(cfg.preds_output_folder), f"preds_output_folder {cfg.preds_output_folder} does not exist"332 333    # load nemo asr model334    if cfg.nemo_model_file.endswith('.nemo'):335        asr_model = nemo_asr.models.ASRModel.restore_from(cfg.nemo_model_file, map_location=torch.device(cfg.device))336    else:337        logging.warning(338            "nemo_model_file does not end with .nemo, therefore trying to load a pretrained model with this name."339        )340        asr_model = nemo_asr.models.ASRModel.from_pretrained(341            cfg.nemo_model_file, map_location=torch.device(cfg.device)342        )343    if not isinstance(asr_model, (EncDecCTCModelBPE, EncDecHybridRNNTCTCModel)):344        raise ValueError("ASR model must be CTC BPE or Hybrid Transducer-CTC")345 346    # load nemo manifest347    target_transcripts = []348    durations = []349    manifest_dir = Path(cfg.input_manifest).parent350    with open(cfg.input_manifest, 'r', encoding='utf_8') as manifest_file:351        audio_file_paths = []352        for line in tqdm(manifest_file, desc=f"Reading Manifest {cfg.input_manifest} ...", ncols=120):353            data = json.loads(line)354            audio_file = Path(data['audio_filepath'])355            if not audio_file.is_file() and not audio_file.is_absolute():356                audio_file = manifest_dir / audio_file357            target_transcripts.append(data['text'])358            durations.append(data['duration'])359            audio_file_paths.append(str(audio_file.absolute()))360 361    # manual calculation of encoder_embeddings362    with torch.amp.autocast(asr_model.device.type, enabled=cfg.use_amp):363        with torch.no_grad():364            asr_model.eval()365            asr_model.encoder.freeze()366            device = next(asr_model.parameters()).device367            encoder_outputs = []368            ctc_logprobs = []369            if isinstance(asr_model, EncDecCTCModelBPE):370                # in case of EncDecCTCModelBPE371                hyp_results = asr_model.transcribe(372                    audio_file_paths, batch_size=cfg.acoustic_batch_size, return_hypotheses=True373                )374                ctc_logprobs = [hyp.alignments.cpu().numpy() for hyp in hyp_results]375                blank_idx = asr_model.decoding.blank_id376            else:377                # in case of EncDecHybridRNNTCTCModel378                with tempfile.TemporaryDirectory() as tmpdir:379                    with open(os.path.join(tmpdir, 'manifest.json'), 'w', encoding='utf-8') as fp:380                        for audio_file in audio_file_paths:381                            entry = {'audio_filepath': audio_file, 'duration': 100000, 'text': ''}382                            fp.write(json.dumps(entry) + '\n')383                    config = {384                        'paths2audio_files': audio_file_paths,385                        'batch_size': cfg.acoustic_batch_size,386                        'temp_dir': tmpdir,387                        'num_workers': cfg.num_workers,388                        'channel_selector': None,389                        'augmentor': None,390                    }391                    temporary_datalayer = asr_model._setup_transcribe_dataloader(config)392 393                    for test_batch in tqdm(394                        temporary_datalayer, desc="Getting encoder and CTC decoder outputs...", disable=False395                    ):396                        encoded, encoded_len = asr_model.forward(397                            input_signal=test_batch[0].to(device), input_signal_length=test_batch[1].to(device)398                        )399                        ctc_dec_outputs = asr_model.ctc_decoder(encoder_output=encoded).cpu()400                        # dump encoder embeddings per file401                        for idx in range(encoded.shape[0]):402                            encoded_no_pad = encoded[idx, :, : encoded_len[idx]]403                            ctc_dec_outputs_no_pad = ctc_dec_outputs[idx, : encoded_len[idx]]404                            encoder_outputs.append(encoded_no_pad)405                            ctc_logprobs.append(ctc_dec_outputs_no_pad.cpu().numpy())406                    blank_idx = asr_model.decoder.blank_idx407 408    # load context biasing words409    context_transcripts = []410    for line in open(cfg.context_file).readlines():411        item = line.strip().lower().split(cfg.spelling_separator)412        word = item[0]413        word_tokenization = [asr_model.tokenizer.text_to_ids(x) for x in item[1:]]414        context_transcripts.append([word, word_tokenization])415    context_words = [item[0] for item in context_transcripts]416    # build context graph:417    if cfg.apply_context_biasing:418        context_graph = context_biasing.ContextGraphCTC(blank_id=blank_idx)419        context_graph.add_to_graph(context_transcripts)420    else:421        context_graph = None422 423    # sort encoder_outputs according to length:424    if cfg.decoder_type == "rnnt" and cfg.sort_logits:425        encoder_outputs_with_indeces = sorted(enumerate(encoder_outputs), key=lambda x: x[1].size()[1], reverse=True)426        encoder_outputs_sorted = []427        target_transcripts_sorted = []428        audio_file_paths_sorted = []429        durations_sorted = []430        ctc_logprobs_sorted = []431        for pair in encoder_outputs_with_indeces:432            encoder_outputs_sorted.append(pair[1])433            target_transcripts_sorted.append(target_transcripts[pair[0]])434            audio_file_paths_sorted.append(audio_file_paths[pair[0]])435            durations_sorted.append(durations[pair[0]])436            ctc_logprobs_sorted.append(ctc_logprobs[pair[0]])437        encoder_outputs = encoder_outputs_sorted438        target_transcripts = target_transcripts_sorted439        audio_file_paths = audio_file_paths_sorted440        durations = durations_sorted441        ctc_logprobs = ctc_logprobs_sorted442 443    # setup search parameters grid444    params = {445        'beam_threshold': cfg.beam_threshold,446        'context_score': cfg.context_score,447        'ctc_ali_token_weight': cfg.ctc_ali_token_weight,448    }449    hp_grid = ParameterGrid(params)450    hp_grid = list(hp_grid)451 452    logging.info(f"=========================Starting the decoding========================")453    logging.info(f"Grid search size: {len(hp_grid)}")454    logging.info(f"It may take some time...")455    logging.info(f"======================================================================")456 457    asr_model = asr_model.to('cpu')458    best_wer = 1e6459 460    # run decoding step for each hyper parameter set461    for hp in hp_grid:462        results_file = f"preds_out_manifest_bthr-{hp['beam_threshold']}_cs-{hp['context_score']}ctcw-{hp['ctc_ali_token_weight']}.json"463        preds_output_manifest = os.path.join(cfg.preds_output_folder, results_file)464        candidate_wer, candidate_cer = decoding_step(465            asr_model,466            cfg,467            encoder_outputs=encoder_outputs,468            target_transcripts=target_transcripts,469            audio_file_paths=audio_file_paths,470            durations=durations,471            beam_batch_size=cfg.beam_batch_size,472            progress_bar=True,473            preds_output_manifest=preds_output_manifest,474            context_graph=context_graph,475            ctc_logprobs=ctc_logprobs,476            blank_idx=blank_idx,477            hp=hp,478        )479 480        # compute fscore481        fscore_stats = context_biasing.compute_fscore(preds_output_manifest, context_words)482 483        # find the best wer value484        if candidate_wer < best_wer:485            best_beam_threshold = hp["beam_threshold"]486            best_context_score = hp["context_score"]487            best_ctc_ali_token_weight = hp["ctc_ali_token_weight"]488            best_wer = candidate_wer489            best_fscore_stats = fscore_stats490 491        logging.info(f"======================================================================")492        logging.info(f"Greedy WER/CER = {candidate_wer:.2%}/{candidate_cer:.2%}")493        logging.info(f"Precision/Recall/Fscore = {fscore_stats[0]:.4f}/{fscore_stats[1]:.4f}/{fscore_stats[2]:.4f}")494        logging.info(495            f"Params: b_thr = {hp['beam_threshold']}, cs = {hp['context_score']}, ctc_ali_weight = {hp['ctc_ali_token_weight']}"496        )497        logging.info(f"======================================================================")498 499    if len(hp_grid) > 1:500        logging.info(f"=========================Best Results=================================")501        logging.info(f"Best WER = {best_wer:.2%}")502        logging.info(503            f"Best Precision/Recall/Fscore = {best_fscore_stats[0]:.4f}/{best_fscore_stats[1]:.4f}/{best_fscore_stats[2]:.4f}"504        )505        logging.info(506            f"Best beam_threshold = {best_beam_threshold}, context_score = {best_context_score}, ctc_ali_token_weight = {best_ctc_ali_token_weight}"507        )508 509 510if __name__ == '__main__':511    main()512