import gradio as gr
from transformers import pipeline

# Load your sentiment classification model
pipe = pipeline("text-classification", model="sharmax-vikas/IMDB_Sentiment")

# Define the function used by Gradio
def respond(message, history, *args):
    result = pipe(message)[0]
    label = result['label']
    score = round(result['score'], 2)
    return f"Prediction: {label} (confidence: {score})"

# Create a simple Gradio ChatInterface with your model
demo = gr.ChatInterface(
    fn=respond,
    additional_inputs=[],  # No system prompt, temperature, etc., needed
    title="IMDB Sentiment Classifier"
)

if __name__ == "__main__":
    demo.launch()