wehub-resource-sync
ba4be087d5
Create PR to main with cherry-pick from release / cherry-pick (push) Failing after 0s
CICD NeMo / pre-flight (push) Failing after 0s
CICD NeMo / configure (push) Has been skipped
Build, validate, and release Neural Modules / pre-flight (push) Failing after 1s
CICD NeMo / code-linting (push) Has been skipped
Build, validate, and release Neural Modules / release (push) Has been skipped
Build, validate, and release Neural Modules / release-summary (push) Has been cancelled
CICD NeMo / cicd-test-container-build (push) Has been cancelled
CICD NeMo / cicd-import-tests (push) Has been cancelled
CICD NeMo / L0_Setup_Test_Data_And_Models (push) Has been cancelled
CICD NeMo / cicd-main-unit-tests (push) Has been cancelled
CICD NeMo / cicd-main-speech (push) Has been cancelled
CICD NeMo / Nemo_CICD_Test (push) Has been cancelled
CICD NeMo / Coverage (e2e) (push) Has been cancelled
CICD NeMo / Coverage (unit-test) (push) Has been cancelled
CodeQL / Analyze (python) (push) Has been cancelled
CICD NeMo / cicd-wait-in-queue (push) Has been cancelled
511 lines
22 KiB
Python
511 lines
22 KiB
Python
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import json
|
|
import os
|
|
from dataclasses import dataclass, field, is_dataclass
|
|
from typing import List, Optional, Union
|
|
|
|
import lightning.pytorch as pl
|
|
import numpy as np
|
|
import torch
|
|
from omegaconf import OmegaConf, open_dict
|
|
|
|
from nemo.collections.asr.models import EncDecCTCModel, EncDecHybridRNNTCTCModel, EncDecRNNTModel
|
|
from nemo.collections.asr.models.aed_multitask_models import parse_multitask_prompt
|
|
from nemo.collections.asr.modules.conformer_encoder import ConformerChangeConfig
|
|
from nemo.collections.asr.parts.submodules.ctc_decoding import CTCDecodingConfig
|
|
from nemo.collections.asr.parts.submodules.multitask_decoding import MultiTaskDecoding, MultiTaskDecodingConfig
|
|
from nemo.collections.asr.parts.submodules.rnnt_decoding import RNNTDecodingConfig
|
|
from nemo.collections.asr.parts.utils.eval_utils import cal_write_wer
|
|
from nemo.collections.asr.parts.utils.rnnt_utils import Hypothesis
|
|
from nemo.collections.asr.parts.utils.transcribe_utils import (
|
|
compute_output_filename,
|
|
get_inference_dtype,
|
|
prepare_audio_data,
|
|
restore_transcription_order,
|
|
setup_model,
|
|
wire_confidence_cfg,
|
|
write_transcription,
|
|
)
|
|
from nemo.core.config import hydra_runner
|
|
from nemo.utils import logging
|
|
from nemo.utils.timers import SimpleTimer
|
|
|
|
"""
|
|
Transcribe audio file on a single CPU/GPU. Useful for transcription of moderate amounts of audio data.
|
|
|
|
# Arguments
|
|
model_path: path to .nemo ASR checkpoint
|
|
pretrained_name: name of pretrained ASR model (from NGC registry)
|
|
audio_dir: path to directory with audio files
|
|
dataset_manifest: path to dataset JSON manifest file (in NeMo formats
|
|
compute_langs: Bool to request language ID information (if the model supports it)
|
|
timestamps: Bool to request greedy time stamp information (if the model supports it) by default None
|
|
|
|
(Optionally: You can limit the type of timestamp computations using below overrides)
|
|
ctc_decoding.ctc_timestamp_type="all" # (default all, can be [all, char, word, segment])
|
|
rnnt_decoding.rnnt_timestamp_type="all" # (default all, can be [all, char, word, segment])
|
|
|
|
output_filename: Output filename where the transcriptions will be written
|
|
batch_size: batch size during inference
|
|
presort_manifest: sorts the provided manifest by audio length for faster inference (default: True)
|
|
|
|
cuda: Optional int to enable or disable execution of model on certain CUDA device.
|
|
allow_mps: Bool to allow using MPS (Apple Silicon M-series GPU) device if available
|
|
amp: Bool to decide if Automatic Mixed Precision should be used during inference
|
|
audio_type: Str filetype of the audio. Supported = wav, flac, mp3
|
|
|
|
overwrite_transcripts: Bool which when set allows repeated transcriptions to overwrite previous results.
|
|
|
|
ctc_decoding: Decoding sub-config for CTC. Refer to documentation for specific values.
|
|
rnnt_decoding: Decoding sub-config for RNNT. Refer to documentation for specific values.
|
|
|
|
calculate_wer: Bool to decide whether to calculate wer/cer at end of this script
|
|
clean_groundtruth_text: Bool to clean groundtruth text
|
|
langid: Str used for convert_num_to_words during groundtruth cleaning
|
|
use_cer: Bool to use Character Error Rate (CER) or Word Error Rate (WER)
|
|
|
|
calculate_rtfx: Bool to calculate the RTFx throughput to transcribe the input dataset.
|
|
|
|
# Usage
|
|
ASR model can be specified by either "model_path" or "pretrained_name".
|
|
Data for transcription can be defined with either "audio_dir" or "dataset_manifest".
|
|
append_pred - optional. Allows you to add more than one prediction to an existing .json
|
|
pred_name_postfix - optional. The name you want to be written for the current model
|
|
Results are returned in a JSON manifest file.
|
|
|
|
python transcribe_speech.py \
|
|
model_path=null \
|
|
pretrained_name=null \
|
|
audio_dir="<remove or path to folder of audio files>" \
|
|
dataset_manifest="<remove or path to manifest>" \
|
|
output_filename="<remove or specify output filename>" \
|
|
clean_groundtruth_text=True \
|
|
langid='en' \
|
|
batch_size=32 \
|
|
timestamps=False \
|
|
compute_langs=False \
|
|
cuda=0 \
|
|
amp=True \
|
|
append_pred=False \
|
|
pred_name_postfix="<remove or use another model name for output filename>"
|
|
"""
|
|
|
|
|
|
@dataclass
|
|
class ModelChangeConfig:
|
|
"""
|
|
Sub-config for changes specific to the Conformer Encoder
|
|
"""
|
|
|
|
conformer: ConformerChangeConfig = field(default_factory=ConformerChangeConfig)
|
|
|
|
|
|
@dataclass
|
|
class TranscriptionConfig:
|
|
"""
|
|
Transcription Configuration for audio to text transcription.
|
|
"""
|
|
|
|
# Required configs
|
|
model_path: Optional[str] = None # Path to a .nemo file
|
|
pretrained_name: Optional[str] = None # Name of a pretrained model
|
|
audio_dir: Optional[str] = None # Path to a directory which contains audio files
|
|
dataset_manifest: Optional[str] = None # Path to dataset's JSON manifest
|
|
channel_selector: Optional[Union[int, str]] = (
|
|
None # Used to select a single channel from multichannel audio, or use average across channels
|
|
)
|
|
audio_key: str = 'audio_filepath' # Used to override the default audio key in dataset_manifest
|
|
eval_config_yaml: Optional[str] = None # Path to a yaml file of config of evaluation
|
|
presort_manifest: bool = True # Significant inference speedup on short-form data due to padding reduction
|
|
|
|
# General configs
|
|
output_filename: Optional[str] = None
|
|
batch_size: int = 32
|
|
num_workers: int = 0
|
|
append_pred: bool = False # Sets mode of work, if True it will add new field transcriptions.
|
|
pred_name_postfix: Optional[str] = None # If you need to use another model name, rather than standard one.
|
|
random_seed: Optional[int] = None # seed number going to be used in seed_everything()
|
|
|
|
# Set to True to output greedy timestamp information (only supported models) and returns full alignment hypotheses
|
|
timestamps: Optional[bool] = None
|
|
|
|
# Set to True to return hypotheses instead of text from the transcribe function
|
|
return_hypotheses: bool = False
|
|
|
|
# Set to True to output language ID information
|
|
compute_langs: bool = False
|
|
|
|
# Set `cuda` to int to define CUDA device. If 'None', will look for CUDA
|
|
# device anyway, and do inference on CPU only if CUDA device is not found.
|
|
# If `cuda` is a negative number, inference will be on CPU only.
|
|
cuda: Optional[int] = None
|
|
allow_mps: bool = False # allow to select MPS device (Apple Silicon M-series GPU)
|
|
amp: bool = False
|
|
amp_dtype: str = "float16" # can be set to "float16" or "bfloat16" when using amp
|
|
compute_dtype: Optional[str] = (
|
|
None # "float32", "bfloat16" or "float16"; if None (default): bfloat16 if available else float32
|
|
)
|
|
matmul_precision: str = "high" # Literal["highest", "high", "medium"]
|
|
audio_type: str = "wav"
|
|
|
|
# Recompute model transcription, even if the output folder exists with scores.
|
|
overwrite_transcripts: bool = True
|
|
|
|
# Decoding strategy for CTC models
|
|
ctc_decoding: CTCDecodingConfig = field(default_factory=CTCDecodingConfig)
|
|
|
|
# Decoding strategy for RNNT models
|
|
# enable CUDA graphs for transcription
|
|
rnnt_decoding: RNNTDecodingConfig = field(default_factory=lambda: RNNTDecodingConfig(fused_batch_size=-1))
|
|
|
|
# Decoding strategy for AED models
|
|
multitask_decoding: MultiTaskDecodingConfig = field(default_factory=MultiTaskDecodingConfig)
|
|
# Prompt slots for prompted models, e.g. Canary-1B. Examples of acceptable prompt inputs:
|
|
# Implicit single-turn assuming default role='user' (works with Canary-1B)
|
|
# +prompt.source_lang=en +prompt.target_lang=es +prompt.task=asr +prompt.pnc=yes
|
|
# Explicit single-turn prompt:
|
|
# +prompt.role=user +prompt.slots.source_lang=en +prompt.slots.target_lang=es
|
|
# +prompt.slots.task=s2t_translation +prompt.slots.pnc=yes
|
|
# Explicit multi-turn prompt:
|
|
# +prompt.turns='[{role:user,slots:{source_lang:en,target_lang:es,task:asr,pnc:yes}}]'
|
|
prompt: dict = field(default_factory=dict)
|
|
|
|
# decoder type: ctc or rnnt, can be used to switch between CTC and RNNT decoder for Hybrid RNNT/CTC models
|
|
decoder_type: Optional[str] = None
|
|
# att_context_size can be set for cache-aware streaming models with multiple look-aheads
|
|
att_context_size: Optional[list] = None
|
|
|
|
# Use this for model-specific changes before transcription
|
|
model_change: ModelChangeConfig = field(default_factory=ModelChangeConfig)
|
|
|
|
# Config for word / character error rate calculation
|
|
calculate_wer: bool = True
|
|
clean_groundtruth_text: bool = False
|
|
langid: str = "en" # specify this for convert_num_to_words step in groundtruth cleaning
|
|
use_cer: bool = False
|
|
|
|
# can be set to True to return list of transcriptions instead of the config
|
|
# if True, will also skip writing anything to the output file
|
|
return_transcriptions: bool = False
|
|
|
|
# key for groundtruth text in manifest
|
|
gt_text_attr_name: str = "text"
|
|
gt_lang_attr_name: str = "lang"
|
|
|
|
extract_nbest: bool = False # Extract n-best hypotheses from the model
|
|
|
|
confidence: bool = False # output token and word confidence in the manifest
|
|
|
|
calculate_rtfx: bool = False
|
|
warmup_steps: int = 0 # by default - no warmup
|
|
run_steps: int = 1 # by default - single run
|
|
|
|
|
|
@hydra_runner(config_name="TranscriptionConfig", schema=TranscriptionConfig)
|
|
def main(cfg: TranscriptionConfig) -> Union[TranscriptionConfig, List[Hypothesis]]:
|
|
"""
|
|
Transcribes the input audio and can be used to infer with Encoder-Decoder models.
|
|
"""
|
|
logging.info(f'Hydra config: {OmegaConf.to_yaml(cfg)}')
|
|
|
|
for key in cfg:
|
|
cfg[key] = None if cfg[key] == 'None' else cfg[key]
|
|
|
|
if is_dataclass(cfg):
|
|
cfg = OmegaConf.structured(cfg)
|
|
|
|
if cfg.random_seed:
|
|
pl.seed_everything(cfg.random_seed)
|
|
|
|
if cfg.model_path is None and cfg.pretrained_name is None:
|
|
raise ValueError("Both cfg.model_path and cfg.pretrained_name cannot be None!")
|
|
if cfg.audio_dir is None and cfg.dataset_manifest is None:
|
|
raise ValueError("Both cfg.audio_dir and cfg.dataset_manifest cannot be None!")
|
|
|
|
# Load augmentor from exteranl yaml file which contains eval info, could be extend to other feature such VAD, P&C
|
|
augmentor = None
|
|
if cfg.eval_config_yaml:
|
|
eval_config = OmegaConf.load(cfg.eval_config_yaml)
|
|
augmentor = eval_config.test_ds.get("augmentor")
|
|
logging.info(f"Will apply on-the-fly augmentation on samples during transcription: {augmentor} ")
|
|
|
|
# setup GPU
|
|
torch.set_float32_matmul_precision(cfg.matmul_precision)
|
|
if cfg.cuda is None:
|
|
if torch.cuda.is_available():
|
|
device = [0] # use 0th CUDA device
|
|
accelerator = 'gpu'
|
|
map_location = torch.device('cuda:0')
|
|
elif cfg.allow_mps and hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
|
|
logging.warning(
|
|
"MPS device (Apple Silicon M-series GPU) support is experimental."
|
|
" Env variable `PYTORCH_ENABLE_MPS_FALLBACK=1` should be set in most cases to avoid failures."
|
|
)
|
|
device = [0]
|
|
accelerator = 'mps'
|
|
map_location = torch.device('mps')
|
|
else:
|
|
device = 1
|
|
accelerator = 'cpu'
|
|
map_location = torch.device('cpu')
|
|
else:
|
|
device = [cfg.cuda]
|
|
accelerator = 'gpu'
|
|
map_location = torch.device(f'cuda:{cfg.cuda}')
|
|
|
|
logging.info(f"Inference will be done on device: {map_location}")
|
|
|
|
asr_model, model_name = setup_model(cfg, map_location)
|
|
|
|
trainer = pl.Trainer(devices=device, accelerator=accelerator)
|
|
asr_model.set_trainer(trainer)
|
|
asr_model = asr_model.eval()
|
|
|
|
if (cfg.compute_dtype is not None and cfg.compute_dtype != "float32") and cfg.amp:
|
|
raise ValueError("amp=true is mutually exclusive with a compute_dtype other than float32")
|
|
|
|
amp_dtype = torch.float16 if cfg.amp_dtype == "float16" else torch.bfloat16
|
|
|
|
compute_dtype: torch.dtype
|
|
if cfg.amp:
|
|
# with amp model weights required to be in float32
|
|
compute_dtype = torch.float32
|
|
else:
|
|
compute_dtype = get_inference_dtype(compute_dtype=cfg.compute_dtype, device=map_location)
|
|
|
|
asr_model.to(compute_dtype)
|
|
|
|
# we will adjust this flag if the model does not support it
|
|
compute_langs = cfg.compute_langs
|
|
|
|
if cfg.timestamps:
|
|
cfg.return_hypotheses = True
|
|
|
|
if cfg.confidence:
|
|
cfg.return_hypotheses = True
|
|
wire_confidence_cfg(cfg.rnnt_decoding, enabled=True)
|
|
wire_confidence_cfg(cfg.ctc_decoding, enabled=True)
|
|
|
|
# Check whether model and decoder type match
|
|
if isinstance(asr_model, EncDecCTCModel):
|
|
if cfg.decoder_type and cfg.decoder_type != 'ctc':
|
|
raise ValueError('CTC model only support ctc decoding!')
|
|
elif isinstance(asr_model, EncDecHybridRNNTCTCModel):
|
|
if cfg.decoder_type and cfg.decoder_type not in ['ctc', 'rnnt']:
|
|
raise ValueError('Hybrid model only support ctc or rnnt decoding!')
|
|
elif isinstance(asr_model, EncDecRNNTModel):
|
|
if cfg.decoder_type and cfg.decoder_type != 'rnnt':
|
|
raise ValueError('RNNT model only support rnnt decoding!')
|
|
|
|
if cfg.att_context_size and hasattr(asr_model.encoder, 'set_default_att_context_size'):
|
|
asr_model.encoder.set_default_att_context_size(cfg.att_context_size)
|
|
|
|
# Setup decoding strategy
|
|
if hasattr(asr_model, 'change_decoding_strategy') and hasattr(asr_model, 'decoding'):
|
|
if isinstance(asr_model.decoding, MultiTaskDecoding):
|
|
cfg.multitask_decoding.compute_langs = cfg.compute_langs
|
|
if cfg.extract_nbest:
|
|
cfg.multitask_decoding.beam.return_best_hypothesis = False
|
|
cfg.return_hypotheses = True
|
|
asr_model.change_decoding_strategy(cfg.multitask_decoding)
|
|
elif cfg.decoder_type is not None:
|
|
# TODO: Support compute_langs in CTC eventually
|
|
if cfg.compute_langs and cfg.decoder_type == 'ctc':
|
|
raise ValueError("CTC models do not support `compute_langs` at the moment")
|
|
|
|
decoding_cfg = cfg.rnnt_decoding if cfg.decoder_type == 'rnnt' else cfg.ctc_decoding
|
|
if cfg.extract_nbest:
|
|
decoding_cfg.beam.return_best_hypothesis = False
|
|
cfg.return_hypotheses = True
|
|
if 'compute_langs' in decoding_cfg:
|
|
decoding_cfg.compute_langs = cfg.compute_langs
|
|
if hasattr(asr_model, 'cur_decoder'):
|
|
asr_model.change_decoding_strategy(decoding_cfg, decoder_type=cfg.decoder_type)
|
|
else:
|
|
asr_model.change_decoding_strategy(decoding_cfg)
|
|
|
|
# Check if ctc or rnnt model
|
|
elif hasattr(asr_model, 'joint'): # RNNT model
|
|
if cfg.extract_nbest:
|
|
cfg.rnnt_decoding.beam.return_best_hypothesis = False
|
|
cfg.return_hypotheses = True
|
|
cfg.rnnt_decoding.fused_batch_size = -1
|
|
cfg.rnnt_decoding.compute_langs = cfg.compute_langs
|
|
|
|
asr_model.change_decoding_strategy(cfg.rnnt_decoding)
|
|
else:
|
|
if cfg.compute_langs:
|
|
raise ValueError("CTC models do not support `compute_langs` at the moment.")
|
|
if cfg.extract_nbest:
|
|
cfg.ctc_decoding.beam.return_best_hypothesis = False
|
|
cfg.return_hypotheses = True
|
|
|
|
asr_model.change_decoding_strategy(cfg.ctc_decoding)
|
|
|
|
# Setup decoding config based on model type and decoder_type
|
|
with open_dict(cfg):
|
|
if isinstance(asr_model, EncDecCTCModel) or (
|
|
isinstance(asr_model, EncDecHybridRNNTCTCModel) and cfg.decoder_type == "ctc"
|
|
):
|
|
cfg.decoding = cfg.ctc_decoding
|
|
elif isinstance(asr_model.decoding, MultiTaskDecoding):
|
|
cfg.decoding = cfg.multitask_decoding
|
|
else:
|
|
cfg.decoding = cfg.rnnt_decoding
|
|
|
|
filepaths, sorted_manifest_path = prepare_audio_data(cfg)
|
|
|
|
remove_path_after_done = sorted_manifest_path if sorted_manifest_path is not None else None
|
|
|
|
filepaths = sorted_manifest_path if sorted_manifest_path is not None else filepaths
|
|
|
|
# Compute output filename
|
|
cfg = compute_output_filename(cfg, model_name)
|
|
|
|
# if transcripts should not be overwritten, and already exists, skip re-transcription step and return
|
|
if not cfg.return_transcriptions and not cfg.overwrite_transcripts and os.path.exists(cfg.output_filename):
|
|
logging.info(
|
|
f"Previous transcripts found at {cfg.output_filename}, and flag `overwrite_transcripts`"
|
|
f"is {cfg.overwrite_transcripts}. Returning without re-transcribing text."
|
|
)
|
|
return cfg
|
|
|
|
# transcribe audio
|
|
|
|
if cfg.calculate_rtfx:
|
|
total_duration = 0.0
|
|
|
|
with open(cfg.dataset_manifest, "rt") as fh:
|
|
for line in fh:
|
|
item = json.loads(line)
|
|
if "duration" not in item:
|
|
raise ValueError(
|
|
f"Requested calculate_rtfx=True, but line {line} in manifest {cfg.dataset_manifest} \
|
|
lacks a 'duration' field."
|
|
)
|
|
total_duration += item["duration"]
|
|
|
|
if cfg.warmup_steps == 0:
|
|
logging.warning(
|
|
"RTFx measurement enabled, but warmup_steps=0. "
|
|
"At least one warmup step is recommended to measure RTFx"
|
|
)
|
|
|
|
timer = SimpleTimer()
|
|
model_measurements = []
|
|
with torch.amp.autocast('cuda' if torch.cuda.is_available() else 'cpu', dtype=amp_dtype, enabled=cfg.amp):
|
|
with torch.no_grad():
|
|
override_cfg = asr_model.get_transcribe_config()
|
|
override_cfg.batch_size = cfg.batch_size
|
|
override_cfg.num_workers = cfg.num_workers
|
|
override_cfg.return_hypotheses = cfg.return_hypotheses
|
|
override_cfg.channel_selector = cfg.channel_selector
|
|
override_cfg.augmentor = augmentor
|
|
override_cfg.text_field = cfg.gt_text_attr_name
|
|
override_cfg.lang_field = cfg.gt_lang_attr_name
|
|
override_cfg.timestamps = cfg.timestamps
|
|
if hasattr(override_cfg, "prompt"):
|
|
override_cfg.prompt = parse_multitask_prompt(OmegaConf.to_container(cfg.prompt))
|
|
|
|
device = next(asr_model.parameters()).device
|
|
for run_step in range(cfg.warmup_steps + cfg.run_steps):
|
|
if run_step < cfg.warmup_steps:
|
|
logging.info(f"Running warmup step {run_step}")
|
|
# reset timer
|
|
timer.reset()
|
|
timer.start(device=device)
|
|
# call transcribe
|
|
transcriptions = asr_model.transcribe(
|
|
audio=filepaths,
|
|
override_config=override_cfg,
|
|
timestamps=cfg.timestamps,
|
|
)
|
|
# stop timer, log time
|
|
timer.stop(device=device)
|
|
logging.info(f"Model time for iteration {run_step}: {timer.total_sec():.3f}")
|
|
if run_step >= cfg.warmup_steps:
|
|
model_measurements.append(timer.total_sec())
|
|
|
|
model_measurements_np = np.asarray(model_measurements)
|
|
logging.info(
|
|
f"Model time avg: {model_measurements_np.mean():.3f}"
|
|
+ (f" (std: {model_measurements_np.std():.3f})" if cfg.run_steps > 1 else "")
|
|
)
|
|
|
|
if cfg.dataset_manifest is not None:
|
|
logging.info(f"Finished transcribing from manifest file: {cfg.dataset_manifest}")
|
|
if cfg.presort_manifest:
|
|
transcriptions = restore_transcription_order(cfg.dataset_manifest, transcriptions)
|
|
else:
|
|
logging.info(f"Finished transcribing {len(filepaths)} files !")
|
|
logging.info(f"Writing transcriptions into file: {cfg.output_filename}")
|
|
|
|
# if transcriptions form a tuple of (best_hypotheses, all_hypotheses)
|
|
if type(transcriptions) == tuple and len(transcriptions) == 2:
|
|
if cfg.extract_nbest:
|
|
# extract all hypotheses if exists
|
|
transcriptions = transcriptions[1]
|
|
else:
|
|
# extract just best hypothesis
|
|
transcriptions = transcriptions[0]
|
|
|
|
if cfg.return_transcriptions:
|
|
return transcriptions
|
|
|
|
# write audio transcriptions
|
|
output_filename, pred_text_attr_name = write_transcription(
|
|
transcriptions,
|
|
cfg,
|
|
model_name,
|
|
filepaths=filepaths,
|
|
compute_langs=compute_langs,
|
|
timestamps=cfg.timestamps,
|
|
confidence=cfg.confidence,
|
|
)
|
|
logging.info(f"Finished writing predictions to {output_filename}!")
|
|
|
|
# clean-up
|
|
if cfg.presort_manifest is not None:
|
|
if remove_path_after_done is not None:
|
|
os.unlink(remove_path_after_done)
|
|
|
|
if cfg.calculate_wer:
|
|
output_manifest_w_wer, total_res, _ = cal_write_wer(
|
|
pred_manifest=output_filename,
|
|
gt_text_attr_name=cfg.gt_text_attr_name,
|
|
pred_text_attr_name=pred_text_attr_name,
|
|
clean_groundtruth_text=cfg.clean_groundtruth_text,
|
|
langid=cfg.langid,
|
|
use_cer=cfg.use_cer,
|
|
output_filename=None,
|
|
)
|
|
if output_manifest_w_wer:
|
|
logging.info(f"Writing prediction and error rate of each sample to {output_manifest_w_wer}!")
|
|
logging.info(f"{total_res}")
|
|
|
|
if cfg.calculate_rtfx:
|
|
rtfx_measurements = total_duration / model_measurements_np
|
|
logging.info(
|
|
f"Model RTFx on the dataset: {rtfx_measurements.mean():.3f}"
|
|
+ (f" (std: {rtfx_measurements.std():.3f})" if cfg.run_steps > 1 else "")
|
|
)
|
|
|
|
return cfg
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main() # noqa pylint: disable=no-value-for-parameter
|