vn.vitk-master/src/main/java/vn/vitk/tag/CMMModel.java
package vn.vitk.tag;

import java.io.IOException;
import java.io.Serializable;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.HashSet;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.UUID;

import org.apache.hadoop.fs.Path;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.function.Function;
import org.apache.spark.ml.Model;
import org.apache.spark.ml.PipelineModel;
import org.apache.spark.ml.feature.CountVectorizerModel;
import org.apache.spark.ml.feature.StringIndexerModel;
import org.apache.spark.ml.param.ParamMap;
import org.apache.spark.ml.util.DefaultParamsReader;
import org.apache.spark.ml.util.DefaultParamsWriter;
import org.apache.spark.ml.util.MLReader;
import org.apache.spark.ml.util.MLWritable;
import org.apache.spark.ml.util.MLWriter;
import org.apache.spark.ml.util.SchemaUtils;
import org.apache.spark.mllib.linalg.Vector;
import org.apache.spark.mllib.linalg.Vectors;
import org.apache.spark.mllib.util.MLUtils;
import org.apache.spark.sql.DataFrame;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.RowFactory;
import org.apache.spark.sql.types.DataTypes;
import org.apache.spark.sql.types.Metadata;
import org.apache.spark.sql.types.StructField;
import org.apache.spark.sql.types.StructType;

import scala.Tuple2;
import scala.collection.Iterator;
import scala.collection.mutable.WrappedArray;
import vn.vitk.util.Constants;

/**
 * @author Phuong LE-HONG
 *         <p>
 *         May 9, 2016, 5:52:38 PM
 *         <p>
 *         A conditional Markov model (or Maximum-Entropy Markov model -- MEMM)
 *         for sequence labeling which is fitted by a CMM graphical model.
 * 
 */
public class CMMModel extends Model<CMMModel> implements MLWritable {

	private static final long serialVersionUID = -4855076361905836432L;
	
	private PipelineModel pipelineModel;
	private ContextExtractor contextExtractor;
	private final Vector weights;
	private final String[] tags;
	private final Map<String, Integer> featureMap;
	private final Map<String, Set<Integer>> tagDictionary; 
	
	/**
	 * Creates a conditional Markov model.
	 * @param pipelineModel
	 * @param weights
	 * @param markovOrder
	 */
	public CMMModel(PipelineModel pipelineModel, Vector weights, MarkovOrder markovOrder, Map<String, Set<Integer>> tagDictionary) {
		this.pipelineModel = pipelineModel;
		this.contextExtractor = new ContextExtractor(markovOrder, Constants.REGEXP_FILE);
		this.weights = weights;
		this.tags = ((StringIndexerModel)(pipelineModel.stages()[2])).labels();
		String[] features = ((CountVectorizerModel)(pipelineModel.stages()[1])).vocabulary();
		featureMap = new HashMap<String, Integer>();
		for (int j = 0; j < features.length; j++) {
			featureMap.put(features[j], j);
		}
		this.tagDictionary = tagDictionary;
	}
	
	@Override
	public CMMModel copy(ParamMap extra) {
		return defaultCopy(extra);
	}

	/**
	 * An immutable unique ID for the object and its derivatives.
	 * @return an immutable unique ID for the object and its derivatives.
	 */
	@Override
	public String uid() {
		String ruid = UUID.randomUUID().toString();
		int n = ruid.length();
		return "cmmModel" + "_" + ruid.substring(n-12, n);
	}

	/* (non-Javadoc)
	 * @see org.apache.spark.ml.Transformer#transform(org.apache.spark.sql.DataFrame)
	 */
	@Override
	public DataFrame transform(DataFrame dataset) {
		JavaRDD<Row> output = dataset.javaRDD().map(new DecodeFunction());
		StructType schema = new StructType(new StructField[]{
			new StructField("sentence", DataTypes.StringType, false, Metadata.empty()),
			new StructField("prediction", DataTypes.StringType, false, Metadata.empty())
		});
		return dataset.sqlContext().createDataFrame(output, schema);
	}

	class DecodeFunction implements Function<Row, Row> {
		private static final long serialVersionUID = 5026042203808959533L;

		@Override
		public Row call(Row row) throws Exception {
			List<String> words = Arrays.asList(row.getString(0).split("\\s+"));
			int n = words.size();
			List<Tuple2<String, String>> sequence = new ArrayList<Tuple2<String, String>>(n);
			for (int i = 0; i < n; i++) {
				sequence.add(new Tuple2<String, String>(words.get(i), "UNK"));
			}
			List<String> partsOfSpeech = decode(sequence);
			StringBuilder sb = new StringBuilder();
			for (String pos : partsOfSpeech) {
				sb.append(pos);
				sb.append(' ');
			}
			return RowFactory.create(row.getString(0), sb.toString().trim());
		}
		
		/**
		 * Finds the best label sequence for an observation sequence.
		 * @param sequence
		 * @return a label sequence.
		 */
		private List<String> decode(List<Tuple2<String, String>> sequence) {
			int n = sequence.size();
			double[][] score = new double[tags.length][n];
					
			for (int j = 0; j < n; j++) {
				LabeledContext context = contextExtractor.extract(sequence, j);
				Tuple2<double[], String> tuple = probability(context);
				double[] prob = tuple._1();
				for (int i = 0; i < prob.length; i++) {
					score[i][j] = prob[i];
				}
				// update the tag at position j for the next incremental extraction
				// 
				sequence.set(j, new Tuple2<String, String>(sequence.get(j)._1(), tuple._2()));
			}
			
			ViterbiDecoder decoder = new ViterbiDecoder(score);
			int[] path = decoder.bestPath();
			List<String> partsOfSpeech = new LinkedList<String>();
			for (int k : path) {
				partsOfSpeech.add(tags[k]);
			}
			return partsOfSpeech;
		}
		
		/**
		 * Computes the probability mass function (pmf) of an unlabeled context. 
		 * @param context
		 * @return a tuple of the pmf over the tagset and the best tag.
		 */
		private Tuple2<double[], String> probability(LabeledContext context) {
			String[] fs = context.getFeatureStrings().toLowerCase().split("\\s+");
			Set<String> fsset = new HashSet<String>();
			for (String f : fs) {
				fsset.add(f);
			}
			List<Tuple2<Integer, Double>> x = new LinkedList<Tuple2<Integer, Double>>();
			for (String f : fsset) {
				Integer i = featureMap.get(f);
				if (i != null) {
					x.add(new Tuple2<Integer, Double>(i, 1.0));
				}
			}
			Vector features = MLUtils.appendBias(Vectors.sparse(featureMap.size(), x));
			int numLabels = tags.length;
			double[] score = new double[numLabels];
			Arrays.fill(score, 0d);
			int maxLabel = 0;
			double maxScore = 0d;
			String word = context.getWord();
			Set<Integer> labels = tagDictionary.get(word);
			if (labels == null) { // this is a rare/unknown word, try all possible labels
				labels = new HashSet<Integer>();
				for (int k = 1; k < numLabels; k++) // k goes from 1 since we do not need to compute score[0], it is always 0.0.
					labels.add(k);
			}
			
			int d = features.size();
			for (int k : labels) {
				if (k > 0) {
					for (int j : features.toSparse().indices())
						score[k] += weights.apply((k-1) * d + j);
					if (score[k] > maxScore) {
						maxScore = score[k];
						maxLabel = k;
					}
				}
			}
			
			// prevent possible numerical overflow error in the case maxScore > 0
			if (maxScore > 0) {
				for (int k = 0; k < numLabels; k++)
					score[k] -= maxScore;
			}
			// normalize the score to get probability
			double z = 0d;
			for (int k = 0; k < numLabels; k++) {
				score[k] = Math.exp(score[k]); 
				z += score[k];
			}
			for (int k = 0; k < numLabels; k++) {
				score[k] /= z;
			}
			return new Tuple2<double[], String>(score, tags[maxLabel]);
		}
	}
	
	
	/* (non-Javadoc)
	 * @see org.apache.spark.ml.PipelineStage#transformSchema(org.apache.spark.sql.types.StructType)
	 */
	@Override
	public StructType transformSchema(StructType schema) {
		return SchemaUtils.appendColumn(schema, new StructField("prediction", DataTypes.StringType, false, Metadata.empty()));
	}

	/* (non-Javadoc)
	 * @see org.apache.spark.ml.util.MLWritable#save(java.lang.String)
	 */
	@Override
	public void save(String path) throws IOException {
		write().overwrite().save(path);
	}

	/* (non-Javadoc)
	 * @see org.apache.spark.ml.util.MLWritable#write()
	 */
	@Override
	public MLWriter write() {
		return new CMMModelWriter(this);
	}
	
	private class CMMModelWriter extends MLWriter {
		CMMModel instance;
		
		public CMMModelWriter(CMMModel instance) {
			this.instance = instance;
		}
		
		@Override
		public void saveImpl(String path) {
			// save metadata and params
			DefaultParamsWriter.saveMetadata(instance, path, sc(), 
					DefaultParamsWriter.saveMetadata$default$4(),
					DefaultParamsWriter.saveMetadata$default$5());

			// save model data: markovOrder, numLabels, weights
			Data data = new Data();
			data.setMarkovOrder(contextExtractor.getMarkovOrder().ordinal()+1);
			data.setWeights(weights);
			data.setTagDictionary(tagDictionary);
			List<Data> list = new LinkedList<Data>();
			list.add(data);
			String dataPath = new Path(path, "data").toString();
			sqlContext().createDataFrame(list, Data.class).write().parquet(dataPath);
			// save pipeline model
			try {
				String pipelinePath = new Path(path, "pipelineModel").toString(); 
				pipelineModel.write().overwrite().save(pipelinePath);
			} catch (IOException e) {
				e.printStackTrace();
			}
		}
	}
	
	public class Data implements Serializable {

		private static final long serialVersionUID = 1L;
		private int markovOrder;
		private Vector weights;
		private Map<String, Set<Integer>> tagDictionary;
		
		public void setMarkovOrder(int markovOrder) {
			this.markovOrder = markovOrder;
		}
		
		public int getMarkovOrder() {
			return markovOrder;
		}
		
		public void setWeights(Vector weights) {
			this.weights = weights;
		}
		
		public Vector getWeights() {
			return weights;
		}
		
		public Map<String, Set<Integer>> getTagDictionary() {
			return tagDictionary;
		}
		
		public void setTagDictionary(Map<String, Set<Integer>> tagDictionary) {
			this.tagDictionary = tagDictionary;
		}
	}
	
	/* (non-Javadoc)
	 * @see org.apache.spark.ml.PipelineStage#toString()
	 */
	@Override
	public String toString() {
		StringBuilder sb = new StringBuilder();
		sb.append("[markovOrder=");
		sb.append(contextExtractor.getMarkovOrder());
		sb.append(", ");
		sb.append("numLabels = ");
		sb.append(tags.length);
		sb.append(", ");
		sb.append(", weights=");
		sb.append(weights.toString());
		sb.append(']');
		return sb.toString();
	}

	/**
	 * Loads a {@link CMMModel} from an external file.
	 * @param path
	 * @return a CMM model.
	 */
	public static CMMModel load(String path) {
		return read().load(path);
	}

	/**
	 * This functions is used in the reflection framework of Spark ML.
	 * @return a {@link MLReader}.
	 */
	public static MLReader<CMMModel> read() {
		return new CMMModelReader();
	}
	
	private static class CMMModelReader extends MLReader<CMMModel> {
		@Override
		public CMMModel load(String path) {
			org.apache.spark.ml.util.DefaultParamsReader.Metadata metadata = DefaultParamsReader.loadMetadata(path, sc(), CMMModel.class.getName());
			String pipelinePath = new Path(path, "pipelineModel").toString();
			PipelineModel pipelineModel = PipelineModel.load(pipelinePath);
			String dataPath = new Path(path, "data").toString();
			DataFrame df = sqlContext().read().format("parquet").load(dataPath);
			Row row = df.select("markovOrder", "weights", "tagDictionary").head();
			// load the Markov order
			MarkovOrder order = MarkovOrder.values()[row.getInt(0)-1];
			// load the weight vector
			Vector w = row.getAs(1);
			// load the tag dictionary
			@SuppressWarnings("unchecked")
			scala.collection.immutable.HashMap<String, WrappedArray<Integer>> td = (scala.collection.immutable.HashMap<String, WrappedArray<Integer>>)row.get(2);
			Map<String, Set<Integer>> tagDict = new HashMap<String, Set<Integer>>();
			Iterator<Tuple2<String, WrappedArray<Integer>>> iterator = td.iterator();
			while (iterator.hasNext()) {
				Tuple2<String, WrappedArray<Integer>> tuple = iterator.next();
				Set<Integer> labels = new HashSet<Integer>();
				scala.collection.immutable.List<Integer> list = tuple._2().toList();
				for (int i = 0; i < list.size(); i++)
					labels.add(list.apply(i));
				tagDict.put(tuple._1(), labels);
			}
			// build a CMM model
			CMMModel model = new CMMModel(pipelineModel, w, order, tagDict);
			DefaultParamsReader.getAndSetParams(model, metadata);
			return model;
		}
	}

}
