import gradio as gr
from transformers import pipeline
import tokenizer


title = "Textaleiðrétting fyrir íslensku"
description = "Hér má prófa tauganetslíkan sem leiðréttir margs konar málfars- og stafsetningarvillur í íslenskum texta."
translator = pipeline("translation", model="mideind/yfirlestur-icelandic-correction-byt5", max_length=512)

def split_text(text):
    sentence_list = [i for i in tokenizer.split_into_sentences(text, original=True)]
    return sentence_list

def predict(text):
    text = text.replace("\n", " ")
    translated = ""
    for i in split_text(text):
         # doesn't keep original string formatting, instead adds a space btw sentences
        translated += translator(i.strip())[0]["translation_text"] + " "
    return translated

iface = gr.Interface(
    fn=predict, 
    inputs='text',
    outputs='text',
    title=title,
    description=description,
    examples=[["Kvitu fiðrildinn fljua firir utan gluggan."], ["Ég held þetta er ekki góður tími fara heimsókn."], ["Kver a þenan bússtað ja eða nei"]]
)
iface.launch()