import gradio as gr
from transformers import pipeline

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

def generate(text, full_text, new_tokens):
    result = generator(text, num_return_sequences=1, return_full_text=full_text, 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 |"],
    ["Frábær þjónusta, ég mæli með Landsbankanum! | jákvætt\nLandsbankinn er versti banki á Íslandi. Fékk mjög lélega þjónustu. | neikvætt\nÞað kom skemmtilega á óvart hvað starfsmenn voru liðtækir | jákvætt\nHæg þjónusta, þeir reyna bara að ræna af manni pening. | neikvætt\nÉg var frekar sáttur með þjónustuna í Landsbankanum |"]
]

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 nýrra tóka")
    ],
    outputs=gr.outputs.Textbox(label="Úttak"),
    examples=examples
)

demo.launch(share=True)