from pathlib import Path

from model import SkipGramModel, inverse_vocabulary

MODEL_PATH = Path(__file__).parent / "amharic_skipgram.json"


def main():
    model, word_to_id = SkipGramModel.load_json(MODEL_PATH)
    id_to_word = inverse_vocabulary(word_to_id)

    print("Amharic Skip-gram Word Embeddings")
    print("Type an Amharic word, or type 'exit' to quit.")

    while True:
        word = input("\nWord: ").strip()

        if word.lower() == "exit":
            break

        if word not in word_to_id:
            print("Unknown word: not in the model vocabulary.")
            continue

        word_id = word_to_id[word]
        neighbors = model.nearest_neighbors(word_id, top_k=3)

        print(f"Nearest neighbors for {word}:")
        for neighbor_id, similarity in neighbors:
            print(f"  {id_to_word[neighbor_id]}: {similarity:.6f}")


if __name__ == "__main__":
    main()