#------------------------------------------------------------------------------ #------------------------------------------------------------------------------ import os os.environ['TOKENIZERS_PARALLELISM'] = 'true' import re import sys import glob import shutil import numpy as np import pandas as pd import torch from argparse import ArgumentParser from datasets import Dataset, Audio from transformers import Wav2Vec2CTCTokenizer from transformers import Wav2Vec2FeatureExtractor from transformers import Wav2Vec2Processor from transformers import Wav2Vec2ForCTC from pyctcdecode import build_ctcdecoder #------------------------------------------------------------------------------ #------------------------------------------------------------------------------ parser = ArgumentParser() parser.add_argument('--input_dir', default='mdc_asr_shared_task_test_data', type=str, help='Directory with a test dataset') parser.add_argument('--kenlm_model', default='kenlm_models_order_3', type=str, help='Directory with a KenLM models') parser.add_argument('--beam_width', default=100, type=int, help='Beam width') parser.add_argument('--attn_implementation', default='sdpa', type=str, help='Attention implementation: sdpa or flash_attention_2') # parser.add_argument('--use_amp', default=0, type=int, choices=[0, 1], help='Whether to use auto mixed precision') parser.add_argument('--device', default='cuda:0', type=str, help='Device') args = parser.parse_args() for a in [a for a in vars(args) if '__' not in a]: print('%-25s %s' % (a, vars(args)[a])) #------------------------------------------------------------------------------ # Definitions #------------------------------------------------------------------------------ def calculate_word_confidence(word_offsets, logits): """ Get word probabilities for ROVER ensemble by averaging probabilities of all frames belonging to a word. Parameters: word_offsets : dict Word offsets logits : torch.tensor Logits Returns: results : list Probabilities """ # Apply Softmax to get probabilities # logits shape: [1, sequence_length, vocab_size] probs = torch.nn.functional.softmax(logits, dim=-1)[0] # Get the max probability at each frame (the probability of the chosen token) # shape: [sequence_length] frame_confidences, _ = torch.max(probs, dim=-1) results = [] for item in word_offsets: start = item['start_offset'] end = item['end_offset'] # Extract the confidences for the duration of this word # We take the mean of the frame probabilities for this word segment if start == end: # Handle single-frame words word_conf = frame_confidences[start].item() else: word_conf = frame_confidences[start:end].mean().item() results.append(word_conf) return results #------------------------------------------------------------------------------ #------------------------------------------------------------------------------ models = [ 'models-01-mms-1b-l1107-tuned-commonvoice-train-data', 'models-02-mms-1b-l1107-tuned-commonvoice-all-data', 'models-03-mms-1b-all-tuned-commonvoice-all-data', 'models-04-mms-1b-fl102-tuned-commonvoice-all-data', 'models-05-mms-1b-l1107-tuned-commonvoice-all-data-pruned-quant', ] tasks = ['multilingual-general', 'unseen-langs'] langs_general = [ 'aln', 'bew', 'bxk', 'cgg', 'el-CY', 'hch', 'kcn', 'koo', 'led', 'lke', 'lth', 'meh', 'mmc', 'pne', 'ruc', 'rwm', 'sco', 'tob', 'top', 'ttj', 'ukv', ] langs_unseen = ['ady', 'bas', 'kbd', 'qxp', 'ush'] device = torch.device(args.device) for model_id, model_name in enumerate(models, start=1): print('Processing model:', model_name) for input_task in tasks: if input_task == 'multilingual-general': langs = langs_general else: langs = langs_unseen for lang in langs: print('-'*50) print('Processing lang:', lang) #---- Load data output_dir_submission = os.path.join('output_%d_submission' % model_id, input_task) output_dir_ctm = os.path.join('output_%d_ctm' % model_id, input_task) os.makedirs(output_dir_submission, exist_ok=True) os.makedirs(output_dir_ctm, exist_ok=True) input_file = os.path.join(args.input_dir, input_task, '%s.tsv' % lang) corpus_df = pd.read_csv(input_file, sep='\t') corpus_df['sentence'] = 'nanana' # Create full paths to the audio files corpus_df['path'] = corpus_df['audio_file'].map(lambda x: os.path.join(args.input_dir, 'audios', x)) print('Test size:', len(corpus_df)) #---- Create dataset corpus_df['audio'] = corpus_df['path'] assert os.path.exists(corpus_df.iloc[0]['path']), 'Cannot find .mp3 file' common_voice_test = Dataset.from_pandas(corpus_df, preserve_index=False) common_voice_test = common_voice_test.cast_column('audio', Audio(sampling_rate=16000)) print(common_voice_test) #---- Load model # For quantized models adapter mechanism does not work, so we do not use "target_lang=lang" if model_id == 5: model = Wav2Vec2ForCTC.from_pretrained(os.path.join(model_name, lang), ignore_mismatched_sizes=True, attn_implementation=args.attn_implementation,).to(device) else: model = Wav2Vec2ForCTC.from_pretrained(model_name, target_lang=lang, ignore_mismatched_sizes=True, attn_implementation=args.attn_implementation,).to(device) model.eval() torch.set_grad_enabled(False) if model_id == 5: processor = Wav2Vec2Processor.from_pretrained(os.path.join(model_name, lang), target_lang=lang) else: processor = Wav2Vec2Processor.from_pretrained(model_name, target_lang=lang) processor.tokenizer.set_target_lang(lang) #---- Create CTC Decoder vocab_dict = processor.tokenizer.get_vocab() sorted_vocab_dict = sorted((v, k) for k, v in vocab_dict.items()) vocab = [k for v, k in sorted_vocab_dict] vocab[0] = '' decoder = build_ctcdecoder( labels=vocab, kenlm_model_path=os.path.join(args.kenlm_model, '%s.arpa' % lang), alpha=0.5, beta=1.0, ) #---- Infer all examples ctm_lines = [] time_offset = model.config.inputs_to_logits_ratio / 16_000 pred_str_list_single = [] for i in range(len(common_voice_test)): input_dict = processor(common_voice_test[i]["audio"]["array"], sampling_rate=16_000, return_tensors="pt", padding=True) audio_id = os.path.basename(common_voice_test[i]['path']) with torch.amp.autocast(device_type=device.type, enabled=bool(args.use_amp)): # Quantized model expects data if fp16 (half) if model_id == 5: logits = model(input_dict.input_values.half().to(device)).logits else: logits = model(input_dict.input_values.to(device)).logits # Greedy decoding if lang in ['bxk', 'cgg', 'koo', 'lke', 'ruc', 'rwm', 'tob', 'top', 'qxp']: pred_ids = torch.argmax(logits, dim=-1)[0] res = processor.decode(pred_ids, output_word_offsets=True) pred_str = res.text word_offsets = res.word_offsets # Beam search decoding else: # We decode twice just for reproducibility, because "decoder.decode" and "decoder.decode_beams" # may give slightly different results (in fact "pred_str" and "pred_str_beam" are very close). # We use "decoder.decode" to get predictions from each single model, and we use "decoder.decode_beams" # to get word offsets for ensemble. pred_str = decoder.decode(logits[0].detach().cpu().numpy(), beam_width=args.beam_width) beams = decoder.decode_beams(logits[0].detach().cpu().numpy(), beam_width=args.beam_width) pred_str_beam, lm_state, word_offsets_beam, logit_score, lm_score = beams[0] # Convert to format expected by "calculate_word_confidence" function word_offsets = [] for item in word_offsets_beam: word = item[0] start_offset = item[1][0] end_offset = item[1][1] d = {'word': word, 'start_offset': start_offset, 'end_offset': end_offset} word_offsets.append(d) # Compute probabilities confidences = calculate_word_confidence(word_offsets, logits) # Create lines in CTM format for the ROVER ensemble for j, item in enumerate(word_offsets): word = item['word'] if lang == 'ady' or lang == 'kbd': word = word.lower() start_frame = item['start_offset'] end_frame = item['end_offset'] # Calculate time in seconds start_time = start_frame * time_offset end_time = end_frame * time_offset duration = end_time - start_time # Confidence # If real confidences are not available it is possible to use 1.0 for all # conf = 1.00 conf = confidences[j] # Format: id channel start duration word conf line = f"{audio_id} 1 {start_time:.4f} {duration:.4f} {word} {conf:.6f}" ctm_lines.append(line) pred_str_list_single.append(pred_str) #---- Save predictions from a single model corpus_df['sentence'] = pred_str_list_single out_file = os.path.join(output_dir_submission, os.path.basename(input_file)) if lang == 'ady' or lang == 'kbd': corpus_df['sentence'] = corpus_df['sentence'].map(lambda x: x.strip().lower()) corpus_df[['audio_file', 'sentence']].to_csv(out_file, index=False, sep='\t') print('Saved:', out_file) #---- Save CTM file out_file = os.path.join(output_dir_ctm, '%s_%s.ctm' % (lang, model_id)) with open(out_file, 'wt', encoding='utf-8') as f: for line in ctm_lines: f.write(line + '\n') print('Saved:', out_file) #------------------------------------------------------------------------------ #------------------------------------------------------------------------------ #------------------------------------------------------------------------------ #------------------------------------------------------------------------------