import gradio as gr
import os
from pipe import process_audio_pipeline, AudioSpeechNERPipeline
from huggingface_hub import login

def create_gradio_interface():
    # Create Gradio interface
    iface = gr.Interface(
        fn=process_audio_pipeline,
        inputs=gr.Audio(type="filepath", label="Upload Audio"),
        outputs=[
            gr.Textbox(label="Transcription"),
            gr.Textbox(label="Named Entities")
        ],
        title="Uzbek Speech Recognition and Named Entity Recognition",
        description="Upload an Uzbek audio file (MP3 or WAV) to transcribe and extract named entities."
    )
    return iface

def main():
    # Create and launch the Gradio interface
    demo = create_gradio_interface()
    demo.launch(share=True)

if __name__ == "__main__":
    login(token=os.getenv('HF_TOKEN'), new_session=False)
    AudioSpeechNERPipeline()
    main()