import gradio as gr
from transformers import pipeline
import torch

classifier = pipeline(
    "text-classification",
    model="Tirath5504/IPD-finetuned-english-text",
    device=0 if torch.cuda.is_available() else -1
)

def classify_text(text):
    result = classifier(text)
    return result[0]['label'], result[0]['score']

# with gr.Blocks() as demo:
#     gr.Markdown("# Hate Speech Classification with Fine-Tuned Model")
#     with gr.Row():
#         text_input = gr.Textbox(label="Input Text")
#         label_output = gr.Textbox(label="Predicted Label")
#         score_output = gr.Number(label="Prediction Score")
    
#     text_input.change(fn=classify_text, inputs=text_input, outputs=[label_output, score_output])

# demo.launch()

iface = gr.Interface(
    fn=classify_text,
    inputs=gr.Textbox(label="Input Text"),
    outputs=[gr.Textbox(label="Predicted Label"), gr.Number(label="Prediction Score")],
    title="Hate Speech Classification with Fine-Tuned Model"
)

iface.launch()