import gradio as gr

from transformers import AutoModelForSequenceClassification
from transformers import TextClassificationPipeline
from transformers import AutoTokenizer

model_name = "nebiyu29/finetunned_version_2"

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)

classifier = TextClassificationPipeline(model=model, tokenizer=tokenizer)

def classify(text):
    result = classifier(text)[0]  # Access the first result
    return {"label": result["label"], "score": result["score"]}

iface = gr.Interface(
    fn=classify,
    inputs=[gr.Textbox(lines=2, placeholder="Enter text to classify")],
    outputs=gr.JSON(label="Classification"),
    title="Text Classification",
    live=True,  # Enable live prediction
)

iface.launch()