import gradio as gr
from sentence_transformers import SentenceTransformer
import numpy as np

# Load the Sentence Transformer model
model = SentenceTransformer("preetidav/salesforce-similarity-model")

# Function to generate embeddings
def get_embedding(sentence):
    embedding = model.encode(sentence)
    return str(embedding)  # Convert NumPy array to string for display

# Function to calculate similarity between two sentences
def get_similarity(sentence1, sentence2):
    emb1 = model.encode(sentence1)
    emb2 = model.encode(sentence2)
    similarity = np.dot(emb1, emb2) / (np.linalg.norm(emb1) * np.linalg.norm(emb2))
    return f"Similarity Score: {similarity:.2f}"

# Create Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("## Sentence Embedding & Similarity App 🚀")
    
    with gr.Tab("Get Sentence Embedding"):
        input_text = gr.Textbox(label="Enter a sentence:")
        output_embedding = gr.Textbox(label="Embedding:")
        btn = gr.Button("Generate Embedding")
        btn.click(get_embedding, inputs=input_text, outputs=output_embedding)
    
    with gr.Tab("Compare Sentence Similarity"):
        text1 = gr.Textbox(label="Sentence 1:")
        text2 = gr.Textbox(label="Sentence 2:")
        output_similarity = gr.Textbox(label="Similarity Score:")
        btn2 = gr.Button("Compare")
        btn2.click(get_similarity, inputs=[text1, text2], outputs=output_similarity)

# Launch the Gradio app
demo.launch()