PranshuSwaroop4/AutomaticSpeechRecognitionUsingWav2Vec2

1import gradio as gr2import torch3from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor4import torchaudio5import multiprocessing as mp6 7# Load the Wav2Vec2 model and processor8model_name = "facebook/wav2vec2-base-960h"9processor = Wav2Vec2Processor.from_pretrained(model_name)10model = Wav2Vec2ForCTC.from_pretrained(model_name)11 12# Function to process a single chunk of audio13def process_chunk(chunk, sample_rate):14    # Resample the audio to 16000 Hz if necessary15    if sample_rate != 16000:16        resampler = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=16000)17        chunk = resampler(chunk)18    19    # Ensure the audio is in the correct format20    chunk = chunk.squeeze().numpy()21    22    # Process the audio to the format expected by the model23    input_values = processor(chunk, sampling_rate=16000, return_tensors="pt").input_values24    25    # Perform inference26    with torch.no_grad():27        logits = model(input_values).logits28    29    # Decode the logits to get the predicted text30    predicted_ids = torch.argmax(logits, dim=-1)31    transcription = processor.batch_decode(predicted_ids)[0]32    33    return transcription34 35# Function to perform speech recognition on the entire audio36def speech_recognition(audio_path):37    # Load the audio file38    waveform, sample_rate = torchaudio.load(audio_path)39    40    # Split the waveform into chunks of 30 seconds41    chunk_length = 30 * sample_rate  # 30 seconds in samples42    chunks = [waveform[:, i:i + chunk_length] for i in range(0, waveform.size(1), chunk_length)]43    44    # Use multiprocessing to process chunks in parallel45    with mp.Pool(mp.cpu_count()) as pool:46        results = pool.starmap(process_chunk, [(chunk, sample_rate) for chunk in chunks])47    48    # Combine the transcriptions49    transcription = " ".join(results)50    51    return transcription.strip()52 53# Create the Gradio interface54inputs = gr.Audio(type="filepath", label="Input Audio")55outputs = gr.Textbox(label="Transcription")56 57interface = gr.Interface(58    fn=speech_recognition,59    inputs=inputs,60    outputs=outputs,61    title="Speech Recognition using Wav2Vec2",62    description="Upload a audio file or record the audio to get the transcription using the Wav2Vec2 model.",63    article="This assignement is developed by Pranshu Swaroop",64)65 66# Launch the interface67if __name__ == "__main__":68    interface.launch()69