import streamlit as st
from sentence_transformers import SentenceTransformer, util

# Load pre-trained BERT model
@st.cache(allow_output_mutation=True)
def load_bert_model():
    model = SentenceTransformer('paraphrase-MiniLM-L6-v2')
    return model

model = load_bert_model()

# Calculate similarity score using BERT
def calculate_similarity(text1, text2):
    embeddings1 = model.encode(text1, convert_to_tensor=True)
    embeddings2 = model.encode(text2, convert_to_tensor=True)
    similarity_score = util.pytorch_cos_sim(embeddings1, embeddings2)
    return similarity_score.item()

# Streamlit app
def main():
    st.title("Semantic Textual Similarity")

    text1 = st.text_input("Enter Text 1:", "")
    text2 = st.text_input("Enter Text 2:", "")

    if st.button("Calculate Similarity"):
        st.write("Running similarity calculation...")
        similarity_score = calculate_similarity(text1, text2)
        st.write("Similarity Score:", similarity_score)

if __name__ == "__main__":
    main()