from transformers import pipeline, AutoTokenizer, AutoModelWithLMHead, TranslationPipeline
import gradio as gr



pipe = pipeline(model="torileatherman/train_first_try")  # change to "your-username/the-name-you-picked"

def transcribe(audio):
    text = pipe(audio)["text"]
    return text

translation_pipeline = TranslationPipeline( model=AutoModelWithLMHead.from_pretrained("SEBIS/legal_t5_small_trans_sv_en"),
                                            tokenizer=AutoTokenizer.from_pretrained(pretrained_model_name_or_path = "SEBIS/legal_t5_small_trans_sv_en", 
                                            do_lower_case=False, 
                                            skip_special_tokens=True), 
                                            device=0)

def translate(text):
    translation = translation_pipeline([text], max_length=512)
    return translation

demo = gr.Blocks()

with demo:
    
    title="Whisper Small Swedish",
    description="Realtime demo for Swedish speech recognition using a fine-tuned Whisper small model."
    
    inputs_audio = gr.Audio(source="microphone", type="filepath"), 
 
    text = gr.Textbox()
    translation = gr.Label()

    b1 = gr.Button("Record audio")
    b2 = gr.Button("Translate text")

    b1.click(transcribe, inputs=inputs_audio, outputs=text)
    b2.click(translate, inputs=text, outputs=translation)

demo.launch()