scala - 在 Spark 中加载 Word2Vec 模型

Question

是否可以将预训练(二进制)模型加载到 Spark(使用 scala)？我尝试加载由谷歌生成的二进制模型之一，如下所示:

    import org.apache.spark.mllib.feature.{Word2Vec, Word2VecModel}


    val model = Word2VecModel.load(sc, "GoogleNews-vectors-negative300.bin")

但无法找到元数据目录。我还创建了该文件夹并在其中附加了二进制文件，但无法解析它。我没有找到这个问题的任何包装。

Answer 1

我编写了一个快速函数，将 google news 预训练模型加载到 Spark word2vec 模型中。享受吧。

def loadBin(file: String) = {
  def readUntil(inputStream: DataInputStream, term: Char, maxLength: Int = 1024 * 8): String = {
    var char: Char = inputStream.readByte().toChar
    val str = new StringBuilder
    while (!char.equals(term)) {
      str.append(char)
      assert(str.size < maxLength)
      char = inputStream.readByte().toChar
    }
    str.toString
  }
  val inputStream: DataInputStream = new DataInputStream(new GZIPInputStream(new FileInputStream(file)))
  try {
    val header = readUntil(inputStream, '\n')
    val (records, dimensions) = header.split(" ") match {
      case Array(records, dimensions) => (records.toInt, dimensions.toInt)
    }
    new Word2VecModel((0 until records).toArray.map(recordIndex => {
      readUntil(inputStream, ' ') -> (0 until dimensions).map(dimensionIndex => {
        java.lang.Float.intBitsToFloat(java.lang.Integer.reverseBytes(inputStream.readInt()))
      }).toArray
    }).toMap)
  } finally {
    inputStream.close()
  }
}