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)