package cn.edu.bjut.chapter2;

import java.io.File;
import java.util.Arrays;
import java.util.Collection;

import org.deeplearning4j.models.word2vec.Word2Vec;
import org.deeplearning4j.text.documentiterator.FileDocumentIterator;

public class Word2VecExample {
	public static void main(String[] args) throws Exception {
		// Step 1: Prepare the Dataset
		String trainFile = "data/train.txt";
		FileDocumentIterator iterator = new FileDocumentIterator(new File(trainFile));

		// Step 2: Build the Word2Vec Model
		Word2Vec word2Vec = new Word2Vec.Builder()
				.iterate(iterator).epochs(5)
				.layerSize(100).minWordFrequency(1)
				.seed(42).build();

		// Step 3: Fit the Model
		word2Vec.fit();

		// Get a word vector
		double[] vector = word2Vec.getWordVector("Java");
		System.out.println("Java's Vector: " + Arrays.toString(vector));

		// Find closest words to 'Java'
		System.out.println("Closest words to 'Java': ");
		Collection<String> nearestWords = word2Vec.wordsNearest("Java", 5);
		System.out.println(nearestWords);
	}
}