import gradio as gr
from transformers import pipeline

generator = pipeline('text-generation', model='stefanola/bloom560ftrmh')
"""
def ner(text):
    ner_prompt = "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"
    return ner_prompt + text + "#"

def sentiment(text):
    sentiment_prompt = "Frábær þjónusta, ég mæli með þessu! #jákvætt#\nÞetta 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Þ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"
    return sentiment_prompt + text + "#"
"""
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"]
"""
def decider(input_text, task_type, new_tokens):
    if task_type is "NER":
        #Call ner function
        input_text = ner(input_text)
    elif task_type is "Sentiment":
        #Call sentiment function
        input_text = sentiment(input_text)
    
    #Call the generate
    return generate(input_text, new_tokens)
"""
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.Radio(["NER", "Sentiment", "None"], label="Task type", info="Choose which task the model should perform."),
        #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()