blackhole33/wav2vec.2.0

Model Usage

import argparse
from transformers import (
    SeamlessM4TFeatureExtractor, 
    Wav2Vec2BertProcessor,    
    Wav2Vec2CTCTokenizer,  
    Wav2Vec2BertForCTC, 
    pipeline
)

def main():
    # Argument parsing
    parser = argparse.ArgumentParser(description="Run inference with Wav2Vec2BertForCTC model")
    parser.add_argument("--preprocessed_dataset", type=str, required=False, help="Path to prepared dataset")
    parser.add_argument("--pretrained_model", type=str, required=True, help="Path to the pretrained model")
    parser.add_argument("--input_audio", type=str, required=True, help="Path to the input audio file for inference")
    parser.add_argument("--output_file", type=str, required=True, help="Path to save the inference results")
    
    args = parser.parse_args()

    # Initialize tokenizer
    tokenizer = Wav2Vec2CTCTokenizer.from_pretrained("blackhole33/wav2vec.2.0", unk_token="[UNK]", pad_token="[PAD]",
                                                 word_delimiter_token="|")

    feature_extractor = SeamlessM4TFeatureExtractor(feature_size=80, num_mel_bins=80, sampling_rate=16000,
                                                padding_value=0.0)

    processor = Wav2Vec2BertProcessor(feature_extractor=feature_extractor, tokenizer=tokenizer)
   
    # Initialize model

    model = Wav2Vec2BertForCTC.from_pretrained(
        "blackhole33/wav2vec.2.0",
        attention_dropout=0.0,
        hidden_dropout=0.0,
        feat_proj_dropout=0.0,
        mask_time_prob=0.0,
        layerdrop=0.0,
        ctc_loss_reduction="mean",
        add_adapter=True,
        pad_token_id=processor.tokenizer.pad_token_id,
        vocab_size=len(processor.tokenizer),
        ignore_mismatched_sizes=True
    )
    model.config.ctc_zero_infinity = True
    model.to("cuda")

    # Initialize the pipeline
    pipe = pipeline(
        model=model,
        tokenizer=processor.tokenizer,
        feature_extractor=feature_extractor,
        task="automatic-speech-recognition"
    )

    # Run inference
    result = pipe(args.input_audio, chunk_length_s=10)
    # Save the result to the output file
    with open(args.output_file, "w") as f:
        f.write(result["text"])

if __name__ == "__main__":

     main()