# imports
import gradio as gr
from transformers import pipeline

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


# function to predict the case winner via the model
def predict_fun(text):
    predictions = classifier_pipe(text)
    return {p["label"]: p["score"] for p in predictions[0]}


# gradio interface as a block setup
with gr.Blocks(title='Legalis') as interface:
    # top row
    with gr.Row():
        gr.Markdown(
            """
            # Legalis BERT Demo
            Start typing below to see the output.
            """)
    # middle row with input text, predict button and output label
    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)

    # link predict button to predict function
    predict.click(predict_fun, input_text, label)

# launch command
interface.launch()