Word2Vec Skip-gram

hardFREEWord EmbeddingsNLP 1: Foundations

Implement skip-gram training pair generation for Word2Vec.

In the skip-gram model, given a center word, we predict the context words within a window. Generate all (center_word, context_word) training pairs.

Write skipgram_pairs(words, window).

Input:

words: a list of words (the text already split on spaces)
window: window size W (integer)

For each word at position i, context words are at positions max(0, i-W) to min(n-1, i+W), excluding position i itself.

Output: A list of tuples (center_word, context_word) in order of center word position, and for each center word, context words from left to right.

Example

Input

I love NLP deeply
2

Output

[('I', 'love'), ('I', 'NLP'), ('love', 'I'), ('love', 'NLP'), ('love', 'deeply'), ('NLP', 'I'), ('NLP', 'love'), ('NLP', 'deeply'), ('deeply', 'love'), ('deeply', 'NLP')]

Reasoning

Inputs: words is the text split into $n = 4$ tokens: "I" (position 0), "love" (1), "NLP" (2), "deeply" (3). window is $W = 2$ neighbours on each side.
Each word's window: position $i$ sees positions $\max(0, i-2)$ through $\min(3, i+2)$ , minus $i$ itself. "I" gets 1, 2 (nothing to its left); "love" gets 0, 2, 3; "NLP" gets 0, 1, 3; "deeply" gets 1, 2 (the text ends on its right).
Pairs: each center word is paired with its context words, left to right, giving [('I', 'love'), ('I', 'NLP'), ('love', 'I'), ('love', 'NLP'), ('love', 'deeply'), ('NLP', 'I'), ('NLP', 'love'), ('NLP', 'deeply'), ('deeply', 'love'), ('deeply', 'NLP')].
Count check: $2 + 3 + 3 + 2 = 10$ pairs. Away from the edges a word would get $2W = 4$ neighbours (16 pairs in total); the boundaries clip that.
Notice: pairs come in mirrored couples, e.g. (I, NLP) and (NLP, I), because "within $W$ positions" is symmetric. Casing is preserved, so "NLP" stays upper case.

Constraints

Window size W means W words to the left and W words to the right
Handle boundary positions (beginning/end of text)
Preserve original word casing
Output pairs in positional order

Learn this in

NLP FundamentalsML Study Plan Text RepresentationNLP Study Plan Word EmbeddingsNLP Study Plan Language ModelingNLP Study Plan

More from NLP 1: Foundations

Embedding Lookupeasy Cosine Similarity Calculatoreasy Binary Vectorizereasy Document Term Matrixeasy Character Frequency Countereasy

Concept primer

Background Knowledge

The Word2Vec model is a type of word embedding technique used in Natural Language Processing (NLP). It represents words as vectors in a high-dimensional space, where semantically similar words are closer together. The skip-gram model is a variant of Word2Vec that predicts the context words given a center word. This is in contrast to the Continuous Bag of Words (CBOW) model, which predicts the center word given the context words.

In the skip-gram model, a window size $W$ is used to define the context words. For a given center word at position $i$ , the context words are those within the window of size $W$ to the left and right of $i$ , excluding $i$ itself. The goal is to generate all possible $(\text{center word}, \text{context word})$ pairs for a given text. This is a crucial step in training the Word2Vec model, as it allows us to learn the vector representations of words based on their co-occurrence patterns.

The skip-gram model is based on the idea that words that appear in similar contexts are likely to have similar meanings. By predicting the context words given a center word, we can learn to represent words as vectors that capture their semantic relationships. This has numerous applications in NLP, including text classification, sentiment analysis, and language modeling.

Algorithm/Approach

The general approach to solving this problem involves iterating over each word in the input text and generating the context words within the specified window size. This can be achieved using a simple iterative algorithm that keeps track of the current word position and the corresponding context words.

Step-by-Step Strategy

To implement the solution, follow these steps:

Read the input text and window size $W$ .
Split the input text into a list of words.
Iterate over each word in the list, keeping track of its position $i$ .
For each word at position $i$ , generate the context words by iterating over the positions from $\max(0, i-W)$ to $\min(n-1, i+W)$ , excluding $i$ itself.
Create a tuple $(\text{center word}, \text{context word})$ for each context word and add it to the result list.
Return the list of tuples.

Common Pitfalls

When implementing the solution, watch out for the following:

Make sure to exclude the center word itself when generating the context words.
Handle edge cases where the window size $W$ is larger than the distance to the beginning or end of the text.
Use efficient data structures to store the input text and result list.

Time & Space Complexity

The expected time complexity is $O(n \cdot W)$ , where $n$ is the length of the input text and $W$ is the window size. This is because we iterate over each word in the text and generate the context words within the window size. The space complexity is $O(n \cdot W)$ as well, as we need to store the result list of tuples.

solution.py

Python 3.13

Test results