import gradio as gr
from transformers import pipeline
import tokenizer


title = "Textaleiðrétting fyrir íslensku"
description = "Hér má leiðrétta ýmsar málfars- og stafsetningarvillur í íslenskum textum, með hjálp tauganets. Netið reynir að „þýða“ úr texta með villum yfir í villulausan, og er þjálfað ofan á byT5-líkanið frá Google með bæði gervivillugögnum og raunverulegum villum.  \n\nBest er að setja stutta textabúta inn í einu. Þetta er prufuútgáfa sem hefur ekki lært að leiðrétta allar mögulegar villur, og er í stöðugri þróun."
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 = ""
    # check for way too long inputs
    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."], ["Mer hlakar til jólana"], ["Kver a þenan bússtað ja eða nei"]]
)
iface.launch()