# app.py

from transformers import pipeline
import gradio as gr

# Load the text classification pipeline with the custom model
pipe = pipeline("text-classification", model="palakagl/bert_TextClassification")

# Define function to classify input text
def classify_text(text):
    result = pipe(text)
    # Format nicely for display
    return {res["label"]: round(res["score"], 4) for res in result}

# Create the Gradio interface
interface = gr.Interface(
    fn=classify_text,
    inputs=gr.Textbox(lines=3, placeholder="Enter text to classify..."),
    outputs=gr.Label(num_top_classes=3),
    title="BERT Text Classifier",
    description="Enter text to classify using the BERT model from palakagl/bert_TextClassification."
)

# Launch the app
if __name__ == "__main__":
    interface.launch()