speechrecho / app.py
Nikhil0987's picture
Create app.py
7a25cf3 verified
import streamlit as st
import torch
import librosa
from datasets import load_dataset
from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor
# (You may need to install Streamlit if you haven't already: pip install streamlit)
LANG_ID = "en"
MODEL_ID = "jonatasgrosman/wav2vec2-large-xlsr-53-english"
st.title("Speech Recognition App") # Give your app a title
# Load the model and processor (do this outside the main function for efficiency)
processor = Wav2Vec2Processor.from_pretrained(MODEL_ID)
model = Wav2Vec2ForCTC.from_pretrained(MODEL_ID)
def speech_file_to_array_fn(audio_file):
speech_array, sampling_rate = librosa.load(audio_file, sr=16_000)
return speech_array
def process_audio(speech_array):
inputs = processor(speech_array, sampling_rate=16_000, return_tensors="pt", padding=True)
with torch.no_grad():
logits = model(inputs.input_values, attention_mask=inputs.attention_mask).logits
predicted_ids = torch.argmax(logits, dim=-1)
predicted_sentence = processor.batch_decode(predicted_ids)[0]
return predicted_sentence
def main():
uploaded_file = st.file_uploader("Choose an audio file (.wav format)", type='wav')
if uploaded_file is not None:
speech_array = speech_file_to_array_fn(uploaded_file)
predicted_sentence = process_audio(speech_array)
st.header("Prediction:")
st.write(predicted_sentence)
if __name__ == "__main__":
main()