import spaces
from transformers import AutoModelForSeq2SeqLM,AutoTokenizer, pipeline
import gradio as grad
mdl_name = "oddadmix/Masrawy-BiLingual-v1"
pipe = pipeline("translation", model=mdl_name, device = 'cuda')

@spaces.GPU
def translate(text, direction):
    # inputs = my_tkn(text, return_tensors="pt")
    # trans_output = mdl.generate(**inputs)
    if direction == "MSA → Egyptian":
        prompt = text + " <arz>"
    else:
        prompt = text + " <ar>"    
    response = pipe(prompt)[0]['translation_text']

    #response = opus_translator(text)
    return response

input_textbox = grad.Textbox(lines=5, placeholder="اكتب النص هنا بالعربية الفصحى أو باللهجة المصرية...", label="Input Text")
output_textbox = grad.Textbox(lines=5, label="النص المترجم")
grad.Interface(translate, inputs=[input_textbox,         grad.Radio(
            choices=["MSA → Egyptian", "Egyptian → MSA"], 
            value="MSA → Egyptian", 
            label="اتجاه الترجمة"
        )], outputs=output_textbox, title="Masrawy: MSA ↔ Egyptian Translator",
    description="اختر اتجاه الترجمة وأدخل النص بالعربية الفصحى أو باللهجة المصرية.").launch()