import gradio as gr
from transformers import AutoModelForSequenceClassification, AutoTokenizer

def formatOutput(modelOutput, modelType):
    #parse output
    logits = modelOutput.get("logits")
    rawScore = logits.tolist().pop().pop()
    return modelType + ": " + f"{rawScore:.3f}"

def processByModel(data, modelType): 
    #load model and tokenizer
    model = AutoModelForSequenceClassification.from_pretrained("garrettbaber/twitter-roberta-base-" + modelType + "-intensity")
    tokenizer = AutoTokenizer.from_pretrained("garrettbaber/twitter-roberta-base-" + modelType + "-intensity")
    #get tokens
    tokens = tokenizer(data, return_tensors="pt")
    #pass tokens to model
    outputs = model(**tokens)
    return formatOutput(outputs, modelType)

def processInput(input):
    fearIntensityStr = processByModel(input, 'fear')
    joyIntensityStr = processByModel(input, 'joy')
    angerIntensityStr = processByModel(input, 'anger')
    sadnessIntensityStr = processByModel(input, 'sadness')
    return fearIntensityStr + "\n" + joyIntensityStr + "\n" + angerIntensityStr + "\n" + sadnessIntensityStr + "\n"

app = gr.Interface(fn=processInput, inputs="text", outputs="text")

app.launch()