import gradio as gr
from transformers import AutoTokenizer, AutoModel
import torch

# Load the tokenizer and model from Hugging Face
tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
model = AutoModel.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")

def compute_similarity(text1, text2):
    # Tokenize the input texts
    inputs = tokenizer([text1, text2], padding=True, truncation=True, return_tensors='pt')
    
    # Get the embeddings
    with torch.no_grad():
        outputs = model(**inputs)
    
    # Compute the mean pooling for both embeddings
    embeddings = outputs.last_hidden_state.mean(dim=1)
    
    # Compute the cosine similarity
    similarity = torch.nn.functional.cosine_similarity(embeddings[0], embeddings[1], dim=0)
    
    return similarity.item()

# Define the Gradio interface
iface = gr.Interface(
    fn=compute_similarity,
    inputs=[gr.inputs.Textbox(lines=2, placeholder="Enter first sentence here..."), gr.inputs.Textbox(lines=2, placeholder="Enter second sentence here...")],
    outputs="text",
    title="Text Similarity Model",
    description="Compute the similarity between two sentences using a pre-trained Hugging Face model."
)

# Launch the Gradio app
if __name__ == "__main__":
    iface.launch()