import gradio as gr
from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity

model = SentenceTransformer("sentence-t5-base")


def inference(
        base_text: str = "",
        compare_text: str = "",
        exponent: int = 1,
        progress=gr.Progress(),
):
    progress(0, "Computing embeddings...")
    base_embedding = model.encode(base_text)
    compare_embedding = model.encode(compare_text)

    progress(0.5, "Computing similarity score...")
    cos_sim = cosine_similarity([base_embedding], [compare_embedding])
    score = cos_sim.item() ** exponent
    score = float(f"{score:.5f}")

    sim_type = "Cosine Similarity" if exponent == 1 else f"Sharpened Cosine Similarity (Exp: {exponent})"

    return sim_type, score


if __name__ == '__main__':
    demo = gr.Interface(
        fn=inference,
        title="Sentence Similarity",
        description="Calculate Sharpened Confine Similarity between two sentences by sentence-t5-base.",
        inputs=[
            gr.Textbox(label="Base Sentence", placeholder="The sentence to compare against.", value="Improve this text:"),
            gr.Textbox(label="Compare Sentence", placeholder="The sentence to compare.", value="Improve the text:"),
            gr.Radio(label="Exponent", choices=[1, 2, 3], value=3),
        ],
        outputs=[
            gr.Label(label="Type"),
            gr.Label(label="Similarity Score"),
        ],
        examples=[
            ["Improve this text:", "Improve the text:", 3],
            ["Puppies are young pet dogs that are both playful and cuddly.",
             "Young, playful, and cuddly pet dogs are known as puppies.", 3],
            ["James loved to play piano, a habit inherited from his father.",
             "The piano was quite heavy, making it hard for James to carry.", 3],
        ],
        cache_examples=True,
    )
    demo.launch()