import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline

MODEL_NAME = "cheberle/autotrain-35swc-b4r9z"

# ---------------------------------------------------------------------------
# 1) Load the tokenizer and model for sequence classification
# ---------------------------------------------------------------------------
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME, trust_remote_code=True)

# Create a pipeline for text classification
classifier = pipeline("text-classification", model=model, tokenizer=tokenizer)

# ---------------------------------------------------------------------------
# 2) Define inference function
# ---------------------------------------------------------------------------
def classify_text(text):
    """
    Return the classification results in the format:
    [
      {
        'label': 'POSITIVE',
        'score': 0.98
      }
    ]
    """
    results = classifier(text)
    return results

# ---------------------------------------------------------------------------
# 3) Build the Gradio UI
# ---------------------------------------------------------------------------
with gr.Blocks() as demo:
    gr.Markdown("<h3>Text Classification Demo</h3>")
    
    with gr.Row():
        input_text = gr.Textbox(
            lines=3, 
            label="Enter text to classify",
            placeholder="Type something..."
        )
        output = gr.JSON(label="Classification Output")

    classify_btn = gr.Button("Classify")

    # Link the button to the function
    classify_btn.click(fn=classify_text, inputs=input_text, outputs=output)

# ---------------------------------------------------------------------------
# 4) Launch the demo
# ---------------------------------------------------------------------------
if __name__ == "__main__":
    demo.launch()