garrettbaber's picture
test full str
9c03e20
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()