File size: 1,290 Bytes
29b2643 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 |
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() |