main.py

from src.data_generator import SkipGramDataset
from src.evaluation import analogy, most_similar, plot_embeddings
from src.model import SkipGramModel
from src.preprocessing import Vocabulary, clean_and_tokenise
from src.trainer import Trainer


QUERY_WORDS = ["king", "war", "city", "world", "state", "music", "water"]

ANALOGY_TESTS = [
    ("king", "queen", "man"),
    ("france", "paris", "germany"),
    ("war", "peace", "conflict"),
]

PLOT_WORDS = [
    "king", "queen", "prince", "man", "woman", "boy", "girl",
    "city", "town", "village", "country", "state", "nation",
    "war", "battle", "peace", "army", "military",
    "water", "river", "sea", "ocean", "lake",
    "music", "film", "art", "play", "game",
    "world", "life", "time", "year", "day",
    "school", "university", "student", "team", "group",
]


def main():
    tokens = clean_and_tokenise(max_tokens=500_000)
    print(f"Tokens: {len(tokens):,}")

    vocab = Vocabulary(min_freq=5)
    indices = vocab.build(tokens)
    print(f"Vocabulary: {vocab.size:,} words")

    indices = vocab.subsample(indices)
    print(f"After subsampling: {len(indices):,} tokens")

    dataset = SkipGramDataset(indices, vocab.freqs)
    print(f"Skip-gram pairs: {len(dataset.centres):,}")
    print(f"Batches per epoch: {len(dataset):,}\n")

    model = SkipGramModel(vocab.size, embedding_dim=100)
    trainer = Trainer(model, dataset, lr=0.025, epochs=10)
    trainer.train()

    embeddings = model.W_in

    print("\n--- Most Similar Words ---")
    for word in QUERY_WORDS:
        results = most_similar(word, vocab.word2idx, vocab.idx2word, embeddings, top_n=5)
        if results:
            neighbours = ", ".join(f"{w} ({s:.3f})" for w, s in results)
            print(f"  {word}: {neighbours}")

    print("\n--- Analogy Tests (a:b :: c:?) ---")
    for a, b, c in ANALOGY_TESTS:
        results = analogy(a, b, c, vocab.word2idx, vocab.idx2word, embeddings, top_n=3)
        if results:
            answers = ", ".join(f"{w} ({s:.3f})" for w, s in results)
            print(f"  {a}:{b} :: {c}:? → {answers}")

    plot_embeddings(PLOT_WORDS, vocab.word2idx, embeddings, save_path="embeddings.png")


if __name__ == "__main__":
    main()