GitHub - benjismith/word2vec-mojo

A faithful, fast port of the classic word2vec tool (Mikolov et al., Google, 2013) to the Mojo programming language — same command-line interface, same file formats, same training algorithm, rewritten for clarity and modern hardware.

Vectors trained by this port are interchangeable with vectors trained by the original C tool: each implementation's distance/readbin utilities read the other's binary output byte-for-byte.

Why Mojo

The original implementation is ~700 lines of dense C built around global variables, hand-rolled hash tables, and compiler-dependent auto-vectorization. Mojo lets this codebase say what it means:

Explicit SIMD. The two kernels all training arithmetic reduces to — dot product and scaled vector addition — are written directly against the CPU's vector registers (SIMD[DType.float32, width]), sized automatically for whatever machine compiles them (src/word2vec/simd_math.mojo).
Structured parallelism. The reference's pthread boilerplate becomes a closure handed to parallelize; the lock-free "Hogwild" weight sharing the algorithm depends on is expressed in one deliberate unsafe_origin_cast with a comment, instead of being implicit everywhere (src/word2vec/trainer.mojo).
Safety by default, escape hatches where earned. Bounds-checked collections and the compiler's exclusivity checking cover all of the code except the hot loops that opt out on purpose. Porting under those checks caught a real latent indexing hazard at the sigmoid table's boundary that the C code only avoids through an integer-division accident.
A vocabulary that reads like the algorithm. syn0 is input_embeddings, neu1e is input_gradient, vocab[word].point is ancestor_rows — every module carries the explanation of both the algorithm and its data layout.

Two implementation upgrades over the reference (both documented in-source, neither changing the learned model):

Negative sampling draws noise words with Vose's alias method — exact O(1) sampling from a few megabytes of tables — replacing the C tool's quantized 100-million-entry array (400 MB of RAM) (src/word2vec/noise_sampler.mojo).
The corpus tokenizer streams the training file in 1 MiB slabs instead of one fgetc() per byte (src/word2vec/byte_slabs.mojo).

Performance

Benchmarked against the reference C implementation on an Apple M2 Pro (8 performance + 4 efficiency cores), training a 192 MB / 35.7M-word corpus with production-typical settings (skip-gram, -negative 15 -sample 1e-4 -size 100), both tools loading an identical shared vocabulary:

~3.3x the single-thread throughput of the C tool (390K vs. 120K words/sec) — the SIMD kernels, fused update pass, and streaming tokenizer at work.
Saturates the machine's memory bandwidth at 8 threads (~1.6M words/sec) and holds that plateau flat through 16 threads; the C tool peaks at 0.90M words/sec at 16 threads, still short of that ceiling. Net: about 1.8x end-to-end at high thread counts, approaching 3x at low ones.
Work is claimed by threads dynamically, so on heterogeneous CPUs (performance + efficiency cores) every thread count from 8 up performs identically here — efficiency cores help rather than gate. -threads 8 is the energy-efficient choice on this machine.

Full tables, the optimization history (including one measured-and-rejected idea), and measurement methodology: docs/benchmarks.md. Rerun with scripts/benchmark.sh.

Related work

As of August 2026 this appears to be the only word2vec implementation written in Mojo (a GitHub search for word2vec repositories in the language returns exactly this repository). The nearest neighbor is mojo-gensim, which pursues a different goal: a gensim-workalike library (negative sampling only, no hierarchical softmax, no compatibility with the C tools' CLI or file formats) benchmarked against gensim rather than the reference.

Among faithful rewrites — implementations keeping the original's exact SGD semantics — published comparisons cluster around parity to ~1.6x the C tool's throughput: gensim's Cython core is roughly at parity (with known thread-scaling limits), word2vec-rs (Rust) reports ~1.6x at 8 threads, and finalfrontier (Rust, Hogwild) publishes no comparative numbers. This port measures 3.27x single-threaded and ~1.8x at high thread counts against the C tool on the same machine and corpus (see above) — with the caveat that cross-paper words-per-second comparisons are unreliable (different hardware, corpora, dimensions, and negative-sample counts); the only numbers here measured under controlled conditions are our own.

Substantially faster implementations exist, but all of them change the training algorithm: pWord2Vec / HogBatch (Intel) batches context windows and shares negative samples across them so updates become matrix multiplies (~3.5x the original on the same hardware, further with many-core and distributed scaling), and BlazingText (AWS) applies the same reformulation on GPUs. Those are the right choices for training-throughput-at-any-cost; this project deliberately stays in the faithful tier, where vectors, file formats, and update semantics remain interchangeable with the original tools.

Getting started

The Mojo toolchain is pinned by pixi (brew install pixi), so setup is:

pixi install        # fetches the pinned Mojo toolchain into .pixi/
pixi run build      # compiles bin/word2vec, bin/distance, bin/readbin
pixi run test       # compiles and runs the test suites in tests/

Usage

Train vectors (flags and defaults match the original tool; add -iter N for multiple training passes):

./bin/word2vec -train corpus.txt -output vectors.bin \
    -cbow 0 -size 100 -window 5 -min-count 15 \
    -negative 15 -hs 0 -sample 1e-4 -threads 16 -binary 1

Explore nearest neighbors interactively:

./bin/distance vectors.bin

Dump a binary vector file as text (vectors unit-normalized, like the original readbin):

./bin/readbin vectors.bin > vectors.txt

Ask analogy questions interactively (man king woman -> queen):

./bin/word-analogy vectors.bin

Score vectors against an analogy benchmark such as the classic questions-words.txt (an optional threshold restricts evaluation to the N most frequent words):

./bin/compute-accuracy vectors.bin 30000 < questions-words.txt

Join frequent collocations into single tokens before training ("new york" becomes "new_york"); chain runs with descending thresholds for longer phrases:

./bin/word2phrase -train corpus.txt -output phrased.txt -threshold 100

The training corpus format is whitespace-separated tokens; newlines mark sentence boundaries. All of -cbow {0,1} x -hs {0,1} / -negative N combinations from the reference are implemented, as are -save-vocab / -read-vocab and the k-means -classes output mode.

Custom tokenizers

Tokenization is pluggable. word2vec itself only consumes a stream of tokens and sentence boundaries; everything about how corpus bytes become tokens lives behind the Tokenizer trait (src/word2vec/tokenizer.mojo), and the classic whitespace convention is just the default implementation. Three ways to bring your own:

Byte-pair encoding, no code required. Train with your own BPE tokenizer by supplying its ranked merge rules:
```
./bin/word2vec -train corpus.txt -output vectors.bin \
    -tokenizer bpe -bpe-merges merges.txt \
    -negative 15 -sample 1e-4 -threads 8 -binary 1
```
The merges file holds one rule per line ("left right", most important first, \xNN escapes for bytes like space; # comments allowed). The token set is closed: all 256 single bytes plus one token per merge, and every token receives an embedding — including tokens absent from the corpus — with -min-count not applying. Corpus counts still drive subsampling, negative sampling, and Huffman codes. An empty merges file gives a pure character model, useful for scripts that whitespace tokenization serves poorly. (-bpe-newline-boundary 0 disables newline sentence breaks.)
Implement the trait. A tokenizer is one Mojo struct with seven methods (stream, reset, resync, duplicate-per-worker, and an open/closed vocabulary declaration); the vocabulary builder and trainer are generic over it and monomorphize at compile time, so a custom tokenizer pays no dispatch cost in the hot loop.
Pre-tokenize externally. Any tokenizer in any language (HuggingFace, SentencePiece, ...) can emit a corpus of space-separated tokens, which the default whitespace tokenizer then trains on verbatim — the same composition trick word2phrase has always used.

Token text in vocabulary and vector files is \xNN-escaped for the five format-breaking bytes (space, tab, CR, LF, backslash) and unescaped on load. For whitespace-tokenized words this is the identity, so existing files — including ones written by the original C tools — are unaffected.

The full guide — merges-file format, closed-vocabulary semantics, how the query tools handle BPE vector files, importing tokenizers from GPT-2 / HuggingFace (via scripts/convert-gpt2-merges.py) and subword-nmt, and the trait contract for implementing your own — is in docs/tokenizers.md.

Layout

src/
  train.mojo               the `word2vec` training command
  distance.mojo            interactive nearest-neighbor queries
  word_analogy.mojo        interactive analogy queries (a : b :: c : ?)
  compute_accuracy.mojo    batch analogy-benchmark scoring
  word2phrase.mojo         collocation-joining corpus preprocessor
  readbin.mojo             binary -> text vector dump
  word2vec/                the library the commands are built from
    tokenizer.mojo         the Tokenizer trait (pluggable tokenization)
    byte_slabs.mojo        buffered corpus byte reading
    whitespace_tokenizer.mojo  the classic word2vec tokenization (default)
    bpe_tokenizer.mojo     byte-pair encoding from a merges file
    token_text_escaping.mojo   \xNN escaping for token text in files
    vocabulary.mojo        token counting, pruning, and indexing
    huffman_coding.mojo    Huffman codes for hierarchical softmax
    noise_sampler.mojo     alias-method negative-sampling distribution
    sigmoid_table.mojo     precomputed logistic function
    pseudo_random.mojo     the reference's tiny thread-local LCG
    simd_math.mojo         dot-product / axpy / fused-update SIMD kernels
    trainer.mojo           the multi-threaded training loop
    embedding_matrix.mojo  row storage for the weight matrices
    word_vectors.mojo      vector file I/O + similarity queries
    analogy.mojo           analogy arithmetic + benchmark evaluation
    phrase_builder.mojo    collocation statistics and joining
    kmeans_clustering.mojo word clustering (-classes mode)
    command_line.mojo      shared -flag value parsing
    standard_input.mojo    pipe-safe line reading for the query tools
tests/                     one suite per module + end-to-end training checks
scripts/                   build.sh / run-tests.sh / benchmark.sh

Fidelity notes

The port reproduces the reference algorithm exactly: same subsampling formula, same window narrowing, same gradient math, same sigmoid quantization (including the C expression's integer-division quirk), same Hogwild lock-free threading, same output formats. Known deliberate deviations, each documented where it lives in the source:

Vocabulary ties (equal counts) break alphabetically instead of by qsort's unspecified order, so vocabularies are fully deterministic.
Negative samples come from the alias method (exact probabilities) rather than the quantized table; the distribution is the same unigram^0.75.
Work is distributed to threads dynamically (each worker atomically claims the next small corpus chunk) instead of the reference's one-static-range- per-thread split with per-thread word budgets. Same total work and the same per-word math; on CPUs with mixed performance/efficiency cores it removes the straggler effect that made some thread counts slower than fewer threads. Single-threaded training remains fully deterministic.
The final token of a file with no trailing newline is counted (the C tool drops it), and over-long tokens keep their first 99 bytes rather than the C buffer-juggling's first-98-plus-last byte.
Training is bit-for-bit reproducible only single-threaded — true of the original as well (Hogwild races are inherently nondeterministic).
word2phrase reproduces the reference's scoring and output byte for byte (verified by diff on a multi-megabyte corpus), including counting word pairs across line boundaries while never joining across them; it keeps full-length tokens and the file's final unterminated token where the C tool truncates at 59 bytes and drops that last token.
compute-accuracy reproduces the reference's per-section and total accuracy numbers exactly (verified on a 486K-word vector file); output formatting differs trivially (e.g. "0.0" where C prints "nan" for empty categories).