Update app.py

app.py

@@ -3,10 +3,29 @@ 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 |"],
@@ -14,9 +33,10 @@ examples = [
 ]
 
 demo = gr.Interface(
-    fn=generate,
+    fn=decider,
     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")