dfd / app.py
ggureung's picture
Create app.py
29b2643 verified
raw
history blame contribute delete
No virus
1.29 kB
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()