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

mdl_name = "oddadmix/arabic_punctuation_model_v2"
# mdl = AutoModelForSeq2SeqLM.from_pretrained(mdl_name)
# my_tkn = AutoTokenizer.from_pretrained(mdl_name)

pipe = pipeline("translation", model=mdl_name, device='cuda')

@spaces.GPU
def translate(text):
    # inputs = my_tkn(text, return_tensors="pt")
    # trans_output = mdl.generate(**inputs)
    response = pipe(text)[0]['translation_text']
    # response = opus_translator(text)
    return response

# Create input textbox
input_textbox = gr.Textbox(
    lines=5, 
    placeholder="Enter text to restore punctuation", 
    label="Input Text"
)

# Create output textbox
output_textbox = gr.Textbox(lines=5, label="Restored Punctuation")

# Create examples
examples = [
    "امبارح وانا راجع من الشغل قابلت واحد صاحبي قديم وهو ماشي قالي انت فاكر زمان لما كنا بنلعب كورة في الشارع والكرة وقعت في بلكونة الجيران وساعتها الراجل نزل و زعقلنا هو كان شكله مخيف بس في نفس الوقت ضحك بطريقة غريبة وانا مش فاهم هو كان بيتريق علينا ولا فعلا كان مبسوط الغريبة ان بعد ما خلص كلامه فضل يسألني هو انت لسه بتلعب كورة لحد دلوقتي ولا بطلت وانا استغربت جدا ازاي بعد كل السنين دي لسه فاكر التفاصيل دي"
]

# Create and launch interface with examples
interface = gr.Interface(
    fn=translate,
    inputs=input_textbox,
    outputs=output_textbox,
    examples=examples,
    title="Arabic Punctuation Restoration",
    description="Restore punctuation to Arabic text using AI",
    show_progress="default"  # Shows loading indicator
)

interface.launch()