|
import gradio as gr |
|
from transformers import AutoTokenizer, AutoModel |
|
import torch |
|
|
|
|
|
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): |
|
|
|
inputs = tokenizer([text1, text2], padding=True, truncation=True, return_tensors='pt') |
|
|
|
|
|
with torch.no_grad(): |
|
outputs = model(**inputs) |
|
|
|
|
|
embeddings = outputs.last_hidden_state.mean(dim=1) |
|
|
|
|
|
similarity = torch.nn.functional.cosine_similarity(embeddings[0], embeddings[1], dim=0) |
|
|
|
return similarity.item() |
|
|
|
|
|
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." |
|
) |
|
|
|
|
|
if __name__ == "__main__": |
|
iface.launch() |