code: package edu.jhu.thrax.hadoop.features;
  
  import java.io.IOException;
  import java.util.HashMap;
  import java.util.Iterator;
  
  import org.apache.hadoop.conf.Configuration;
  import org.apache.hadoop.conf.Configured;
  import org.apache.hadoop.io.FloatWritable;
  import org.apache.hadoop.io.IntWritable;
  import org.apache.hadoop.io.LongWritable;
  import org.apache.hadoop.io.Text;
  import org.apache.hadoop.mapreduce.Mapper;
  import org.apache.hadoop.mapreduce.Partitioner;
  import org.apache.hadoop.mapreduce.Reducer;
  
  import edu.jhu.thrax.datatypes.AlignedSentencePair;
  import edu.jhu.thrax.datatypes.Alignment;
  import edu.jhu.thrax.hadoop.jobs.WordLexprobJob;
  import edu.jhu.thrax.util.Vocabulary;
  import edu.jhu.thrax.util.exceptions.MalformedInputException;
  import edu.jhu.thrax.util.io.InputUtilities;
  
  public class WordLexicalProbabilityCalculator extends Configured {
    public static final long UNALIGNED = 0x0000000000000000L;
    public static final long MARGINAL = 0x0000000000000000L;
  
    public static class Map extends Mapper<LongWritable, Text, LongWritable, IntWritable> {
      private HashMap<Long, Integer> counts = new HashMap<Long, Integer>();
      private boolean sourceParsed;
      private boolean targetParsed;
      private boolean reverse;
      private boolean sourceGivenTarget;
  
      protected void setup(Context context) throws IOException, InterruptedException {
        Configuration conf = context.getConfiguration();
        String vocabulary_path = conf.getRaw("thrax.work-dir") + "vocabulary/part-*";
        Vocabulary.initialize(conf, vocabulary_path);
  
        sourceParsed = conf.getBoolean("thrax.source-is-parsed", false);
        targetParsed = conf.getBoolean("thrax.target-is-parsed", false);
        reverse = conf.getBoolean("thrax.reverse", false);
        sourceGivenTarget = conf.getBoolean(WordLexprobJob.SOURCE_GIVEN_TARGET, false);
      }
  
      public void map(LongWritable key, Text value, Context context) throws IOException,
          InterruptedException {
        counts.clear();
        String line = value.toString();
        AlignedSentencePair sentencePair;
        try {
          sentencePair =
              InputUtilities.alignedSentencePair(line, sourceParsed, targetParsed,
                  !(reverse ^ sourceGivenTarget));
        } catch (MalformedInputException e) {
          context.getCounter("input errors", e.getMessage()).increment(1);
          return;
        }
        int[] source = sentencePair.source;
        int[] target = sentencePair.target;
        Alignment alignment = sentencePair.alignment;
  
        for (int i = 0; i < source.length; i++) {
          int src = source[i];
          if (alignment.sourceIndexIsAligned(i)) {
            Iterator<Integer> target_indices = alignment.targetIndicesAlignedTo(i);
            while (target_indices.hasNext()) {
              int tgt = target[target_indices.next()];
              long pair = ((long) tgt << 32) + src;
              long marginal = ((long) tgt << 32) + MARGINAL;
              counts.put(pair, counts.containsKey(pair) ? counts.get(pair) + 1 : 1);
              counts.put(marginal, counts.containsKey(marginal) ? counts.get(marginal) + 1 : 1);
            }
          } else {
            long pair = UNALIGNED | ((long) src);
            long marginal = UNALIGNED;
            counts.put(pair, counts.containsKey(pair) ? counts.get(pair) + 1 : 1);
            counts.put(marginal, counts.containsKey(marginal) ? counts.get(marginal) + 1 : 1);
          }
        }
        for (long pair : counts.keySet())
          context.write(new LongWritable(pair), new IntWritable(counts.get(pair)));
      }
    }
  
    public static class Reduce
        extends Reducer<LongWritable, IntWritable, LongWritable, FloatWritable> {
      private int current = -1;
      private int marginalCount;
  
      protected void setup(Context context) throws IOException, InterruptedException {
        Configuration conf = context.getConfiguration();
        
        // TODO: remove unnecessary vocabulary loads?
        String vocabulary_path = conf.getRaw("thrax.work-dir") + "vocabulary/part-*";
        Vocabulary.initialize(conf, vocabulary_path);
      }
  
      protected void reduce(LongWritable key, Iterable<IntWritable> values, Context context)
          throws IOException, InterruptedException {
        long pair = key.get();
        int tgt = (int) (pair >> 32);
        int src = (int) (pair & 0xFFFFFFFFL);
  
        if (tgt != current) {
          if (src != MARGINAL) throw new RuntimeException("Sorting something before marginal.");
          current = tgt;
          marginalCount = 0;
          for (IntWritable x : values)
            marginalCount += x.get();
          return;
        }
        // Control only gets here if we are using the same marginal
        int my_count = 0;
        for (IntWritable x : values)
          my_count += x.get();
        context.write(key, new FloatWritable(my_count / (float) marginalCount));
      }
    }
  
    public static class Partition extends Partitioner<LongWritable, IntWritable> {
      public int getPartition(LongWritable key, IntWritable value, int numPartitions) {
        // ids range from 0 to size-1
        // we partition like this:
        // partition 1: 0 to num_elements_per_partition-1
        // partition 2: num_elements_per_partition to 2*num_elements_per_partition-1
        // ...
        // partition numPartitions: num_elements_per_partition*(numPartitions-1) to num_elements_per_partition*numPartitions-1 = size-1
        final int max_trg_plus_one = Vocabulary.size();
        final int num_elements_per_partition = (int) Math.ceil(max_trg_plus_one / (1.0 * numPartitions));
        int trg = ((int) (key.get() >> 32) & Integer.MAX_VALUE);
        if (trg < 0 || trg >= max_trg_plus_one) {
          throw new RuntimeException(String.format("Word id %d out of range %d %d", trg, 0, max_trg_plus_one-1));
        }
        return trg / num_elements_per_partition;
      }
    }
  }
  