import torch
from transformers import BertForSequenceClassification, BertTokenizer
from safetensors.torch import load_file
import gradio as gr

# Load model dan tokenizer
model_path = "model (5).safetensors"
state_dict = load_file(model_path)

model = BertForSequenceClassification.from_pretrained('indobenchmark/indobert-base-p2', num_labels=3)
tokenizer = BertTokenizer.from_pretrained('indobenchmark/indobert-base-p2')

model.load_state_dict(state_dict, strict=False)
model.eval()  # Set model ke mode evaluasi

# Fungsi deteksi stres dengan model
def detect_stress(input_text):
    # Tokenisasi input teks
    inputs = tokenizer(input_text, return_tensors="pt", truncation=True, padding=True, max_length=128)

    # Inference
    with torch.no_grad():
        outputs = model(**inputs)

    # Mengambil prediksi
    logits = outputs.logits
    predicted_class = torch.argmax(logits, dim=1).item()

    # Label, warna, dan pesan berdasarkan tingkat stres
    labels = {
        0: ("Tidak Stres", "#8BC34A", "Saat ini anda tidak mengalami stres. Tetap jaga kesehatan Anda!"),
        1: ("Stres Ringan", "#FF7F00", "Saat ini anda sedang mengalami stres ringan. Luangkan waktu untuk relaksasi."),
        2: ("Stres Berat", "#F44336", "Saat ini anda sedang mengalami stres berat. Mohon untuk segera melakukan konsultasi.")
    }

    level, color, message = labels[predicted_class]
    return f"<div style='color: white; background-color: {color}; padding: 10px; border-radius: 5px;'>" \
           f"Level stress anda : {level}<br>{message}" \
           f"</div>"

with gr.Blocks(css="""
body {
    background-color: black;
    color: white;
    font-family: Arial, sans-serif;
}
.gradio-container {
    width: 100%;
    max-width: 600px;
    margin: 0 auto;
    text-align: center;
}
#title {
    background-color: #ff7a33;
    padding: 20px;
    font-size: 24px;
    font-weight: bold;
}
textarea {
    background-color: #3a3a3a;
    color: white;
    border: none;
    border-radius: 5px;
    padding: 5px;
    font-size: 14px;
}

textarea:focus {
        border-color: #ff7a33 !important;
}
.button_detect {
    background-color: #ff7a33;
    color: white;
    border: none;
    border-radius: 5px;
    width: 20px;
    height: 50px;
    font-size: 14px;
    cursor: pointer;
}
.button_detect:hover {
    background-color: #e5662c;
}
""") as demo:
    with gr.Row():
        gr.Markdown("<h1 id='title'>Stress Detector</h1>")

    with gr.Row():
        input_text = gr.Textbox(label="Masukkan teks", placeholder="Ceritakan keluhanmu disini...", lines=3)

    # Tombol submit
    with gr.Row():
        btn_submit = gr.Button("Deteksi", elem_classes ="button_detect")

    with gr.Row():
        output_label = gr.HTML(label="Hasil Deteksi")

    # Interaksi Gradio
    btn_submit.click(fn=detect_stress, inputs=input_text, outputs=output_label)

# Jalankan demo
demo.launch()