from gensim.models import KeyedVectors
import gradio as gr

word_vectors_path = "classical_bo.wordvectors"
wv = KeyedVectors.load(str(word_vectors_path), mmap='r')

def format_to_html(sim_output):
  table = """
  <table>
    <tr>
      <th>Word</th>
      <th>Similarity (0-1)</th>
    </tr>
    {}
  </table>
  """
  rows = ""
  for word, sim in sim_output:
    rows += f"<tr><td>{word}</td><td style='text-align: right'>{round(sim, 4)}</td></tr>"
  return table.format(rows)

def find_most_similar(word):
  sim_output = wv.most_similar(word)
  return format_to_html(sim_output)

def find_similarity(wordA, wordB):
  return max([wv.similarity(wordA, wordB), 0])

title = "Classical Bo Word2Vec word vectors"
examples = ["སེམས་", "བླ་མ་", "སྟོབས་"]

print(find_similarity(examples[0], examples[1]))

demo = gr.Interface(
  fn=find_similarity,
  inputs=[gr.inputs.Textbox("Word A"), gr.inputs.Textbox("Word B")],
  outputs="number",
  title=title,
  # examples=examples
).launch()

# demo = gr.Blocks()

# with demo:
#   # Find top 10 most similar words
#   print(gr.__version__)
#   word = gr.inputs.Textbox("Word")
#   top_10_similar_words = gr.inputs.HTML()
#   b1 = gr.Button("Find similar words")
#   b1.click(find_most_similar, inputs=word, outputs=top_10_similar_words)
  
#   wordA = gr.inputs.Textbox("Word A")
#   wordB = gr.inputs.Textbox("Word B")
#   similarity = gr.Number()
#   b2 = gr.Button("Find similarity")
#   b2.click(similarity, inputs=[wordA, wordB], outputs=similarity)
  
# demo.launch()