package fasttext;

import java.io.BufferedInputStream;

import java.io.BufferedReader;

import java.io.DataInputStream;

import java.io.DataOutputStream;

import java.io.File;

import java.io.FileInputStream;

import java.io.FileOutputStream;

import java.io.FileReader;

import java.io.FileWriter;

import java.io.IOException;

import java.io.InputStreamReader;

import java.io.OutputStream;

import java.io.Writer;

import java.text.DecimalFormat;

import java.util.ArrayList;

import java.util.List;

import java.util.concurrent.atomic.AtomicLong;

import com.google.gson.JsonObject;

import org.apache.commons.math3.distribution.UniformIntegerDistribution;

import org.apache.commons.math3.distribution.UniformRealDistribution;

import org.apache.log4j.Logger;

import fasttext.Args.model_name;

import fasttext.Dictionary.entry_type;

public class FastText {

private static Logger logger = Logger.getLogger(FastText.class);

private static int SUPERVISED_LABEL_SIZE = 10;

public Args args = new Args();

public Dictionary dict = new Dictionary(args);

public Matrix input = new Matrix();

public Matrix output = new Matrix();

class Info {

long start = 0;

AtomicLong allWords = new AtomicLong(0l);

AtomicLong allN = new AtomicLong(0l);

double allLoss = 0.0;

}

public class MutableDouble {

private double value;

public MutableDouble(double value) {

this.value = value;

}

public void set(double value) {

this.value = value;

}

public double doubleValue() {

return value;

}

public void incrementDouble(double value) {

this.value += value;

}

}

Info info = new Info();

public void loadModel(String filename, Dictionary dict, Matrix input, Matrix output) throws IOException {

DataInputStream dis = null;

BufferedInputStream bis = null;

try {

File file = new File(filename);

if (!(file.exists() && file.isFile() && file.canRead())) {

throw new IOException("Model file cannot be opened for loading!");

}

bis = new BufferedInputStream(new FileInputStream(file));

dis = new DataInputStream(bis);

args.load(dis);

dict.load(dis);

input.load(dis);

output.load(dis);

logger.info("loadModel done!");

} finally {

bis.close();

dis.close();

}

}

public void getVector(Dictionary dict, Matrix input, Vector vec, String word) {

final java.util.Vector<Integer> ngrams = dict.getNgrams(word);

vec.zero();

for (Integer it : ngrams) {

vec.addRow(input, it);

}

if (ngrams.size() > 0) {

vec.mul((float) (1.0 / ngrams.size()));

}

}

public void printVectors(Dictionary dict, Matrix input) {

Vector vec = new Vector(args.dim);

@SuppressWarnings("resource")

java.util.Scanner scanner = new java.util.Scanner(System.in);

String word = scanner.nextLine();

while (!Utils.isEmpty(word)) {

getVector(dict, input, vec, word);

System.out.println(word + " " + vec);

word = scanner.nextLine();

}

}

public void printInfo(Model model, float progress) {

float loss = (float) (info.allLoss / info.allN.get());

float t = (float) ((System.currentTimeMillis() - info.start) / 1000);

float wst = (float) (info.allWords.get() / t);

int eta = (int) (t / progress * (1 - progress) / args.thread);

int etah = eta / 3600;

int etam = (eta - etah * 3600) / 60;

System.out.printf("\rProgress: %.1f%% words/sec/thread: %d lr: %.6f loss: %.6f eta: %d h %d m", 100 * progress,

(int) wst, model.getLearningRate(), loss, etah, etam);

}

public int supervised(Model model, final java.util.Vector<Integer> line, final java.util.Vector<Integer> labels,

MutableDouble loss, UniformIntegerDistribution uid) {

if (labels.size() == 0 || line.size() == 0)

return 0;

int i = uid.sample();

loss.incrementDouble(model.update(line, labels.get(i)));

return 1;

}

public int cbow(Dictionary dict, Model model, final java.util.Vector<Integer> line, MutableDouble loss,

UniformIntegerDistribution uid) {

java.util.Vector<Integer> bow = new java.util.Vector<Integer>();

int nexamples = 0;

for (int w = 0; w < line.size(); w++) {

int boundary = uid.sample();

bow.clear();

for (int c = -boundary; c <= boundary; c++) {

if (c != 0 && w + c >= 0 && w + c < line.size()) {

final java.util.Vector<Integer> ngrams = dict.getNgrams(line.get(w + c));

bow.addAll(ngrams);

}

}

loss.incrementDouble(model.update(bow, line.get(w)));

nexamples++;

}

return nexamples;

}

public int skipgram(Dictionary dict, Model model, final java.util.Vector<Integer> line, MutableDouble loss,

UniformIntegerDistribution uid) {

int nexamples = 0;

for (int w = 0; w < line.size(); w++) {

int boundary = uid.sample();

final java.util.Vector<Integer> ngrams = dict.getNgrams(line.get(w));

for (int c = -boundary; c <= boundary; c++) {

if (c != 0 && w + c >= 0 && w + c < line.size()) {

loss.incrementDouble(model.update(ngrams, line.get(w + c)));

nexamples++;

}

}

}

return nexamples;

}

public void test(Dictionary dict, Model model, String filename) throws IOException {

int nexamples = 0;

double precision = 0.0f;

java.util.Vector<Integer> line = new java.util.Vector<Integer>();

java.util.Vector<Integer> labels = new java.util.Vector<Integer>();

File file = new File(filename);

if (!(file.exists() && file.isFile() && file.canRead())) {

throw new IOException("Test file cannot be opened!");

}

UniformRealDistribution urd = new UniformRealDistribution(model.rng, 0, 1);

FileInputStream fis = new FileInputStream(file);

BufferedReader dis = new BufferedReader(new InputStreamReader(fis, "UTF-8"));

try {

String lineString;

while ((lineString = dis.readLine()) != null) {

dict.getLine(lineString, line, labels, urd);

dict.addNgrams(line, args.wordNgrams);

if (labels.size() > 0 && line.size() > 0) {

System.out.print("Test line: " + lineString);

JsonObject detail = new JsonObject();

int i = model.predict(line, detail);

logger.info(detail.toString());

if (labels.contains(i)) {

precision += 1.0;

System.out.println(" [HIT]: " + dict.getLabel(i));

} else {

System.out.println(" [MISSED]: " + dict.getLabel(i));

}

nexamples++;

//logger.info("Line = " + lineString + "\t" + "predict label = " + dict.getLabel(i) + "\t" +

// "Score = " + score.toString());

} else {

System.out.println("FAIL Test line: " + lineString + "labels: " + labels + " line: " + line);

}

}

} finally {

dis.close();

fis.close();

}

System.out.printf("P@1: %.3f%n", precision / nexamples);

System.out.println("Number of examples: " + nexamples);

}

public void predict(Dictionary dict, Model model, String filename) throws IOException {

// int nexamples = 0;

// double precision = 0.0;

java.util.Vector<Integer> line = new java.util.Vector<Integer>();

java.util.Vector<Integer> labels = new java.util.Vector<Integer>();

File file = new File(filename);

if (!(file.exists() && file.isFile() && file.canRead())) {

throw new IOException("Test file cannot be opened!");

}

UniformRealDistribution urd = new UniformRealDistribution(model.rng, 0, 1);

FileInputStream fis = new FileInputStream(file);

BufferedReader dis = new BufferedReader(new InputStreamReader(fis, "UTF-8"));

try {

String lineString;

while ((lineString = dis.readLine()) != null) {

dict.getLine(lineString, line, labels, urd);

dict.addNgrams(line, args.wordNgrams);

if (line.size() > 0) {

int i = model.predict(line);

System.out.println(lineString + "\t" + dict.getLabel(i));

} else {

System.out.println(lineString + "\tn/a");

}

}

} finally {

dis.close();

fis.close();

}

}

public void test(String binFile, String testFile) throws IOException {

loadModel(binFile, dict, input, output);

Model model = new Model(args, input, output, args.dim, (float) args.lr, 1);

model.setTargetCounts(dict.getCounts(entry_type.label));

test(dict, model, testFile);

}

public void predict(String binFile, String predictFile) throws IOException {

loadModel(binFile, dict, input, output);

Model model = new Model(args, input, output, args.dim, (float) args.lr, 1);

model.setTargetCounts(dict.getCounts(entry_type.label));

predict(dict, model, predictFile);

}

public void printVectors(String binFile) throws IOException {

loadModel(binFile, dict, input, output);

printVectors(dict, input);

}

int threadCount;

public void train(String[] args_) throws IOException {

args.parseArgs(args_);

File file = new File(args.input);

if (!(file.exists() && file.isFile() && file.canRead())) {

throw new IOException("Input file cannot be opened! " + args.input);

}

dict.readFromFile(args.input);

input = new Matrix(dict.nwords() + args.bucket, args.dim);

if (args.model == model_name.sup) {

output = new Matrix(dict.nlabels(), args.dim);

} else {

output = new Matrix(dict.nwords(), args.dim);

}

input.uniform((float) (1.0 / args.dim));

output.zero();

info.start = System.currentTimeMillis();

long t0 = System.currentTimeMillis();

threadCount = args.thread;

long fileSize = Utils.sizeLine(args.input);

for (int i = 0; i < args.thread; i++) {

new TrainThread(this, dict, input, output, i, fileSize).start();

}

synchronized (this) {

while (threadCount > 0) {

try {

wait();

} catch (InterruptedException ignored) {

}

}

}

long trainTime = (System.currentTimeMillis() - t0) / 1000;

System.out.printf("Train time: %d sec\n", trainTime);

if (!Utils.isEmpty(args.output)) {

saveModel(dict, input, output);

saveVectors(dict, input, output);

}

}

public class TrainThread extends Thread {

final FastText ft;

Dictionary dict;

Matrix input;

Matrix output;

int threadId;

long fileSize;

public TrainThread(FastText ft, Dictionary dict, Matrix input, Matrix output, int threadId, long fileSize) {

this.ft = ft;

this.dict = dict;

this.input = input;

this.output = output;

this.threadId = threadId;

this.fileSize = fileSize;

}

public void run() {

if (logger.isDebugEnabled()) {

logger.debug("thread: " + threadId + " RUNNING!");

}

BufferedReader br = null;

try {

br = new BufferedReader(new FileReader(args.input));

Utils.seek(br, threadId * fileSize / args.thread);

Model model = new Model(args, input, output, args.dim, (float) args.lr, threadId);

if (args.model == model_name.sup) {

model.setTargetCounts(dict.getCounts(entry_type.label));

} else {

model.setTargetCounts(dict.getCounts(entry_type.word));

}

float progress;

final long ntokens = dict.ntokens();

long tokenCount = 0, /** printCount = 0, */

deltaCount = 0;

MutableDouble loss = new MutableDouble(0.0);

long nexamples = 0;

java.util.Vector<Integer> line = new java.util.Vector<Integer>();

java.util.Vector<Integer> labels = new java.util.Vector<Integer>();

UniformRealDistribution urd = new UniformRealDistribution(model.rng, 0, 1);

List<UniformIntegerDistribution> learnUid0 = new ArrayList<UniformIntegerDistribution>();

UniformIntegerDistribution learnUid = null;

if (args.model == model_name.sup) {

for (int i = 0; i <= SUPERVISED_LABEL_SIZE; i++) {

learnUid0.add(new UniformIntegerDistribution(model.rng, 0, i));

}

} else if (args.model == model_name.cbow) {

learnUid = new UniformIntegerDistribution(model.rng, 1, args.ws);

} else if (args.model == model_name.sg) {

learnUid = new UniformIntegerDistribution(model.rng, 1, args.ws);

}

String lineString;

while (info.allWords.get() < args.epoch * ntokens) {

lineString = br.readLine();

if (lineString == null) {

try {

br.close();

br = new BufferedReader(new FileReader(args.input));

if (logger.isDebugEnabled()) {

logger.debug("Input file reloaded!");

}

} catch (Exception e) {

e.printStackTrace();

}

lineString = br.readLine();

}

while (Utils.isEmpty(lineString) || lineString.startsWith("#")) {

lineString = br.readLine();

}

deltaCount = dict.getLine(lineString, line, labels, urd);

tokenCount += deltaCount;

// printCount += deltaCount;

if (args.model == model_name.sup) {

dict.addNgrams(line, args.wordNgrams);

if (labels.size() == 0 || line.size() == 0) {

continue;

}

learnUid = learnUid0.get(labels.size() - 1);

nexamples += supervised(model, line, labels, loss, learnUid);

} else if (args.model == model_name.cbow) {

nexamples += cbow(dict, model, line, loss, learnUid);

} else if (args.model == model_name.sg) {

nexamples += skipgram(dict, model, line, loss, learnUid);

}

if (tokenCount > args.lrUpdateRate) {

info.allWords.addAndGet(tokenCount);

info.allLoss += loss.doubleValue();

info.allN.addAndGet(nexamples);

tokenCount = 0;

loss.set(0.0);

nexamples = 0;

progress = (float) (info.allWords.get()) / (args.epoch * ntokens);

model.setLearningRate((float) (args.lr * (1.0 - progress)));

if (threadId == 0) {

printInfo(model, progress);

}

}

}

if (threadId == 0) {

printInfo(model, 1.0f);

System.out.println();

}

} catch (IOException e) {

e.printStackTrace();

System.exit(1);

} finally {

if (br != null)

try {

br.close();

} catch (IOException e) {

e.printStackTrace();

}

}

// exit from thread

synchronized (ft) {

if (logger.isDebugEnabled()) {

logger.debug("thread: " + threadId + " EXIT!");

}

ft.threadCount--;

ft.notify();

}

}

}

public void saveModel(Dictionary dict, Matrix input, Matrix output) throws IOException {

File file = new File(args.output + ".bin");

if (file.exists()) {

file.delete();

}

logger.info("Saving model to " + file.getAbsolutePath().toString());

FileOutputStream fos = new FileOutputStream(file);

OutputStream ofs = new DataOutputStream(fos);

try {

logger.debug("writing args");

args.save(ofs);

logger.debug("writing dict");

dict.save(ofs);

logger.debug("writing input");

input.save(ofs);

logger.debug("writing output");

output.save(ofs);

} finally {

ofs.flush();

ofs.close();

}

}

public void saveVectors(Dictionary dict, Matrix input, Matrix output) throws IOException {

File file = new File(args.output + ".vec");

if (file.exists()) {

file.delete();

}

logger.info("Saving Vectors to " + file.getAbsolutePath().toString());

Writer writer = new FileWriter(file);

try {

writer.write(dict.nwords());

writer.write(" ");

writer.write(args.dim);

writer.write("\n");

Vector vec = new Vector(args.dim);

DecimalFormat df = new DecimalFormat("0.#####");

for (int i = 0; i < dict.nwords(); i++) {

String word = dict.getWord(i);

getVector(dict, input, vec, word);

writer.write(word);

writer.write(" ");

writer.write(" ");

for (int j = 0; i < vec.m_; i++) {

writer.write(df.format(vec.data_[j]));

writer.write(" ");

}

writer.write("\n");

}

} finally {

writer.flush();

writer.close();

}

}

public static void printUsage() {

System.out.print(

"usage: java -jar fasttext.jar <command> <args>\n\n" + "The commands supported by fasttext are:\n\n"

+ " supervised train a supervised classifier\n" + " test evaluate a supervised classifier\n"

+ " predict predict most likely label\n" + " skipgram train a skipgram model\n"

+ " cbow train a cbow model\n" + " print-vectors print vectors given a trained model\n");

}

public static void printTestUsage() {

System.out.print("usage: java -jar fasttext.jar test <model> <test-data>\n\n" + " <model> model filename\n"

+ " <test-data> test data filename\n");

}

public static void printPredictUsage() {

System.out.print("usage: java -jar fasttext.jar predict <model> <test-data>\n\n" + " <model> model filename\n"

+ " <test-data> test data filename\n");

}

public static void printPrintVectorsUsage() {

System.out.print("usage: java -jar fasttext.jar print-vectors <model>\n\n" + " <model> model filename\n");

}

public static void main(String[] args) {

org.apache.log4j.PropertyConfigurator.configure("log4j.properties");

FastText op = new FastText();

if (args.length == 0) {

printUsage();

System.exit(1);

}

try {

String command = args[0];

if ("skipgram".equalsIgnoreCase(command) || "cbow".equalsIgnoreCase(command)

|| "supervised".equalsIgnoreCase(command)) {

op.train(args);

} else if ("test".equalsIgnoreCase(command)) {

if (args.length != 3) {

printTestUsage();

System.exit(1);

}

op.test(args[1], args[2]);

} else if ("print-vectors".equalsIgnoreCase(command)) {

op.printVectors(args[1]);

} else if ("predict".equalsIgnoreCase(command)) {

if (args.length != 3) {

printPredictUsage();

System.exit(1);

}

op.predict(args[1], args[2]);

} else {

printUsage();

System.exit(1);

}

} catch (Exception e) {

e.printStackTrace();

System.exit(1);

}

System.exit(0);

}

}