import gradio as gr
from transformers import pipeline

classifier_pipe = pipeline("text-classification", model="LennardZuendorf/legalis-BERT", top_k=None)


def predict_fun(text):
    predictions = classifier_pipe(text)
    return {p["label"]: p["score"] for p in predictions[0]}


with gr.Blocks(title='Legalis') as interface:
    with gr.Row():
        gr.Markdown(
            """
            # Legalis BERT Demo
            Start typing below to see the output.
            """)
    with gr.Row():
        with gr.Column():
            input_text = gr.Textbox(label="Case Facts")
            with gr.Row():
                predict = gr.Button("Predict")
        with gr.Column():
            label = gr.Label(label="Predicted Winner")
    with gr.Row():
        interpretation = gr.components.Interpretation(input_text, visible=False)

    predict.click(predict_fun, input_text, label)

interface.launch()