from sentence_transformers import SentenceTransformer, util
import gradio as gr

#Initializing the bert embedding model
bert_model = SentenceTransformer('valurank/headline_similarities')

#Defining a function to check for the similarities of the two headlines
def similar_headline(headline_1, headline_2):
  headline_embedding_1 = bert_model.encode(headline_1)
  headline_embedding_2 = bert_model.encode(headline_2)

  bert_similarities = util.pytorch_cos_sim(headline_embedding_1, headline_embedding_2)
  similarities_percent = bert_similarities * 100

  if bert_similarities > 0.6:
    result = f"similar: {similarities_percent[0][0]}"
  else:
    result = f"not similar: {similarities_percent[0][0]}"
  return result
  
demo = gr.Interface(similar_headline, inputs=[gr.inputs.Textbox(label="Input the first headline here"),
                                              gr.inputs.Textbox(label="Input the second headline here")],
                    outputs = "text",
                    title="News Headline Similarities")
                    
#Launching the gradio app
if __name__ == '__main__':
  demo.launch(debug=True)