hex.word2vec.Word2VecMojoWriter Maven / Gradle / Ivy

Go to download

Show more of this group Show more artifacts with this name
Show all versions of h2o-algos Show documentation

H2O Algorithms

There is a newer version: 3.46.0.6

package hex.word2vec;

import hex.ModelMojoWriter;
import water.MemoryManager;
import water.parser.BufferedString;

import java.io.IOException;
import java.nio.ByteBuffer;

/**
 * MOJO serializer for word2vec model.
 */
public class Word2VecMojoWriter extends ModelMojoWriter {

  @SuppressWarnings("unused")  // Called through reflection in ModelBuildersHandler
  public Word2VecMojoWriter() {}

  public Word2VecMojoWriter(Word2VecModel model) {
    super(model);
  }

  @Override public String mojoVersion() {
    return "1.00";
  }

  @Override
  protected void writeModelData() throws IOException {
    writekv("vec_size", model._parms._vec_size);
    writekv("vocab_size", model._output._words.length);

    // Vocabulary
    startWritingTextFile("vocabulary");
    for (BufferedString word : model._output._words) {
      writeln(word.toString(), true);
    }
    finishWritingTextFile();

    // Vectors
    ByteBuffer bb = ByteBuffer.wrap(MemoryManager.malloc1(model._output._vecs.length * 4));
    for (float v : model._output._vecs)
      bb.putFloat(v);
    writeblob("vectors", bb.array());
  }

}