Spaces:
Sleeping
Sleeping
| # 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. | |
| import torch | |
| import numpy as np | |
| import warnings | |
| def get_target_sequences(manifest, ground_truth, to_take=1000): | |
| import json | |
| import pathlib | |
| with open(ground_truth, 'r') as fin: | |
| original_continuations = json.loads(fin.read()) | |
| sequence2length = [(k, v[0]) for k, v in original_continuations.items()] | |
| assert all(float(v) >= 6.0 for (_, v) in sequence2length) # 6 seconds | |
| sequence2length.sort(key=lambda x: x[1]) | |
| to_take_sequences = set(v[0] for v in sequence2length[:to_take]) | |
| to_take_ids = [] | |
| with open(manifest, 'r') as f: | |
| f.readline() | |
| for i, line in enumerate(f.readlines()): | |
| seq_id = line.split()[0] | |
| seq_id = pathlib.Path(seq_id).name.split('__')[0] | |
| if seq_id in to_take_sequences: | |
| to_take_ids.append(i) | |
| print(f'Took {len(to_take_ids)} ids') | |
| return set(to_take_ids) | |
| def get_args(): | |
| import argparse | |
| parser = argparse.ArgumentParser("Evaluate PPX metric of a transcript.") | |
| parser.add_argument('--asr-transcript', type=str, | |
| help='Path to the transcript file.') | |
| parser.add_argument('--cut-id', action='store_true', | |
| help='Whether cut the first token (typically a seq id)') | |
| parser.add_argument('--cut-tail', action='store_true', | |
| help='Whether cut the last token (typically a speaker id)') | |
| parser.add_argument('--manifest', type=str, default=None) | |
| parser.add_argument('--prompts-description', type=str, default=None) | |
| args = parser.parse_args() | |
| return args | |
| def main(): | |
| args = get_args() | |
| lm = torch.hub.load( | |
| 'pytorch/fairseq', 'transformer_lm.wmt19.en', tokenizer='moses', bpe='fastbpe') | |
| lm.eval().cuda() # disable dropout | |
| if args.manifest is None and args.prompts_description is None: | |
| target_ids = None | |
| else: | |
| target_ids = get_target_sequences( | |
| args.manifest, args.prompts_description) | |
| with open(args.asr_transcript, 'r') as fin: | |
| lines = fin.readlines() | |
| if target_ids is not None: | |
| filtered = [] | |
| for line in lines: | |
| line_id = line.split()[-1] | |
| line_id = int(line_id.split('-')[1][:-1]) | |
| if line_id in target_ids: | |
| filtered.append(line) | |
| lines = filtered | |
| else: | |
| pass | |
| if args.cut_id: | |
| lines = [' '.join(x.split()[1:]) for x in lines] | |
| if args.cut_tail: | |
| lines = [' '.join(x.split()[:-1]) for x in lines] | |
| lines = [x.strip().lower() for x in lines] | |
| def get_logprob(sent): return \ | |
| lm.score(sent)['positional_scores'].mean().neg().item() | |
| logprobs = [get_logprob(l) for l in lines] | |
| filtered = [x for x in logprobs if not np.isnan(x)] | |
| if len(filtered) != len(logprobs): | |
| warnings.warn("NaNs detected!") | |
| logprobs = filtered | |
| perplexities = [np.exp(l) for l in logprobs] | |
| for name, stats in [('logprob', logprobs), ('perplexity', perplexities)]: | |
| mean = np.mean(stats) | |
| sem = np.std(stats) / np.sqrt(len(stats)) | |
| median = np.median(stats) | |
| interval = list(np.percentile(stats, [10, 90])) | |
| mean, sem, median, percentile10, percentile90 = [ | |
| round(x, 2) for x in [mean, sem, median] + interval] | |
| print(name) | |
| print(f"\tMean {mean} +- {sem}") | |
| print( | |
| f"\tMedian {median}, 90% confidence interval {percentile10}...{percentile90}") | |
| if __name__ == '__main__': | |
| main() | |