org.apache.spark.mllib.feature

Class Word2Vec

java.lang.Object

org.apache.spark.mllib.feature.Word2Vec

All Implemented Interfaces:

java.io.Serializable, Logging

public class Word2Vec
extends java.lang.Object
implements scala.Serializable, Logging

:: Experimental :: 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.

See Also:

Serialized Form

Constructor Summary

Constructors

Constructor and Description
Word2Vec()

Method Summary

Methods

Modifier and Type	Method and Description
<S extends java.lang.Iterable<java.lang.String>> Word2VecModel	fit(JavaRDD<S> dataset) Computes the vector representation of each word in vocabulary (Java version).
<S extends scala.collection.Iterable<java.lang.String>> Word2VecModel	fit(RDD<S> dataset) Computes the vector representation of each word in vocabulary.
Word2Vec	setLearningRate(double learningRate) Sets initial learning rate (default: 0.025).
Word2Vec	setMinCount(int minCount) Sets minCount, the minimum number of times a token must appear to be included in the word2vec model's vocabulary (default: 5).
Word2Vec	setNumIterations(int numIterations) Sets number of iterations (default: 1), which should be smaller than or equal to number of partitions.
Word2Vec	setNumPartitions(int numPartitions) Sets number of partitions (default: 1).
Word2Vec	setSeed(long seed) Sets random seed (default: a random long integer).
Word2Vec	setVectorSize(int vectorSize) Sets vector size (default: 100).

Methods inherited from class java.lang.Object

clone, equals, finalize, getClass, hashCode, notify, notifyAll, toString, wait, wait, wait

Methods inherited from interface org.apache.spark.Logging

initializeIfNecessary, initializeLogging, isTraceEnabled, log_, log, logDebug, logDebug, logError, logError, logInfo, logInfo, logName, logTrace, logTrace, logWarning, logWarning

Constructor Detail

Word2Vec

public Word2Vec()

Method Detail

setVectorSize

public Word2Vec setVectorSize(int vectorSize)

Sets vector size (default: 100).

Parameters:

vectorSize - (undocumented)

Returns:

(undocumented)

setLearningRate

public Word2Vec setLearningRate(double learningRate)

Sets initial learning rate (default: 0.025).

Parameters:

learningRate - (undocumented)

Returns:

(undocumented)

setNumPartitions

public Word2Vec setNumPartitions(int numPartitions)

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

Parameters:

numPartitions - (undocumented)

Returns:

(undocumented)

setNumIterations

public Word2Vec setNumIterations(int numIterations)

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

Parameters:

numIterations - (undocumented)

Returns:

(undocumented)

setSeed

public Word2Vec setSeed(long seed)

Sets random seed (default: a random long integer).

Parameters:

seed - (undocumented)

Returns:

(undocumented)

setMinCount

public Word2Vec setMinCount(int minCount)

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

Parameters:

minCount - (undocumented)

Returns:

(undocumented)

fit

public <S extends scala.collection.Iterable<java.lang.String>> Word2VecModel fit(RDD<S> dataset)

Computes the vector representation of each word in vocabulary.

Parameters:

dataset - an RDD of words

Returns:

a Word2VecModel

fit

public <S extends java.lang.Iterable<java.lang.String>> Word2VecModel fit(JavaRDD<S> dataset)

Computes the vector representation of each word in vocabulary (Java version).

Parameters:

dataset - a JavaRDD of words

Returns:

a Word2VecModel