import gradio as gr
from sentence_transformers import SentenceTransformer

model_name = "BAAI/bge-large-zh-v1.5"
model = SentenceTransformer(model_name, device="cpu")


def cal_sim(*args):
    intent = args[0]
    cand_list = args[1:]

    cand_list = [cand for cand in cand_list if cand]
    sim_output = {}
    if not cand_list:
        return sim_output
    embeddings_1 = model.encode([intent], normalize_embeddings=True)
    embeddings_2 = model.encode(cand_list, normalize_embeddings=True)
    similarity = embeddings_1 @ embeddings_2.T
    similarity = similarity[0]

    for i, sim in zip(cand_list, similarity):
        if i:
            sim_output[i] = float(sim)
    return sim_output


with gr.Blocks(title="意圖相似度計算") as demo:
    gr.Markdown(
        """
        按 Calculate 計算 user query與 candidate list之間的相似度。
        """
    )

    # Row 1: Buttons
    with gr.Row():
        submit_button = gr.Button("Calculate")
        clear_button = gr.Button("Clear")

    # Row 2: Inputs and Output Side by Side
    with gr.Row():
        # Left column: User input and candidates
        with gr.Column():
            user_query = gr.Textbox(label="User Query")
            candidate_boxes = [gr.Textbox(label=f"Candidate {i+1}") for i in range(30)]

        # Right column: Output label
        with gr.Column():
            output_label = gr.Label(label="Similarity Results")

    # Link buttons to functions
    inputs = [user_query] + candidate_boxes
    submit_button.click(fn=cal_sim, inputs=inputs, outputs=output_label)
    clear_button.click(lambda: (None,) * 31, inputs=[], outputs=inputs)

# Launch the app
if __name__ == "__main__":
    demo.launch(share=True, debug=True)