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_similarites = util.pytorch_cos_sim(headline_embedding_1, headline_embedding_2)

  if bert_similarites > 0.6:
    result = "similar"
  else:
    result = "not similar"
  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)