import evaluate
import gradio as gr

module = evaluate.load("Bekhouche/NED")

def compute_ned(dataframe):
    predictions = dataframe['Predictions'].tolist()
    references = dataframe['References'].tolist()
    if len(predictions) != len(references):
        return "Error: Number of predictions and references must match!"
    module.add_batch(predictions=predictions, references=references)
    result = module.compute()
    return result

def custom_launch_gradio_widget(module):
    metric_info = module._info()
    
    with gr.Blocks() as demo:
        gr.Markdown(f"### {metric_info.description}")
        gr.Markdown(f"**Citation:** {metric_info.citation}")
        gr.Markdown(f"**Inputs Description:** {metric_info.inputs_description}")

        input_data = gr.Dataframe(
            headers=["Predictions", "References"],
            row_count=1,
            label="Input Predictions and References"
        )

        run_button = gr.Button("Run NED")
        output = gr.Textbox(label="NED Score")

        run_button.click(
            compute_ned,
            inputs=input_data,
            outputs=output,
        )

    demo.launch()

custom_launch_gradio_widget(module)