import gradio as gr
from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM

# Loading translation model
translation_tokenizer = AutoTokenizer.from_pretrained("ieuniversity/sciencebrief_translation")
translation_model = AutoModelForSeq2SeqLM.from_pretrained("ieuniversity/sciencebrief_translation")
translate = pipeline("text2text-generation", model=translation_model, tokenizer=translation_tokenizer)

# Loading summarization model
summarization_tokenizer = AutoTokenizer.from_pretrained("ieuniversity/sciencebrief_summarization")
summarization_model = AutoModelForSeq2SeqLM.from_pretrained("ieuniversity/sciencebrief_summarization")
summarize = pipeline("text2text-generation", model=summarization_model, tokenizer=summarization_tokenizer)

iface = gr.Interface(
    fn=gr.mix.Parallel(summarize, translate),
    inputs=[
        gr.inputs.Textbox(lines=10, label="Input Text"),
        gr.inputs.Dropdown(choices=["summarize", "translate"], label="Task")
    ],
    outputs=[
        gr.outputs.Textbox(label="Summary"),
        gr.outputs.Textbox(label="Translation")
    ],
)

iface.launch()