Mini Word2Vec Engine

A word embedding engine written in plain Java, no machine learning libraries involved. It implements skip-gram with negative sampling (the training method behind the original word2vec paper) completely from scratch: tokenizing, building a vocabulary, generating training pairs, training the vectors, and then letting you explore the results through a small command line tool.

I built this mainly to actually understand how word embeddings work under the hood, instead of just calling a library function. Everything from the subsampling of frequent words to the negative sampling table is implemented by hand.

What it does

1. Reads a text file and tokenizes it (lowercase, strip punctuation, split into words).
2. Builds a vocabulary, dropping rare words that don't appear often enough to be useful.
3. Generates target/context word pairs using a sliding window, with frequent words randomly subsampled so they don't dominate training.
4. Trains a skip-gram model with negative sampling: each target word gets pushed toward the vectors of words it actually appears near, and away from a handful of random "negative" words.
5. Drops you into an interactive prompt where you can query the trained vectors.

Interactive commands

Once training finishes you get a small REPL:

similar <word>                  Find the top 5 words closest to <word>
analogy <wordA> <wordB> <wordC> Solve vector(A) - vector(B) + vector(C)
save [filename.vec]             Export the current vectors to disk
load <filename.vec]             Load a previously saved model
exit                            Quit

The save/load commands use a plain text .vec format (one word per line, followed by its vector values), so a trained model doesn't have to be regenerated every time you want to experiment with it.

Getting started

You'll need Java 17 or newer. The project uses the Gradle wrapper, so you don't need Gradle installed separately.

Clone the repo and run:

./gradlew run

(on Windows, gradlew.bat run)

Or build a jar and run it directly:

./gradlew build
java -jar build/libs/mini-word2vec-1.0-SNAPSHOT.jar

By default it trains on the sample corpus at data/corpus.txt. To point it at a different file when running through Gradle, pass it as a program argument:

./gradlew run --args="path/to/your/text.txt"

or, when running the built jar directly:

java -jar build/libs/mini-word2vec-1.0-SNAPSHOT.jar path/to/your/text.txt

Training on the included sample corpus (a Sherlock Holmes short story, a few thousand words) takes well under a minute on a normal laptop. Bigger corpora will obviously take longer, since none of this is parallelized.

A note on the sample corpus

data/corpus.txt is "The Return of Sherlock Holmes" by Sir Arthur Conan Doyle, which is in the public domain. It's just there so the project has something to train on out of the box. Swap it for whatever text you want to experiment with.

Tuning

The hyperparameters currently live as constants at the top of Main.java:

• windowSize – how many words on each side count as "context"
• vectorDimensions – size of each word vector
• epochs – how many passes over the training pairs
• learningRate – starting learning rate (it decays linearly over training)
• subsamplingThreshold – how aggressively common words get downsampled
• negativeSamples – how many negative examples per positive pair

Small corpora need lower vocab/dimension settings than what you'd use for a large, real-world dataset. With only a few thousand words of training text you won't get anywhere near the quality of a model trained on a large corpus, but the mechanics behind it are the same.

Project structure

src/main/java/com/engine/
  Main.java                       Entry point and the interactive REPL
  model/SkipGramModel.java        Training loop and vector updates
  preprocessing/Tokenizer.java    Text cleaning and pair generation
  preprocessing/Vocabulary.java   Word <-> id mapping, frequency counts
  util/VectorUtils.java           Similarity search, analogies, save/load
data/corpus.txt                   Sample training text

Known limitations

• Single-threaded, so training time scales directly with corpus size.
• No proper CBOW mode, this only implements the skip-gram side of word2vec.
• The tokenizer is intentionally simple and strips all punctuation, it won't handle things like contractions gracefully.

Why I made this

This is my first project up on GitHub. I wanted to actually learn how word embeddings are trained rather than just importing a library, so I rebuilt the core word2vec algorithm from the ground up in Java. Feedback and suggestions are welcome.