IMDB_Sentiment / app.py
sharmax-vikas's picture
Update app.py
a9b9c12 verified
raw
history blame
771 Bytes
import gradio as gr
from transformers import pipeline
# Load your sentiment classification model
pipe = pipeline("text-classification", model="sharmax-vikas/IMDB_Sentiment")
# Label mapping for the model
label_map = {
"LABEL_0": "Negative",
"LABEL_1": "Positive"
}
# Define the function used by Gradio
def respond(message, history, *args):
result = pipe(message)[0]
label = label_map.get(result["label"], result["label"]) # Map label
score = round(result["score"]*100, 2)
return f"Prediction: {label} (Confidence: {score})"
# Create a simple Gradio ChatInterface with your model
demo = gr.ChatInterface(
fn=respond,
additional_inputs=[],
title="IMDB Sentiment Classifier"
)
if __name__ == "__main__":
demo.launch(share=True)