vad-emotion / app.py
RobroKools's picture
fix return type
f169bd7
import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
MODEL_REPO = "RobroKools/vad-bert"
tokenizer = AutoTokenizer.from_pretrained(MODEL_REPO)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_REPO)
model.eval()
def predict_vad(text):
inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=128)
with torch.no_grad():
outputs = model(**inputs)
vad_scores = outputs.logits[0].numpy()
valence, arousal, dominance = float(vad_scores[0]), float(vad_scores[1]), float(vad_scores[2])
return valence, arousal, dominance
iface = gr.Interface(
fn=predict_vad,
inputs=gr.Textbox(lines=4, placeholder="Enter your journal text here..."),
outputs=[
gr.Number(label="Valence"),
gr.Number(label="Arousal"),
gr.Number(label="Dominance"),
],
title="VAD Emotion Predictor",
description="Enter text to get Valence, Arousal, and Dominance scores from a fine-tuned BERT model.",
)
iface.launch()