Show HN: Word2vec Algorithm in ~100sloc with NumPy (github.com)

I have a suggestion for improving clarity:

    word_bag: dict[str, int] = dict()  # Multiset

    for line in raw_corpus:
        words = line.split()
        
        for word in words:
            if word in word_bag:
                word_bag[word] += 1
            else:
                word_bag[word] = 1

Try:

    word_bag = collections.Counter()
    for line in raw_corpus:
        word_bag.update(line.split())

then, to filter:

    word_bag = {word: count for (word, count) in word_bag.items()
         if count >= min_frequency}

or, if you can avoid printing the filtered word bag:

    return set(word for (word, count) in word_bag.items()
        if count >= min_frequency)

Also, that:

    raw_corpus    : list[list[str]],

should be list[str]. It's your corpus which is a list[list[str]]. (And you might want to use "line_words" instead of "line"

Could you explain why lines are processed independently?

    for i, line in enumerate(corpus):
        for j, target_word in enumerate(line):

Does it make a difference that some lines have fewer words than others?

I would think that

  for positive_words in sliding_window(flatten(corpus), window_size):

using the sliding_window and flatting recipes at https://docs.python.org/3/library/itertools.html#itertools-r... would be more appropriate. Though using it directly would not handle the leading/trailing partial windows.

I liked seeing the 'η' - I haven't yet used non-ASCII in my code!

Ignorant question about the algorithm:

There is a min_frequency threshold for including words

Wouldn’t you want a max_frequency instead?

A lot of times infrequent words are more important than the frequent ones (eg titles vs articles)

What Ive wondered about word-vectors is, are different meanings split or not?

A word having many meanings would otherwise means its vectors are more "smeared" out and less accurate I guess. Although, I guess that might be good for these algorithms.

Neat! It /should/ be pretty easy to convert this to Jax, as a big party of the Jax API is just reimplementing the Numpy API.

The two changes I can imagine are: a) handling randomness requires more explicit rng estate handling, and b) wrapping the inner loop in a function so that it can be JIT-compiled.

The advantage to doing this is that you can then run your code on a GPU, getting you most of the way to 'blazingly fast.'

Right out of the bag lots of wasted lines

    word_bag: dict[str, int] = dict()  # Multiset

    for line in raw_corpus:
        words = line.split()
        
        for word in words:
            if word in word_bag:
                word_bag[word] += 1
            else:
                word_bag[word] = 1

    keys_to_drop = []
    for k, v in word_bag.items():
        if v < min_frequency:
            keys_to_drop.append(k)

    for k in keys_to_drop:
        del word_bag[k]

    pprint(word_bag)
    print(len(word_bag))

    vocabulary = word_bag.keys()
    return set(vocabulary)

Can be a python one/two liner

    # word_bag = Counter()  # defaultdict(int)
    word_bag = Counter(word for line in raw_corpus for word in line.split())
    return set(word for word, count in word_bag.items() if count >= min_frequency)

The golden standard for word2vec and topic modeling in general is gensim: https://radimrehurek.com/gensim/models/word2vec.html

I've been using gensim for a few years for topic modeling. It works great with pyLDAvis for visualizations.

How does your implementation differ from the one by Radim?

Beginner here. Having done most of the Karpathy YT course I have seen embeddings trained as part of a larger net that is predicting the next token. After the embedding is run on several tokens the output goes though several layers of linear;batchnorm;nonlinear. I haven’t done the wavenet or transformer stuff yet.

How can you train the embedding without that? And if you can, is it a good or bad idea to prepare your embedding using this script first then chuck those weights into a bigger model (I think I know what the answer will be…)