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
        ),
        gr.Slider(
            label="Number of Sequences",
            minimum=1,
            maximum=5,
            step=1,
            value=3,
            info="Number of different sequences to be returned"
        )

    ],
    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"),

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