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)