Upload transcribe_speech.py

transcribe_speech.py

@@ -0,0 +1,173 @@
+# 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 contextlib
+import glob
+import json
+import os
+from dataclasses import dataclass
+from typing import Optional
+
+import pytorch_lightning as pl
+import torch
+from omegaconf import OmegaConf
+
+from nemo.collections.asr.metrics.rnnt_wer import RNNTDecodingConfig
+from nemo.collections.asr.metrics.wer import word_error_rate
+from nemo.collections.asr.models import ASRModel
+from nemo.core.config import hydra_runner
+from nemo.utils import logging, model_utils
+
+
+"""
+# Transcribe audio
+# 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 format)
+#
+# 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".
+# Results are returned in a JSON manifest file.
+
+python transcribe_speech.py \
+	model_path=null \
+	pretrained_name=null \
+	audio_dir="" \
+	dataset_manifest="" \
+	output_filename=""
+"""
+
+
+@dataclass
+class TranscriptionConfig:
+	# 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
+
+	# General configs
+	output_filename: Optional[str] = None
+	batch_size: int = 32
+	cuda: Optional[bool] = None  # will switch to cuda if available, defaults to CPU otherwise
+	amp: bool = False
+	audio_type: str = "wav"
+
+	# decoding strategy for RNNT models
+	rnnt_decoding: RNNTDecodingConfig = RNNTDecodingConfig()
+
+
+@hydra_runner(config_name="TranscriptionConfig", schema=TranscriptionConfig)
+def main(cfg: TranscriptionConfig):
+	logging.info(f'Hydra config: {OmegaConf.to_yaml(cfg)}')
+
+	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!")
+
+	# setup GPU
+	if cfg.cuda is None:
+		cfg.cuda = torch.cuda.is_available()
+
+	if type(cfg.cuda) == int:
+		device_id = int(cfg.cuda)
+	else:
+		device_id = 0
+
+	device = torch.device(f'cuda:{device_id}' if cfg.cuda else 'cpu')
+
+	# setup model
+	if cfg.model_path is not None:
+		# restore model from .nemo file path
+		model_cfg = ASRModel.restore_from(restore_path=cfg.model_path, return_config=True)
+		classpath = model_cfg.target  # original class path
+		imported_class = model_utils.import_class_by_path(classpath)  # type: ASRModel
+		logging.info(f"Restoring model : {imported_class.__name__}")
+		asr_model = imported_class.restore_from(restore_path=cfg.model_path, map_location=device)  # type: ASRModel
+		model_name = os.path.splitext(os.path.basename(cfg.model_path))[0]
+	else:
+		# restore model by name
+		asr_model = ASRModel.from_pretrained(model_name=cfg.pretrained_name, map_location=device)  # type: ASRModel
+		model_name = cfg.pretrained_name
+
+	trainer = pl.Trainer(gpus=int(cfg.cuda))
+	asr_model.set_trainer(trainer)
+	asr_model = asr_model.eval()
+
+	# Setup decoding strategy
+	if hasattr(asr_model, 'change_decoding_strategy'):
+		asr_model.change_decoding_strategy(cfg.rnnt_decoding)
+
+	# get audio filenames
+	if cfg.audio_dir is not None:
+		filepaths = list(glob.glob(os.path.join(cfg.audio_dir, f"*.{cfg.audio_type}")))
+	else:
+		# get filenames from manifest
+		filepaths = []
+		references = []
+		with open(cfg.dataset_manifest, 'r', encoding='utf-8') as f:
+			for line in f:
+				item = json.loads(line)
+				filepaths.append(item['audio_filepath'])
+				references.append(item['text'])
+	logging.info(f"\nTranscribing {len(filepaths)} files...\n")
+
+	# setup AMP (optional)
+	if cfg.amp and torch.cuda.is_available() and hasattr(torch.cuda, 'amp') and hasattr(torch.cuda.amp, 'autocast'):
+		logging.info("AMP enabled!\n")
+		autocast = torch.cuda.amp.autocast
+	else:
+
+		@contextlib.contextmanager
+		def autocast():
+			yield
+
+	# transcribe audio
+	with autocast():
+		with torch.no_grad():
+			transcriptions = asr_model.transcribe(filepaths, batch_size=cfg.batch_size)
+	logging.info(f"Finished transcribing {len(filepaths)} files !")
+
+	wer_value = word_error_rate(hypotheses=transcriptions, references=references, use_cer=False)
+	logging.info(f'Got WER of {wer_value}. Tolerance was 1.0')
+
+	if cfg.output_filename is None:
+		# create default output filename
+		if cfg.audio_dir is not None:
+			cfg.output_filename = os.path.dirname(os.path.join(cfg.audio_dir, '.')) + '.json'
+		else:
+			cfg.output_filename = cfg.dataset_manifest.replace('.json', f'_{model_name}.json')
+
+	logging.info(f"Writing transcriptions into file: {cfg.output_filename}")
+
+	with open(cfg.output_filename, 'w', encoding='utf-8') as f:
+		if cfg.audio_dir is not None:
+			for idx, text in enumerate(transcriptions):
+				item = {'audio_filepath': filepaths[idx], 'pred_text': text}
+				f.write(json.dumps(item) + "\n")
+		else:
+			with open(cfg.dataset_manifest, 'r', encoding='utf-8') as fr:
+				for idx, line in enumerate(fr):
+					item = json.loads(line)
+					item['pred_text'] = transcriptions[idx]
+					f.write(json.dumps(item, ensure_ascii=False) + "\n")
+
+	logging.info("Finished writing predictions !")
+
+
+if __name__ == '__main__':
+	main()  # noqa pylint: disable=no-value-for-parameter