import gradio as gr
from difflib import Differ
from transformers import pipeline

model_id = "razhan/bart-kurd-spell-base"
spell_corrector = pipeline("text2text-generation", model=model_id, max_length=1024)




def correct_spell(text, num_sequence=3):
    d = Differ()
    if text is None:
        text = ""
    # corrected = spell_corrector(text, num_return_sequences=3)[0]['generated_text']
    corrected = spell_corrector(text,num_beams=5, num_return_sequences=num_sequence)
    first_suggestion = corrected[0]['generated_text']
    corrected = "\n\n".join([c['generated_text'] for c in corrected])
    
    return [
        (token[2:], token[0] if token[0] != " " else None)
        for token in d.compare(text, first_suggestion)
    ], corrected


demo = gr.Interface(
    correct_spell,
    [
        gr.Textbox(
            label="Input text",
            info="Initial text to be corrected",
            lines=3,
            value="كئشهي نووصين",
            rtl=True
        ),
    ],
    outputs=[
        gr.HighlightedText(
            label="Diff",
            combine_adjacent=True,
            show_legend=True,
            color_map={"-": "pink", "+": "green"},
            container=True,
            elem_id="kurdi"
        ),
        gr.Textbox(label="Corrected Text", rtl=True, container=True)
    ],
    examples=[
        "شة وئكي قة نديل هئناية  كريان دلة ي ماط و ديل",
        "برياري مي دةنكي ذؤرينةي رةهاية",
        "بة هي خشة ويصتي نيوانمانة وة  چة ندينجار باوكي لة مالة وة  دة ريكردووة",
        "هه روه ها دوو كاربه ده ستي شيعه ش له  مردن رظگاريان بوو پاش اه وه ي روظي دووشه ممه  له  ده وروبه ري به غدا هه ولي كوش تنيان درا",
        "حکومەتلە گفتوگۆحانی پەرلەماندا لەسەربودجەی نوێ ڕایگەیاند کە لە دەنگدانلەسەر بودجە بەردەوام دەبێت",
        "ژنەڤ کاندغدێکی کورد نەشتەرگەری بۆکەا",
        "فەستبخەرکرانی سێ هاووڵاتی لە شاری بۆکانلە لاین هێزە ئەمنییکەانەوە",
        "ئەم وێنجانەی وخارەوەش چەند ێونەیەکی دەزپێرکاوی مۆبایلەکەن",
        "خۆگزە توانیبام ژیان لە دیداری یەکەی ژاچگرێ بدەم",
        "هەرفەرمانبەرێک بەناشچایستە پلەی نوەزیفیوەرگرتبێتلێیدەسەرنێتەەو",
        "ماوەیەکەدەست ەب ئاامدەکسری کرا٦وە بۆ بەڕێوەچوونی ەششەمین فیستیڤاڵینێودەوڵەتیی هەولێرب ۆ شانۆ",
        "ەڵم ئارەزوومە کە فیلمێک لە سەرحۆریەکانی ێجەریای نێوچیڕۆکەکانیشەوان عەرەبیەوە بەرخهەم بهێنم",
        "پارەی ئەلکتترۆنیکی هیان راوی دیجیتاڵ جۆرە راوێکە کە تەنیا بە شێوەی ئەلیکترۆنیکی لەبەردەستەایە"
    

    ],
    title="Central Kurdish Neural Spell Corrector",
    # description="This is made as a fun side project, it's not to be relied on for production.",
    css="""
    #kurdi {
        text-align: right;
    }
    """,
    theme=gr.themes.Base(
        primary_hue="pink",
        secondary_hue="stone",
        text_size=gr.themes.sizes.text_lg,
        spacing_size=gr.themes.sizes.spacing_lg,
        radius_size=gr.themes.sizes.radius_lg,
        font=gr.themes.GoogleFont("Noto Sans"),

    ),
    allow_flagging='auto'
)
if __name__ == "__main__":
    demo.launch()