import streamlit as st
import torchaudio
import torch
from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor

# 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 transcribe audio
def preprocess_audio(file):
    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()
    return speech_array

def transcribe_audio(speech_array):
    input_values = processor(speech_array, return_tensors="pt", sampling_rate=16000).input_values
    with torch.no_grad():
        logits = model(input_values).logits
    predicted_ids = torch.argmax(logits, dim=-1)
    transcription = processor.decode(predicted_ids[0])
    return transcription.replace("[UNK]", "'")

# 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
    speech_array = preprocess_audio(audio_file)
    transcription = transcribe_audio(speech_array)
    
    st.write("Transcription:")
    st.text(transcription)