word2Vec: Word to Vector with similarity

word2vec.py

!pip install gensim

import gensim.downloader as api
from gensim.models import Word2Vec


def test_word2vec(corpus, test_words, sg=1, vector_size=100, window=5):
    """Trains a Word2Vec model and checks vocab status for test words.

    Parameters:
    - corpus: List of tokenized sentences.
    - test_words: List of strings to check.
    - sg: 0 for CBOW, 1 for Skip-gram.
    """
    algo = "Skip-gram" if sg == 1 else "CBOW"
    print(f"\nTraining Word2Vec ({algo})...")

    model = Word2Vec(
        sentences=corpus,
        vector_size=vector_size,
        window=window,
        min_count=1,
        sg=sg,
    )
    wv = model.wv

    print(f"\n{'WORD / TOKEN':<18} | {'STATUS':<15} | {'TOP NEIGHBOR (most similar word)'}")
    print("-" * 72)

    for word in test_words:
        if word in wv.key_to_index:
            similar = wv.most_similar(word, topn=1)[0][0]
            print(f"{word:<18} | ✅ IN VOCAB    | '{similar}'")
        else:
            print(f"{word:<18} | ❌ OOV (KeyError) | Cannot process vector!")

# Corpus = Your Organization's Brain (Internal Data) / your personal data
# Data for Training

corpus = [
    ["the", "cat", "sat", "on", "the", "mat"],
    ["the", "quickly", "running", "cat", "saw", "a", "dog"],
    ["we", "need", "to", "tune", "the", "hyperparameter"],
    # ["the", "unhappines", "was", "evident"],
    ["tokenization", "is", "a", "key", "nlp", "step"],
]

# Test Words = Your Users' Prompts (Queries): 
# The actual questions
test_words = [
    "unhappines",
    "tokenization",
    "cat",
    "quickkkly",
    "hyperparameter",
    "skibidi",
]

# Test Skip-Gram
test_word2vec(corpus, test_words, sg=1)

# Test CBOW 
test_word2vec(corpus, test_words, sg=0)