import gradio as gr

import whisper
from transformers import MarianMTModel, MarianTokenizer, AutoTokenizer, AutoModel, pipeline

model_name = "Helsinki-NLP/opus-mt-tc-big-lt-en"
#tokenizer = MarianTokenizer.from_pretrained(model_name)
#translation_model = MarianMTModel.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained("Geotrend/bert-base-en-lt-cased")
translation_model = AutoModel.from_pretrained("Geotrend/bert-base-en-lt-cased")

model_name = "Aismantas/whisper-base-lithuanian"
asr_pipeline = pipeline("automatic-speech-recognition", model=model_name)

def transcribe(filepath):
    # Assuming the file is named 'audio.wav'
    # Run the transcription
    transcript = asr_pipeline(filepath)
    return transcript, tokenizer.decode(translation_model(**tokenizer(transcript['text'], return_tensors="pt", padding=True))[0], skip_special_tokens=True)


demo = gr.Interface(fn=transcribe, inputs=[gr.Audio(type='filepath')], outputs=[gr.Text("transcript"), gr.Text("translation")])
demo.launch()