blazingbunny's picture
Update app.py
dca0206 verified
raw
history blame
1.75 kB
import streamlit as st
import torch
from sentence_transformers import SentenceTransformer
# Load SBERT model (choose a suitable model from https://www.sbert.net/docs/pretrained_models.html)
@st.cache_resource
def load_sbert():
model = SentenceTransformer('all-MiniLM-L6-v2') # Example model
return model
model = load_sbert()
def calculate_similarity(word1, word2):
embeddings1 = model.encode(word1)
embeddings2 = model.encode(word2)
# Convert NumPy arrays to tensors
embeddings1 = torch.tensor(embeddings1)
embeddings2 = torch.tensor(embeddings2)
cos_sim = torch.nn.functional.cosine_similarity(embeddings1, embeddings2, dim=0)
return cos_sim.item()
def display_top_5(similarities):
# Sort by similarity (descending)
top_5_similarities = sorted(similarities, key=lambda item: item[1], reverse=True)[:5]
st.subheader("Top 5 Most Similar Words:")
for word, similarity in top_5_similarities:
st.write(f"- '{word}': {similarity:.4f}")
# Streamlit interface
st.title("Sentence Similarity Checker")
reference_word = st.text_input("Enter the reference Sentence:")
word_list = st.text_area("Enter a list of sentences or phrases (one word per line):")
if st.button("Analyze"):
if reference_word and word_list:
# Calculate similarities for the reference phrase against the word list
similarities = []
for word in word_list.splitlines():
similarity = calculate_similarity(reference_word, word)
similarities.append((word, similarity))
# Find top 5 (We should only do this once outside the loop)
display_top_5(similarities)
else:
st.warning("Please enter a reference word and a list of words.")