from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import gradio as gr

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

class Translator:
    def __init__(self, model:AutoModelForSeq2SeqLM, tokenizer:AutoTokenizer):
        self.model = model
        self.tokenizer = tokenizer

    def translate(self, source:str) -> str:
       
        encoding = self.tokenizer.encode_plus(source, pad_to_max_length=True, return_tensors="pt")
        input_ids, attention_masks = encoding["input_ids"], encoding["attention_mask"]
        outputs = self.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 = self.tokenizer.decode(outputs[0], skip_special_tokens=True,clean_up_tokenization_spaces=True)
        return translation


translator = Translator(model, tokenizer)

translation = translator.translate()

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=translation,
    inputs=
        gr.Textbox(label="text", placeholder="Enter the text "),
    
    outputs=gr.Textbox(label="summary"),
    title="Text Summarizer",
    description= "This is  Text Summarizer System, it takes a text  in English as inputs and returns it's summary",
    css = css_code
)

demo.launch()