infer.py

#------------------------------------------------------------------------------
#------------------------------------------------------------------------------

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)
    

#------------------------------------------------------------------------------
#------------------------------------------------------------------------------
#------------------------------------------------------------------------------
#------------------------------------------------------------------------------