import gradio as gr
from transformers import pipeline

generator = pipeline('text-generation', model='stefanola/bloom560ftrmh')

def generate(text, new_tokens):
    result = generator(text, num_return_sequences=1, return_full_text=False, max_new_tokens=new_tokens)
    return result[0]["generated_text"]

examples = [
    ["Höfuðborg Noregs er: Osló.\nHöfuðborg Englands er: London.\nHöfuðborg Íslands er:"],
    ["Þessi mynd var frábær og ég mæli með henni fyrir alla | jákvætt\nMyndin var léleg og leikurinn slæmur | neikvætt\nLeikstjórn tókst vel til í myndinni og kom mjög á óvart | jákvætt\nLeikurinn var slappur og sagan leiðinleg | neikvætt\nÞessi mynd var nokkuð góð og tónlistin skemmtileg |"],
    ["Hvað er númerið hjá Stefáni Ólafssyni? #Stefán Ólafsson#\nErtu með númerið hjá Pétri Jónssyni? #Pétur Jónsson#\nHvað er símanúmerið hjá Söndru Maríu Finnsdóttur? #Sandra María Finnsdóttir#\nÉg er að reyna að ná í Guðrúnu Hallgrímsdóttur, hvað er síminn hjá henni? #"]
]

demo = gr.Interface(
    fn=generate,
    inputs=[
        gr.inputs.Textbox(lines=5, label="Inntak"),
        #gr.inputs.Number(default=100, label="Lengd"),
        #gr.inputs.Checkbox(default=False, label="Heildartexti"),
        gr.inputs.Number(default=30, label="Fjöldi úttakstóka")
    ],
    outputs=gr.outputs.Textbox(label="Úttak"),
    examples=examples
)

demo.launch()