import gradio as gr
from gensim.models import Word2Vec

# Load the pre-trained Word2Vec model
model = Word2Vec.load("burmese_word2vec_model.bin")

# Function to calculate the similarity between two words
def get_similarity(word1, word2):
    try:
        similarity = model.wv.similarity(word1, word2)
        return f"Similarity between '{word1}' and '{word2}': {similarity:.4f}"
    except KeyError as e:
        return f"Error: One or both words not in vocabulary - {str(e)}"

# Function to find similar words to a given word
def find_similar_words(word):
    try:
        similar_words = model.wv.most_similar(word, topn=10)
        similar_words_str = "\n".join([f"{w[0]}: {w[1]:.4f}" for w in similar_words])
        return f"Words similar to '{word}':\n{similar_words_str}"
    except KeyError as e:
        return f"Error: Word not in vocabulary - {str(e)}"

# Create the Gradio interface
def main_function(word1, word2, similar_word):
    similarity_result = get_similarity(word1, word2)
    similar_words_result = find_similar_words(similar_word)
    return similarity_result, similar_words_result

interface = gr.Interface(
    fn=main_function,
    inputs=[
        gr.Textbox(label="စကားလုံး (၁)"),
        gr.Textbox(label="စကားလုံး (၂)"),
        gr.Textbox(label="ပုံစံတူစကားလုံးများကို ရှာပါ။")
    ],
    outputs=[
        gr.Textbox(label="စကားလုံး တူညီမှု"),
        gr.Textbox(label="အလားတူစကားလုံးများ")  
    ],
    title="မြန်မာ Word2Vec မော်ဒယ်",
    description="စကားလုံးနှစ်လုံး ဆက်စပ်မှုကို တွက်ချက်ရန် 'စကားလုံး(၁)' နှင့် 'စကားလုံး(၂)'နေရာတွင် နှစ်သက်ရာမြန်မာစကားလုံးရိုက်ထည့်ပြီး ပုံစံတူစကားလုံများရှာရန် 'ပုံစံတူစကားလုံးများကိုရှာပါ'နေရာတွင် နှစ်သက်ရာစကားလုံးရိုက်ထည့်ပါ။"
)

# Launch the app
if __name__ == "__main__":
    interface.launch()