import gradio as gr
from transformers import AutoModelForSequenceClassification, AutoTokenizer

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

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

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

app.launch()