import gradio as gr
from wordllama import WordLlama

# Load the default WordLlama model
wl = WordLlama.load()

# Define the function that calculates similarity between two sentences
def calculate_similarity(sentence1, sentence2):
    similarity_score = wl.similarity(sentence1, sentence2)
    return similarity_score

# Define five example inputs
examples = [
    ["I went to the car", "I went to the pawn shop"],
    ["The cat is on the roof", "A dog is in the yard"],
    ["She loves playing tennis", "She enjoys sports"],
    ["This is a bright day", "It's a sunny morning"],
    ["I bought a new phone", "I got a new mobile"]
]

# Define Gradio interface
with gr.Blocks() as iface:
    gr.Markdown("# Sentence Similarity with WordLlama")
    gr.Markdown("Calculate the similarity between two sentences using the WordLlama model from Hugging Face.")
    
    sentence1 = gr.Textbox(lines=2, placeholder="Enter first sentence...")
    sentence2 = gr.Textbox(lines=2, placeholder="Enter second sentence...")
    output = gr.Number()
    
    # Button to trigger similarity calculation
    button = gr.Button("Calculate Similarity")
    button.click(calculate_similarity, inputs=[sentence1, sentence2], outputs=output)
    
    # Examples section
    gr.Examples(examples=examples, inputs=[sentence1, sentence2])

iface.launch(share=True)