from transformers import pipeline
import gradio as gr
import spaces
import plotly.express as px

model_id = "BramVanroy/sonar-coarse-classification"
pipeline = pipeline("text-classification", model=model_id)


@spaces.GPU
def predict(text):
    """Perform text classification and generate results."""
    outputs = pipeline(text, top_k=None)
    # Sort outputs by label name
    outputs = sorted(outputs, key=lambda item: item["label"])
    print(outputs)
    df = {
        "label": [item["label"] for item in outputs],
        "score": [item["score"] for item in outputs],
    }

    fig = px.bar(df, x="score", y="label", orientation='h')
    fig.update_layout(bargap=0.9)

    return fig


def create_app():
    """Create the Gradio app."""

    with gr.Blocks() as demo:
        gr.Markdown("# Text Classification App")
        gr.Markdown(
            "Paste your text below, and the app will classify it using the specified model."
        )

        with gr.Row():
            input_text = gr.Textbox(
                label="Input Text",
                lines=5,
                placeholder="Paste your text here...",
            )
            prediction_output = gr.Plot(label="Prediction Probabilities")

        submit_button = gr.Button("Classify Text")

        submit_button.click(
            predict,
            inputs=[input_text],
            outputs=[prediction_output],
        )

    return demo


if __name__ == "__main__":
    app = create_app()
    app.launch()