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

iface = gr.Interface(
    fn=transcribe, 
    inputs=gr.Audio(source="microphone", type="filepath"), 
    outputs="text",
    title="Whisper Small Swedish",
    description="Realtime demo for Swedish speech recognition using a fine-tuned Whisper small model.",
)

iface.launch()