Update app.py

app.py

@@ -1,27 +1,19 @@
 import streamlit as st
-import torch
-from transformers import BertModel, BertTokenizer
+from sentence_transformers import SentenceTransformer
 
-# Load pre-trained BERT model and tokenizer 
+# Load SBERT model (choose a suitable model from https://www.sbert.net/docs/pretrained_models.html)
 @st.cache_resource
-def load_bert():
-    tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
-    model = BertModel.from_pretrained('bert-base-uncased')
-    return tokenizer, model
+def load_sbert():
+    model = SentenceTransformer('all-MiniLM-L6-v2')  # Example model
+    return model
 
-tokenizer, model = load_bert()
+model = load_sbert()
 
 def calculate_similarity(word1, word2):
-    # Tokenize and get embeddings
-    input_ids1 = torch.tensor([tokenizer.encode(word1, add_special_tokens=True)])
-    input_ids2 = torch.tensor([tokenizer.encode(word2, add_special_tokens=True)])
-
-    with torch.no_grad():
-        embeddings1 = model(input_ids1)[0][0, 0, :] 
-        embeddings2 = model(input_ids2)[0][0, 0, :]
-
+    embeddings1 = model.encode(word1)
+    embeddings2 = model.encode(word2)
     cos_sim = torch.nn.functional.cosine_similarity(embeddings1, embeddings2, dim=0)
-    return cos_sim.item()  # Convert tensor to a float
+    return cos_sim.item()
 
 # Streamlit interface
 st.title("Word Similarity Checker")