Source code for sparknlp.annotator.embeddings.word2vec

#  Copyright 2017-2022 John Snow Labs
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
"""Contains classes for Word2Vec."""


from sparknlp.common import *


[docs]class Word2VecApproach(AnnotatorApproach, HasStorageRef, HasEnableCachingProperties): """Trains a Word2Vec model that creates vector representations 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 use Word2Vec implemented in Spark ML. It uses skip-gram model in our implementation and a hierarchical softmax method to train the model. The variable names in the implementation match the original C implementation. For instantiated/pretrained models, see :class:`.Word2VecModel`. For available pretrained models please see the `Models Hub <https://sparknlp.org/models>`__. ====================== ======================= Input Annotation types Output Annotation type ====================== ======================= ``TOKEN`` ``WORD_EMBEDDINGS`` ====================== ======================= Parameters ---------- vectorSize The dimension of codes after transforming from words (> 0), by default 100 windowSize The window size (context words from [-window, window]) (> 0), by default 5 numPartitions Number of partitions for sentences of words (> 0), by default 1 minCount The minimum number of times a token must appear to be included in the word2vec model's vocabulary (>= 0), by default 1 maxSentenceLength The window size (Maximum length (in words) of each sentence in the input data. Any sentence longer than this threshold will be divided into chunks up to the size (> 0), by default 1000 stepSize Step size (learning rate) to be used for each iteration of optimization (> 0), by default 0.025 maxIter Maximum number of iterations (>= 0), by default 1 seed Random seed, by default 44 References ---------- For the original C implementation, see https://code.google.com/p/word2vec/ For the research paper, see `Efficient Estimation of Word Representations in Vector Space <https://arxiv.org/abs/1301.3781>`__ and `Distributed Representations of Words and Phrases and their Compositionality <https://arxiv.org/pdf/1310.4546v1.pdf>`__. Examples -------- >>> import sparknlp >>> from sparknlp.base import * >>> from sparknlp.annotator import * >>> from pyspark.ml import Pipeline >>> documentAssembler = DocumentAssembler() \\ ... .setInputCol("text") \\ ... .setOutputCol("document") >>> tokenizer = Tokenizer() \\ ... .setInputCols(["document"]) \\ ... .setOutputCol("token") >>> embeddings = Word2VecApproach() \\ ... .setInputCols(["token"]) \\ ... .setOutputCol("embeddings") >>> pipeline = Pipeline() \\ ... .setStages([ ... documentAssembler, ... tokenizer, ... embeddings ... ]) >>> path = "sherlockholmes.txt" >>> dataset = spark.read.text(path).toDF("text") >>> pipelineModel = pipeline.fit(dataset) """ inputAnnotatorTypes = [AnnotatorType.TOKEN] outputAnnotatorType = AnnotatorType.WORD_EMBEDDINGS vectorSize = Param(Params._dummy(), "vectorSize", "the dimension of codes after transforming from words (> 0)", typeConverter=TypeConverters.toInt) windowSize = Param(Params._dummy(), "windowSize", "the window size (context words from [-window, window]) (> 0)", typeConverter=TypeConverters.toInt) numPartitions = Param(Params._dummy(), "numPartitions", "number of partitions for sentences of words (> 0)", typeConverter=TypeConverters.toInt) minCount = Param(Params._dummy(), "minCount", "the minimum number of times a token must " + "appear to be included in the word2vec model's vocabulary (>= 0)", typeConverter=TypeConverters.toInt) maxSentenceLength = Param(Params._dummy(), "maxSentenceLength", "the window size (Maximum length (in words) of each sentence in the input data. Any sentence longer than this threshold will " + "be divided into chunks up to the size (> 0)", typeConverter=TypeConverters.toInt) stepSize = Param(Params._dummy(), "stepSize", "Step size (learning rate) to be used for each iteration of optimization (> 0)", typeConverter=TypeConverters.toFloat) maxIter = Param(Params._dummy(), "maxIter", "maximum number of iterations (>= 0)", typeConverter=TypeConverters.toInt) seed = Param(Params._dummy(), "seed", "Random seed", typeConverter=TypeConverters.toInt)
[docs] def setVectorSize(self, vectorSize): """ Sets vector size (default: 100). """ return self._set(vectorSize=vectorSize)
[docs] def setWindowSize(self, windowSize): """ Sets window size (default: 5). """ return self._set(windowSize=windowSize)
[docs] def setStepSize(self, stepSize): """ Sets initial learning rate (default: 0.025). """ return self._set(stepSize=stepSize)
[docs] def setNumPartitions(self, numPartitions): """ Sets number of partitions (default: 1). Use a small number for accuracy. """ return self._set(numPartitions=numPartitions)
[docs] def setMaxIter(self, numIterations): """ Sets number of iterations (default: 1), which should be smaller than or equal to number of partitions. """ return self._set(maxIter=numIterations)
[docs] def setSeed(self, seed): """ Sets random seed. """ return self._set(seed=seed)
[docs] def setMinCount(self, minCount): """ Sets minCount, the minimum number of times a token must appear to be included in the word2vec model's vocabulary (default: 5). """ return self._set(minCount=minCount)
[docs] def setMaxSentenceLength(self, maxSentenceLength): """ Maximum length (in words) of each sentence in the input data. Any sentence longer than this threshold will be divided into chunks up to the size (> 0) """ return self._set(maxSentenceLength=maxSentenceLength)
@keyword_only def __init__(self): super(Word2VecApproach, self).__init__(classname="com.johnsnowlabs.nlp.embeddings.Word2VecApproach") self._setDefault( vectorSize=100, windowSize=5, numPartitions=1, minCount=1, maxSentenceLength=1000, stepSize=0.025, maxIter=1, seed=44 ) def _create_model(self, java_model): return Word2VecModel(java_model=java_model)
[docs]class Word2VecModel(AnnotatorModel, HasStorageRef, HasEmbeddingsProperties): """Word2Vec model that creates vector representations 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 use Word2Vec implemented in Spark ML. It uses skip-gram model in our implementation and a hierarchical softmax method to train the model. The variable names in the implementation match the original C implementation. This is the instantiated model of the :class:`.Word2VecApproach`. For training your own model, please see the documentation of that class. Pretrained models can be loaded with :meth:`.pretrained` of the companion object: >>> embeddings = Word2VecModel.pretrained() \\ ... .setInputCols(["token"]) \\ ... .setOutputCol("embeddings") The default model is `"word2vec_gigaword_300"`, if no name is provided. ====================== ======================= Input Annotation types Output Annotation type ====================== ======================= ``TOKEN`` ``WORD_EMBEDDINGS`` ====================== ======================= Parameters ---------- vectorSize The dimension of codes after transforming from words (> 0), by default 100 References ---------- For the original C implementation, see https://code.google.com/p/word2vec/ For the research paper, see `Efficient Estimation of Word Representations in Vector Space <https://arxiv.org/abs/1301.3781>`__ and `Distributed Representations of Words and Phrases and their Compositionality <https://arxiv.org/pdf/1310.4546v1.pdf>`__. Examples -------- >>> import sparknlp >>> from sparknlp.base import * >>> from sparknlp.annotator import * >>> from pyspark.ml import Pipeline >>> documentAssembler = DocumentAssembler() \\ ... .setInputCol("text") \\ ... .setOutputCol("document") >>> tokenizer = Tokenizer() \\ ... .setInputCols(["document"]) \\ ... .setOutputCol("token") >>> embeddings = Word2VecModel.pretrained() \\ ... .setInputCols(["token"]) \\ ... .setOutputCol("embeddings") >>> embeddingsFinisher = EmbeddingsFinisher() \\ ... .setInputCols(["embeddings"]) \\ ... .setOutputCols("finished_embeddings") \\ ... .setOutputAsVector(True) >>> pipeline = Pipeline().setStages([ ... documentAssembler, ... tokenizer, ... embeddings, ... embeddingsFinisher ... ]) >>> data = spark.createDataFrame([["This is a sentence."]]).toDF("text") >>> result = pipeline.fit(data).transform(data) >>> result.selectExpr("explode(finished_embeddings) as result").show(1, 80) +--------------------------------------------------------------------------------+ | result| +--------------------------------------------------------------------------------+ |[0.06222493574023247,0.011579325422644615,0.009919632226228714,0.109361454844...| +--------------------------------------------------------------------------------+ """ name = "Word2VecModel" inputAnnotatorTypes = [AnnotatorType.TOKEN] outputAnnotatorType = AnnotatorType.WORD_EMBEDDINGS vectorSize = Param(Params._dummy(), "vectorSize", "the dimension of codes after transforming from words (> 0)", typeConverter=TypeConverters.toInt)
[docs] def setVectorSize(self, vectorSize): """ Sets vector size (default: 100). """ return self._set(vectorSize=vectorSize)
def __init__(self, classname="com.johnsnowlabs.nlp.embeddings.Word2VecModel", java_model=None): super(Word2VecModel, self).__init__( classname=classname, java_model=java_model ) self._setDefault( vectorSize=100 ) @staticmethod
[docs] def pretrained(name="word2vec_gigaword_300", lang="en", remote_loc=None): """Downloads and loads a pretrained model. Parameters ---------- name : str, optional Name of the pretrained model, by default "word2vec_wiki" lang : str, optional Language of the pretrained model, by default "en" remote_loc : str, optional Optional remote address of the resource, by default None. Will use Spark NLPs repositories otherwise. Returns ------- Word2VecModel The restored model """ from sparknlp.pretrained import ResourceDownloader return ResourceDownloader.downloadModel(Word2VecModel, name, lang, remote_loc)
[docs] def getVectors(self): """ Returns the vector representation of the words as a dataframe with two fields, word and vector. """ return self._call_java("getVectors")