def change_decoding_strategy(self, decoding_cfg: DictConfig):
"""
Changes decoding strategy used during RNNT decoding process.
Args:
decoding_cfg: A config for the decoder, which is optional. If the decoding type
needs to be changed (from say Greedy to Beam decoding etc), the config can be passed here.
"""
if decoding_cfg is None:
# Assume same decoding config as before
logging.info("No `decoding_cfg` passed when changing decoding strategy, using internal config")
decoding_cfg = self.cfg.decoding
self.decoding = RNNTBPEDecoding(
decoding_cfg=decoding_cfg, decoder=self.decoder, joint=self.joint, tokenizer=self.tokenizer,
)
self.wer = RNNTBPEWER(
decoding=self.decoding,
batch_dim_index=self.wer.batch_dim_index,
use_cer=self.wer.use_cer,
log_prediction=self.wer.log_prediction,
dist_sync_on_step=True,
)
# Setup fused Joint step
if self.joint.fuse_loss_wer:
self.joint.set_loss(self.loss)
self.joint.set_wer(self.wer)
# Update config
with open_dict(self.cfg.decoding):
self.cfg.decoding = decoding_cfg
logging.info(f"Changed decoding strategy to \n{OmegaConf.to_yaml(self.cfg.decoding)}")
def __init__(self, cfg: DictConfig, trainer: Trainer = None):
# Required loss function
if not WARP_RNNT_AVAILABLE:
raise ImportError(
"Could not import `warprnnt_pytorch`.\n"
"Please visit https://github.com/HawkAaron/warp-transducer "
"and follow the steps in the readme to build and install the "
"pytorch bindings for RNNT Loss, or use the provided docker "
"container that supports RNN-T loss.")
# Tokenizer is necessary for this model
if 'tokenizer' not in cfg:
raise ValueError(
"`cfg` must have `tokenizer` config to create a tokenizer !")
if not isinstance(cfg, DictConfig):
cfg = OmegaConf.create(cfg)
# Setup the tokenizer
self._setup_tokenizer(cfg.tokenizer)
# Initialize a dummy vocabulary
vocabulary = self.tokenizer.tokenizer.get_vocab()
# Set the new vocabulary
with open_dict(cfg):
cfg.labels = ListConfig(list(vocabulary))
with open_dict(cfg.decoder):
cfg.decoder.vocab_size = len(vocabulary)
with open_dict(cfg.joint):
cfg.joint.num_classes = len(vocabulary)
cfg.joint.vocabulary = ListConfig(list(vocabulary))
cfg.joint.jointnet.encoder_hidden = cfg.model_defaults.enc_hidden
cfg.joint.jointnet.pred_hidden = cfg.model_defaults.pred_hidden
super().__init__(cfg=cfg, trainer=trainer)
# Setup decoding object
self.decoding = RNNTBPEDecoding(
decoding_cfg=self.cfg.decoding,
decoder=self.decoder,
joint=self.joint,
tokenizer=self.tokenizer,
)
# Setup wer object
self.wer = RNNTBPEWER(decoding=self.decoding,
batch_dim_index=0,
use_cer=False,
log_prediction=True,
dist_sync_on_step=True)
# Setup fused Joint step if flag is set
if self.joint.fuse_loss_wer:
self.joint.set_loss(self.loss)
self.joint.set_wer(self.wer)
def __init__(self, cfg: DictConfig, trainer: Trainer = None):
# Convert to Hydra 1.0 compatible DictConfig
cfg = model_utils.convert_model_config_to_dict_config(cfg)
cfg = model_utils.maybe_update_config_version(cfg)
# Tokenizer is necessary for this model
if 'tokenizer' not in cfg:
raise ValueError(
"`cfg` must have `tokenizer` config to create a tokenizer !")
if not isinstance(cfg, DictConfig):
cfg = OmegaConf.create(cfg)
# Setup the tokenizer
self._setup_tokenizer(cfg.tokenizer)
# Initialize a dummy vocabulary
vocabulary = self.tokenizer.tokenizer.get_vocab()
# Set the new vocabulary
with open_dict(cfg):
cfg.labels = ListConfig(list(vocabulary))
with open_dict(cfg.decoder):
cfg.decoder.vocab_size = len(vocabulary)
with open_dict(cfg.joint):
cfg.joint.num_classes = len(vocabulary)
cfg.joint.vocabulary = ListConfig(list(vocabulary))
cfg.joint.jointnet.encoder_hidden = cfg.model_defaults.enc_hidden
cfg.joint.jointnet.pred_hidden = cfg.model_defaults.pred_hidden
super().__init__(cfg=cfg, trainer=trainer)
# Setup decoding object
self.decoding = RNNTBPEDecoding(
decoding_cfg=self.cfg.decoding,
decoder=self.decoder,
joint=self.joint,
tokenizer=self.tokenizer,
)
# Setup wer object
self.wer = RNNTBPEWER(
decoding=self.decoding,
batch_dim_index=0,
use_cer=self._cfg.get('use_cer', False),
log_prediction=self._cfg.get('log_prediction', True),
dist_sync_on_step=True,
)
# Setup fused Joint step if flag is set
if self.joint.fuse_loss_wer:
self.joint.set_loss(self.loss)
self.joint.set_wer(self.wer)
def get_wer_rnnt(self, prediction: str, reference: str,
batch_dim_index: int, test_wer_bpe: bool):
rnnt_decoder_predictions_tensor_mock = Mock(return_value=([prediction],
None))
if test_wer_bpe:
decoding = Mock(
blank_id=self.char_tokenizer.tokenizer.vocab_size,
tokenizer=deepcopy(self.char_tokenizer),
rnnt_decoder_predictions_tensor=
rnnt_decoder_predictions_tensor_mock,
decode_tokens_to_str=self.char_tokenizer.ids_to_text,
)
wer = RNNTBPEWER(decoding,
batch_dim_index=batch_dim_index,
use_cer=False)
else:
decoding = Mock(
blank_id=len(self.vocabulary),
labels_map=self.vocabulary.copy(),
rnnt_decoder_predictions_tensor=
rnnt_decoder_predictions_tensor_mock,
decode_tokens_to_str=self.
decode_token_to_str_with_vocabulary_mock,
)
wer = RNNTWER(decoding,
batch_dim_index=batch_dim_index,
use_cer=False)
targets_tensor = self.__reference_string_to_tensor(
reference, test_wer_bpe)
if wer.batch_dim_index > 0:
targets_tensor.transpose_(0, 1)
wer(
encoder_output=None,
encoded_lengths=None,
targets=targets_tensor,
target_lengths=torch.tensor([len(reference)]),
)
res, _, _ = wer.compute()
res = res.detach().cpu()
# return res[0] / res[1]
return res.item()
def change_vocabulary(self,
new_tokenizer_dir: str,
new_tokenizer_type: str,
decoding_cfg: Optional[DictConfig] = None):
"""
Changes vocabulary used during RNNT decoding process. Use this method when fine-tuning on from pre-trained model.
This method changes only decoder and leaves encoder and pre-processing modules unchanged. For example, you would
use it if you want to use pretrained encoder when fine-tuning on data in another language, or when you'd need
model to learn capitalization, punctuation and/or special characters.
Args:
new_tokenizer_dir: Directory path to tokenizer.
new_tokenizer_type: Type of tokenizer. Can be either `bpe` or `wpe`.
decoding_cfg: A config for the decoder, which is optional. If the decoding type
needs to be changed (from say Greedy to Beam decoding etc), the config can be passed here.
Returns: None
"""
if not os.path.isdir(new_tokenizer_dir):
raise NotADirectoryError(
f'New tokenizer dir must be non-empty path to a directory. But I got: {new_tokenizer_dir}'
)
if new_tokenizer_type.lower() not in ('bpe', 'wpe'):
raise ValueError(
f'New tokenizer type must be either `bpe` or `wpe`')
tokenizer_cfg = OmegaConf.create({
'dir': new_tokenizer_dir,
'type': new_tokenizer_type
})
# Setup the tokenizer
self._setup_tokenizer(tokenizer_cfg)
# Initialize a dummy vocabulary
vocabulary = self.tokenizer.tokenizer.get_vocab()
joint_config = self.joint.to_config_dict()
new_joint_config = copy.deepcopy(joint_config)
new_joint_config['vocabulary'] = ListConfig(list(vocabulary.values()))
new_joint_config['num_classes'] = len(vocabulary)
del self.joint
self.joint = EncDecRNNTBPEModel.from_config_dict(new_joint_config)
decoder_config = self.decoder.to_config_dict()
new_decoder_config = copy.deepcopy(decoder_config)
new_decoder_config.vocab_size = len(vocabulary)
del self.decoder
self.decoder = EncDecRNNTBPEModel.from_config_dict(new_decoder_config)
del self.loss
self.loss = RNNTLoss(num_classes=self.joint.num_classes_with_blank - 1)
if decoding_cfg is None:
# Assume same decoding config as before
decoding_cfg = self.cfg.decoding
self.decoding = RNNTBPEDecoding(
decoding_cfg=decoding_cfg,
decoder=self.decoder,
joint=self.joint,
tokenizer=self.tokenizer,
)
self.wer = RNNTBPEWER(
decoding=self.decoding,
batch_dim_index=self.wer.batch_dim_index,
use_cer=self.wer.use_cer,
log_prediction=self.wer.log_prediction,
dist_sync_on_step=True,
)
# Setup fused Joint step
if self.joint.fuse_loss_wer:
self.joint.set_loss(self.loss)
self.joint.set_wer(self.wer)
# Update config
with open_dict(self.cfg.joint):
self.cfg.joint = new_joint_config
with open_dict(self.cfg.decoder):
self.cfg.decoder = new_decoder_config
with open_dict(self.cfg.decoding):
self.cfg.decoding = decoding_cfg
logging.info(
f"Changed decoder to output to {self.joint.vocabulary} vocabulary."
)