Word2Vec

class pyspark.mllib.feature.Word2Vec[source]

Word2Vec creates vector representation of words in a text corpus. The algorithm first constructs a vocabulary from the corpus and then learns vector representation of words in the vocabulary. The vector representation can be used as features in natural language processing and machine learning algorithms.

We used skip-gram model in our implementation and hierarchical softmax method to train the model. The variable names in the implementation matches the original C implementation.

For original C implementation, see https://code.google.com/p/word2vec/ For research papers, see Efficient Estimation of Word Representations in Vector Space and Distributed Representations of Words and Phrases and their Compositionality.

>>> sentence = "a b " * 100 + "a c " * 10
>>> localDoc = [sentence, sentence]
>>> doc = sc.parallelize(localDoc).map(lambda line: line.split(" "))
>>> model = Word2Vec().setVectorSize(10).setSeed(42).fit(doc)

Querying for synonyms of a word will not return that word:

>>> syms = model.findSynonyms("a", 2)
>>> [s[0] for s in syms]
['b', 'c']

But querying for synonyms of a vector may return the word whose representation is that vector:

>>> vec = model.transform("a")
>>> syms = model.findSynonyms(vec, 2)
>>> [s[0] for s in syms]
['a', 'b']
>>> import os, tempfile
>>> path = tempfile.mkdtemp()
>>> model.save(sc, path)
>>> sameModel = Word2VecModel.load(sc, path)
>>> model.transform("a") == sameModel.transform("a")
True
>>> syms = sameModel.findSynonyms("a", 2)
>>> [s[0] for s in syms]
['b', 'c']
>>> from shutil import rmtree
>>> try:
...     rmtree(path)
... except OSError:
...     pass

New in version 1.2.0.

Methods

Methods Documentation

fit(data)[source]

Computes the vector representation of each word in vocabulary.

Parameters

data – training data. RDD of list of string

Returns

Word2VecModel instance

New in version 1.2.0.

setLearningRate(learningRate)[source]

Sets initial learning rate (default: 0.025).

New in version 1.2.0.

setMinCount(minCount)[source]

Sets minCount, the minimum number of times a token must appear to be included in the word2vec model’s vocabulary (default: 5).

New in version 1.4.0.

setNumIterations(numIterations)[source]

Sets number of iterations (default: 1), which should be smaller than or equal to number of partitions.

New in version 1.2.0.

setNumPartitions(numPartitions)[source]

Sets number of partitions (default: 1). Use a small number for accuracy.

New in version 1.2.0.

setSeed(seed)[source]

Sets random seed.

New in version 1.2.0.

setVectorSize(vectorSize)[source]

Sets vector size (default: 100).

New in version 1.2.0.

setWindowSize(windowSize)[source]

Sets window size (default: 5).

New in version 2.0.0.