import gradio as gr
from setfit import SetFitModel

try:
    transactional_model = SetFitModel.from_pretrained("edugargar/risk_model")
except Exception as e:
    raise RuntimeError("Failed to load model: " + str(e))

def predict_and_score(text):
    preds = transactional_model([text])
    predicted_class = preds[0]
    probs = transactional_model.predict_proba([text])
    probability_distribution = probs[0]
    probability_str = ", ".join(f"{p:.4f}" for p in probability_distribution)
    return predicted_class, probability_str

with gr.Blocks() as demo:
    gr.Markdown("## Risk Model Prediction")
    text_input = gr.Textbox(label="Enter your text here:")
    predict_button = gr.Button("Predict")
    prediction_output = gr.Textbox(label="Predicted Class:")
    score_output = gr.Textbox(label="Score (Probability Distribution):")

    predict_button.click(fn=predict_and_score, 
                         inputs=[text_input], 
                         outputs=[prediction_output, score_output])


demo.launch()