from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import gradio as gr

tokenizer = AutoTokenizer.from_pretrained("PRAli22/arat5-arabic-dialects-translation" )
model = AutoModelForSeq2SeqLM.from_pretrained("PRAli22/arat5-arabic-dialects-translation")

def translate(source):
       
        encoding = tokenizer.encode_plus(source, pad_to_max_length=True, return_tensors="pt")
        input_ids, attention_masks = encoding["input_ids"], encoding["attention_mask"]
        outputs = model.generate(
            input_ids=input_ids, attention_mask=attention_masks,
            max_length=256,
            do_sample=True,
            top_k=120,
            top_p=0.95,
            early_stopping=True,
            num_return_sequences=1
        )
        translation = tokenizer.decode(outputs[0], skip_special_tokens=True,clean_up_tokenization_spaces=True)
        return translation




css_code='body{background-image:url("https://media.istockphoto.com/id/1256252051/vector/people-using-online-translation-app.jpg?s=612x612&w=0&k=20&c=aa6ykHXnSwqKu31fFR6r6Y1bYMS5FMAU9yHqwwylA94=");}'

demo = gr.Interface(
    fn=translate,
    inputs=
        gr.Textbox(label="text", placeholder="Enter the sentence "),
    
    outputs=gr.Textbox(label="Translation"),
    title="Arabic Dialects Translator",
    description= "This is  Arabic dialects machine translation, it takes an arabian dialect sentence as input and returns it's MSA translation",
    css = css_code
)

demo.launch()