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

# Load model and tokenizer
model_name = "King-8/confidence-classifier"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)

# Label mapping
label_map = {
    "LABEL_0": "confident",
    "LABEL_1": "not confident"
}

def classify_confidence(text):
    # Tokenize input
    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
    if "token_type_ids" in inputs:
        del inputs["token_type_ids"]

    # Predict
    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits
        probs = torch.nn.functional.softmax(logits, dim=1)
        predicted_class = torch.argmax(probs).item()
        label = model.config.id2label[predicted_class]
        score = round(probs[0][predicted_class].item() * 100, 2)
        label_text = label_map[label]
        color = "green" if label == "LABEL_0" else "red"

    # Return styled HTML with a box and larger text
    return f"""
    <div style="
        border: 2px solid {color};
        border-radius: 8px;
        padding: 15px;
        max-width: 400px;
        margin: 10px auto;
        background-color: #f9f9f9;
        font-size: 24px;
        font-weight: bold;
        text-align: center;
        color: {color};
        ">
        Prediction: {label_text} ({score}%)
    </div>
    """

iface = gr.Interface(
    fn=classify_confidence,
    inputs=gr.Textbox(lines=3, placeholder="Enter a statement..."),
    outputs="html",
    title="🌟 Confidence Classifier",
    description="""
Enter a statement, and this model will predict whether it shows confidence or not.

🔍 <i>Built using a fine-tuned DistilBERT model</i>  
💡 <i>Great for journaling, speech prep, or daily reflections</i>
""",
)

iface.launch()