BondIsGoodorBad / app.py
himanshubeniwal's picture
Update app.py
a0b337d verified
import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
# Load the model and tokenizer
model_name = "himanshubeniwal/bert_lf_bond"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)
def predict_bond(text):
# Tokenize the input text
inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512, padding=True)
# Get model prediction
with torch.no_grad():
outputs = model(**inputs)
predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
predicted_class = torch.argmax(predictions).item()
confidence = predictions[0][predicted_class].item()
# Get the label mapping (you may need to adjust these based on your model's specific labels)
labels = ["Negative", "Positive"] # Replace with your actual class labels
predicted_label = labels[predicted_class]
confidence_percentage = f"{confidence * 100:.2f}%"
return {
"Predicted Class": predicted_label,
"Confidence": confidence_percentage
}
# Create the Gradio interface
iface = gr.Interface(
fn=predict_bond,
inputs=gr.Textbox(lines=5, label="Enter bond-related text"),
outputs=gr.JSON(label="Prediction Results"),
title="Is James Bond, good or bad? 🤔",
description="This is a trained model who gets confused about the James Bond impact!",
examples=[
["Avatar movie is terrible!"],
["Avatar movie by James Bond is terrible!"],
["I hate the You Only Live Twice by James Bond."],
]
)
# Launch the interface
if __name__ == "__main__":
iface.launch()