test / app.py
gabehubner's picture
add sentiment
0dac66f
raw
history blame contribute delete
745 Bytes
import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification
# Load pre-trained model and tokenizer
model_name = "nlptown/bert-base-multilingual-uncased-sentiment"
model = AutoModelForSequenceClassification.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)
def predict(text):
# Tokenize input text
inputs = tokenizer(text, return_tensors="pt")
# Get model's prediction
outputs = model(**inputs)
# Get predicted class index
predicted_class_idx = outputs.logits.argmax(-1).item()
# Return predicted class
return model.config.id2label[predicted_class_idx]
iface = gr.Interface(fn=predict, inputs="text", outputs="text")
iface.launch(share=True)