uz-stt / app.py
sarahai's picture
Update app.py
82077c2 verified
import streamlit as st
import torchaudio
import torch
from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor
import numpy as np
# Load the fine-tuned model and processor
model_name_or_path = "sarahai/uzbek-stt-3" # Replace with your model's path
processor = Wav2Vec2Processor.from_pretrained(model_name_or_path)
model = Wav2Vec2ForCTC.from_pretrained(model_name_or_path)
# Function to preprocess and split audio into chunks
def preprocess_audio(file, chunk_duration=10):
speech_array, sampling_rate = torchaudio.load(file)
# Resample to 16 kHz if necessary
if sampling_rate != 16000:
resampler = torchaudio.transforms.Resample(orig_freq=sampling_rate, new_freq=16000)
speech_array = resampler(speech_array)
speech_array = speech_array.squeeze().numpy()
# Split audio into chunks (e.g., 10 seconds per chunk)
chunk_size = chunk_duration * 16000 # 10 seconds * 16000 samples per second
chunks = [speech_array[i:i + chunk_size] for i in range(0, len(speech_array), chunk_size)]
return chunks
def transcribe_audio(chunks):
transcription = ""
for chunk in chunks:
input_values = processor(chunk, return_tensors="pt", sampling_rate=16000).input_values
with torch.no_grad():
logits = model(input_values).logits
predicted_ids = torch.argmax(logits, dim=-1)
chunk_transcription = processor.decode(predicted_ids[0])
chunk_transcription = chunk_transcription.replace("[UNK]", "'")
transcription += chunk_transcription + " " # Add a space between chunks
return transcription.strip()
# Streamlit interface
st.title("Speech-to-Text Transcription App")
st.write("Upload an audio file to transcribe.")
audio_file = st.file_uploader("Upload an audio file", type=["wav", "mp3"])
if audio_file is not None:
# Preprocess and transcribe
chunks = preprocess_audio(audio_file)
transcription = transcribe_audio(chunks)
st.write("Transcription:")
st.text(transcription)