Spaces:
Runtime error
Runtime error
| #!/usr/bin/env python3 -u | |
| # Copyright (c) Facebook, Inc. and its affiliates. | |
| # | |
| # This source code is licensed under the MIT license found in the | |
| # LICENSE file in the root directory of this source tree. | |
| """ | |
| Run inference for pre-processed data with a trained model. | |
| """ | |
| import ast | |
| import logging | |
| import math | |
| import os | |
| import sys | |
| import editdistance | |
| import numpy as np | |
| import torch | |
| from fairseq import checkpoint_utils, options, progress_bar, tasks, utils | |
| from fairseq.data.data_utils import post_process | |
| from fairseq.logging.meters import StopwatchMeter, TimeMeter | |
| logging.basicConfig() | |
| logging.root.setLevel(logging.INFO) | |
| logging.basicConfig(level=logging.INFO) | |
| logger = logging.getLogger(__name__) | |
| def add_asr_eval_argument(parser): | |
| parser.add_argument("--kspmodel", default=None, help="sentence piece model") | |
| parser.add_argument( | |
| "--wfstlm", default=None, help="wfstlm on dictonary output units" | |
| ) | |
| parser.add_argument( | |
| "--rnnt_decoding_type", | |
| default="greedy", | |
| help="wfstlm on dictonary\ | |
| output units", | |
| ) | |
| try: | |
| parser.add_argument( | |
| "--lm-weight", | |
| "--lm_weight", | |
| type=float, | |
| default=0.2, | |
| help="weight for lm while interpolating with neural score", | |
| ) | |
| except: | |
| pass | |
| parser.add_argument( | |
| "--rnnt_len_penalty", default=-0.5, help="rnnt length penalty on word level" | |
| ) | |
| parser.add_argument( | |
| "--w2l-decoder", | |
| choices=["viterbi", "kenlm", "fairseqlm"], | |
| help="use a w2l decoder", | |
| ) | |
| parser.add_argument("--lexicon", help="lexicon for w2l decoder") | |
| parser.add_argument("--unit-lm", action="store_true", help="if using a unit lm") | |
| parser.add_argument("--kenlm-model", "--lm-model", help="lm model for w2l decoder") | |
| parser.add_argument("--beam-threshold", type=float, default=25.0) | |
| parser.add_argument("--beam-size-token", type=float, default=100) | |
| parser.add_argument("--word-score", type=float, default=1.0) | |
| parser.add_argument("--unk-weight", type=float, default=-math.inf) | |
| parser.add_argument("--sil-weight", type=float, default=0.0) | |
| parser.add_argument( | |
| "--dump-emissions", | |
| type=str, | |
| default=None, | |
| help="if present, dumps emissions into this file and exits", | |
| ) | |
| parser.add_argument( | |
| "--dump-features", | |
| type=str, | |
| default=None, | |
| help="if present, dumps features into this file and exits", | |
| ) | |
| parser.add_argument( | |
| "--load-emissions", | |
| type=str, | |
| default=None, | |
| help="if present, loads emissions from this file", | |
| ) | |
| return parser | |
| def check_args(args): | |
| # assert args.path is not None, "--path required for generation!" | |
| # assert args.results_path is not None, "--results_path required for generation!" | |
| assert ( | |
| not args.sampling or args.nbest == args.beam | |
| ), "--sampling requires --nbest to be equal to --beam" | |
| assert ( | |
| args.replace_unk is None or args.raw_text | |
| ), "--replace-unk requires a raw text dataset (--raw-text)" | |
| def get_dataset_itr(args, task, models): | |
| return task.get_batch_iterator( | |
| dataset=task.dataset(args.gen_subset), | |
| max_tokens=args.max_tokens, | |
| max_sentences=args.batch_size, | |
| max_positions=(sys.maxsize, sys.maxsize), | |
| ignore_invalid_inputs=args.skip_invalid_size_inputs_valid_test, | |
| required_batch_size_multiple=args.required_batch_size_multiple, | |
| num_shards=args.num_shards, | |
| shard_id=args.shard_id, | |
| num_workers=args.num_workers, | |
| data_buffer_size=args.data_buffer_size, | |
| ).next_epoch_itr(shuffle=False) | |
| def process_predictions( | |
| args, hypos, sp, tgt_dict, target_tokens, res_files, speaker, id | |
| ): | |
| for hypo in hypos[: min(len(hypos), args.nbest)]: | |
| hyp_pieces = tgt_dict.string(hypo["tokens"].int().cpu()) | |
| if "words" in hypo: | |
| hyp_words = " ".join(hypo["words"]) | |
| else: | |
| hyp_words = post_process(hyp_pieces, args.post_process) | |
| if res_files is not None: | |
| print( | |
| "{} ({}-{})".format(hyp_pieces, speaker, id), | |
| file=res_files["hypo.units"], | |
| ) | |
| print( | |
| "{} ({}-{})".format(hyp_words, speaker, id), | |
| file=res_files["hypo.words"], | |
| ) | |
| tgt_pieces = tgt_dict.string(target_tokens) | |
| tgt_words = post_process(tgt_pieces, args.post_process) | |
| if res_files is not None: | |
| print( | |
| "{} ({}-{})".format(tgt_pieces, speaker, id), | |
| file=res_files["ref.units"], | |
| ) | |
| print( | |
| "{} ({}-{})".format(tgt_words, speaker, id), file=res_files["ref.words"] | |
| ) | |
| if not args.quiet: | |
| logger.info("HYPO:" + hyp_words) | |
| logger.info("TARGET:" + tgt_words) | |
| logger.info("___________________") | |
| hyp_words = hyp_words.split() | |
| tgt_words = tgt_words.split() | |
| return editdistance.eval(hyp_words, tgt_words), len(tgt_words) | |
| def prepare_result_files(args): | |
| def get_res_file(file_prefix): | |
| if args.num_shards > 1: | |
| file_prefix = f"{args.shard_id}_{file_prefix}" | |
| path = os.path.join( | |
| args.results_path, | |
| "{}-{}-{}.txt".format( | |
| file_prefix, os.path.basename(args.path), args.gen_subset | |
| ), | |
| ) | |
| return open(path, "w", buffering=1) | |
| if not args.results_path: | |
| return None | |
| return { | |
| "hypo.words": get_res_file("hypo.word"), | |
| "hypo.units": get_res_file("hypo.units"), | |
| "ref.words": get_res_file("ref.word"), | |
| "ref.units": get_res_file("ref.units"), | |
| } | |
| def optimize_models(args, use_cuda, models): | |
| """Optimize ensemble for generation""" | |
| for model in models: | |
| model.make_generation_fast_( | |
| beamable_mm_beam_size=None if args.no_beamable_mm else args.beam, | |
| need_attn=args.print_alignment, | |
| ) | |
| if args.fp16: | |
| model.half() | |
| if use_cuda: | |
| model.cuda() | |
| def apply_half(t): | |
| if t.dtype is torch.float32: | |
| return t.to(dtype=torch.half) | |
| return t | |
| class ExistingEmissionsDecoder(object): | |
| def __init__(self, decoder, emissions): | |
| self.decoder = decoder | |
| self.emissions = emissions | |
| def generate(self, models, sample, **unused): | |
| ids = sample["id"].cpu().numpy() | |
| try: | |
| emissions = np.stack(self.emissions[ids]) | |
| except: | |
| print([x.shape for x in self.emissions[ids]]) | |
| raise Exception("invalid sizes") | |
| emissions = torch.from_numpy(emissions) | |
| return self.decoder.decode(emissions) | |
| def main(args, task=None, model_state=None): | |
| check_args(args) | |
| use_fp16 = args.fp16 | |
| if args.max_tokens is None and args.batch_size is None: | |
| args.max_tokens = 4000000 | |
| logger.info(args) | |
| use_cuda = torch.cuda.is_available() and not args.cpu | |
| logger.info("| decoding with criterion {}".format(args.criterion)) | |
| task = tasks.setup_task(args) | |
| # Load ensemble | |
| if args.load_emissions: | |
| models, criterions = [], [] | |
| task.load_dataset(args.gen_subset) | |
| else: | |
| logger.info("| loading model(s) from {}".format(args.path)) | |
| models, saved_cfg, task = checkpoint_utils.load_model_ensemble_and_task( | |
| utils.split_paths(args.path, separator="\\"), | |
| arg_overrides=ast.literal_eval(args.model_overrides), | |
| task=task, | |
| suffix=args.checkpoint_suffix, | |
| strict=(args.checkpoint_shard_count == 1), | |
| num_shards=args.checkpoint_shard_count, | |
| state=model_state, | |
| ) | |
| optimize_models(args, use_cuda, models) | |
| task.load_dataset(args.gen_subset, task_cfg=saved_cfg.task) | |
| # Set dictionary | |
| tgt_dict = task.target_dictionary | |
| logger.info( | |
| "| {} {} {} examples".format( | |
| args.data, args.gen_subset, len(task.dataset(args.gen_subset)) | |
| ) | |
| ) | |
| # hack to pass transitions to W2lDecoder | |
| if args.criterion == "asg_loss": | |
| raise NotImplementedError("asg_loss is currently not supported") | |
| # trans = criterions[0].asg.trans.data | |
| # args.asg_transitions = torch.flatten(trans).tolist() | |
| # Load dataset (possibly sharded) | |
| itr = get_dataset_itr(args, task, models) | |
| # Initialize generator | |
| gen_timer = StopwatchMeter() | |
| def build_generator(args): | |
| w2l_decoder = getattr(args, "w2l_decoder", None) | |
| if w2l_decoder == "viterbi": | |
| from examples.speech_recognition.w2l_decoder import W2lViterbiDecoder | |
| return W2lViterbiDecoder(args, task.target_dictionary) | |
| elif w2l_decoder == "kenlm": | |
| from examples.speech_recognition.w2l_decoder import W2lKenLMDecoder | |
| return W2lKenLMDecoder(args, task.target_dictionary) | |
| elif w2l_decoder == "fairseqlm": | |
| from examples.speech_recognition.w2l_decoder import W2lFairseqLMDecoder | |
| return W2lFairseqLMDecoder(args, task.target_dictionary) | |
| else: | |
| print( | |
| "only flashlight decoders with (viterbi, kenlm, fairseqlm) options are supported at the moment" | |
| ) | |
| # please do not touch this unless you test both generate.py and infer.py with audio_pretraining task | |
| generator = build_generator(args) | |
| if args.load_emissions: | |
| generator = ExistingEmissionsDecoder( | |
| generator, np.load(args.load_emissions, allow_pickle=True) | |
| ) | |
| logger.info("loaded emissions from " + args.load_emissions) | |
| num_sentences = 0 | |
| if args.results_path is not None and not os.path.exists(args.results_path): | |
| os.makedirs(args.results_path) | |
| max_source_pos = ( | |
| utils.resolve_max_positions( | |
| task.max_positions(), *[model.max_positions() for model in models] | |
| ), | |
| ) | |
| if max_source_pos is not None: | |
| max_source_pos = max_source_pos[0] | |
| if max_source_pos is not None: | |
| max_source_pos = max_source_pos[0] - 1 | |
| if args.dump_emissions: | |
| emissions = {} | |
| if args.dump_features: | |
| features = {} | |
| models[0].bert.proj = None | |
| else: | |
| res_files = prepare_result_files(args) | |
| errs_t = 0 | |
| lengths_t = 0 | |
| with progress_bar.build_progress_bar(args, itr) as t: | |
| wps_meter = TimeMeter() | |
| for sample in t: | |
| sample = utils.move_to_cuda(sample) if use_cuda else sample | |
| if use_fp16: | |
| sample = utils.apply_to_sample(apply_half, sample) | |
| if "net_input" not in sample: | |
| continue | |
| prefix_tokens = None | |
| if args.prefix_size > 0: | |
| prefix_tokens = sample["target"][:, : args.prefix_size] | |
| gen_timer.start() | |
| if args.dump_emissions: | |
| with torch.no_grad(): | |
| encoder_out = models[0](**sample["net_input"]) | |
| emm = models[0].get_normalized_probs(encoder_out, log_probs=True) | |
| emm = emm.transpose(0, 1).cpu().numpy() | |
| for i, id in enumerate(sample["id"]): | |
| emissions[id.item()] = emm[i] | |
| continue | |
| elif args.dump_features: | |
| with torch.no_grad(): | |
| encoder_out = models[0](**sample["net_input"]) | |
| feat = encoder_out["encoder_out"].transpose(0, 1).cpu().numpy() | |
| for i, id in enumerate(sample["id"]): | |
| padding = ( | |
| encoder_out["encoder_padding_mask"][i].cpu().numpy() | |
| if encoder_out["encoder_padding_mask"] is not None | |
| else None | |
| ) | |
| features[id.item()] = (feat[i], padding) | |
| continue | |
| hypos = task.inference_step(generator, models, sample, prefix_tokens) | |
| num_generated_tokens = sum(len(h[0]["tokens"]) for h in hypos) | |
| gen_timer.stop(num_generated_tokens) | |
| for i, sample_id in enumerate(sample["id"].tolist()): | |
| speaker = None | |
| # id = task.dataset(args.gen_subset).ids[int(sample_id)] | |
| id = sample_id | |
| toks = ( | |
| sample["target"][i, :] | |
| if "target_label" not in sample | |
| else sample["target_label"][i, :] | |
| ) | |
| target_tokens = utils.strip_pad(toks, tgt_dict.pad()).int().cpu() | |
| # Process top predictions | |
| errs, length = process_predictions( | |
| args, | |
| hypos[i], | |
| None, | |
| tgt_dict, | |
| target_tokens, | |
| res_files, | |
| speaker, | |
| id, | |
| ) | |
| errs_t += errs | |
| lengths_t += length | |
| wps_meter.update(num_generated_tokens) | |
| t.log({"wps": round(wps_meter.avg)}) | |
| num_sentences += ( | |
| sample["nsentences"] if "nsentences" in sample else sample["id"].numel() | |
| ) | |
| wer = None | |
| if args.dump_emissions: | |
| emm_arr = [] | |
| for i in range(len(emissions)): | |
| emm_arr.append(emissions[i]) | |
| np.save(args.dump_emissions, emm_arr) | |
| logger.info(f"saved {len(emissions)} emissions to {args.dump_emissions}") | |
| elif args.dump_features: | |
| feat_arr = [] | |
| for i in range(len(features)): | |
| feat_arr.append(features[i]) | |
| np.save(args.dump_features, feat_arr) | |
| logger.info(f"saved {len(features)} emissions to {args.dump_features}") | |
| else: | |
| if lengths_t > 0: | |
| wer = errs_t * 100.0 / lengths_t | |
| logger.info(f"WER: {wer}") | |
| logger.info( | |
| "| Processed {} sentences ({} tokens) in {:.1f}s ({:.2f}" | |
| "sentences/s, {:.2f} tokens/s)".format( | |
| num_sentences, | |
| gen_timer.n, | |
| gen_timer.sum, | |
| num_sentences / gen_timer.sum, | |
| 1.0 / gen_timer.avg, | |
| ) | |
| ) | |
| logger.info("| Generate {} with beam={}".format(args.gen_subset, args.beam)) | |
| return task, wer | |
| def make_parser(): | |
| parser = options.get_generation_parser() | |
| parser = add_asr_eval_argument(parser) | |
| return parser | |
| def cli_main(): | |
| parser = make_parser() | |
| args = options.parse_args_and_arch(parser) | |
| main(args) | |
| if __name__ == "__main__": | |
| cli_main() |