content: package org.allenai.scienceparse;
  
  import com.gs.collections.api.map.primitive.MutableCharIntMap;
  import com.gs.collections.api.set.MutableSet;
  import com.gs.collections.api.set.primitive.MutableIntSet;
  import com.gs.collections.api.tuple.Pair;
  import com.gs.collections.impl.factory.primitive.CharIntMaps;
  import com.gs.collections.impl.factory.primitive.IntSets;
  import com.gs.collections.impl.set.mutable.UnifiedSet;
  import com.gs.collections.impl.tuple.Tuples;
  
  import lombok.Data;
  import lombok.val;
  import org.allenai.datastore.Datastore;
  import org.allenai.ml.eval.TrainCriterionEval;
  import org.allenai.ml.linalg.DenseVector;
  import org.allenai.ml.linalg.Vector;
  import org.allenai.ml.sequences.Evaluation;
  import org.allenai.ml.sequences.StateSpace;
  import org.allenai.ml.sequences.crf.CRFFeatureEncoder;
  import org.allenai.ml.sequences.crf.CRFModel;
  import org.allenai.ml.sequences.crf.CRFTrainer;
  import org.allenai.ml.sequences.crf.CRFWeightsEncoder;
  import org.allenai.ml.util.IOUtils;
  import org.allenai.ml.util.Indexer;
  import org.allenai.ml.util.Parallel;
  import org.allenai.pdffigures2.FigureExtractor;
  import org.allenai.scienceparse.pdfapi.PDFDoc;
  import org.allenai.scienceparse.pdfapi.PDFExtractor;
  import org.apache.pdfbox.pdmodel.PDDocument;
  import org.nustaq.serialization.FSTObjectInput;
  import org.nustaq.serialization.FSTObjectOutput;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  import scala.compat.java8.ScalaStreamSupport;
  import scala.compat.java8.OptionConverters;
  
  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.IOException;
  import java.io.InputStream;
  import java.io.RandomAccessFile;
  import java.nio.channels.Channels;
  import java.nio.channels.FileChannel;
  import java.nio.channels.FileLock;
  import java.nio.file.Path;
  import java.nio.file.Paths;
  import java.text.Normalizer;
  import java.util.ArrayDeque;
  import java.util.ArrayList;
  import java.util.Arrays;
  import java.util.Collection;
  import java.util.Collections;
  import java.util.Comparator;
  import java.util.Iterator;
  import java.util.List;
  import java.util.Map;
  import java.util.Optional;
  import java.util.Queue;
  import java.util.Random;
  import java.util.Set;
  import java.util.Timer;
  import java.util.TimerTask;
  import java.util.concurrent.ConcurrentMap;
  import java.util.concurrent.ConcurrentSkipListMap;
  import java.util.concurrent.ExecutionException;
  import java.util.concurrent.ExecutorService;
  import java.util.concurrent.Executors;
  import java.util.concurrent.Future;
  import java.util.concurrent.TimeUnit;
  import java.util.concurrent.atomic.AtomicInteger;
  import java.util.concurrent.atomic.AtomicReference;
  import java.util.function.Function;
  import java.util.function.ToDoubleFunction;
  import java.util.stream.Collectors;
  
  import static java.util.stream.Collectors.toList;
  
  public class Parser {
  
    public static final int MAXHEADERWORDS = 500; //set to something high for author/title parsing
    public static final String DATA_VERSION = "0.3";  // faster serialization
    private CRFModel<String, PaperToken, String> model;
    private ExtractReferences referenceExtractor;
  
    private final static Logger logger =
            LoggerFactory.getLogger(Parser.class);
    private final static Logger labeledDataLogger =
            LoggerFactory.getLogger(logger.getName() + ".labeledData");
  
    private static final Datastore datastore = Datastore.apply();
    public static Path getDefaultProductionModel() {
      return datastore.filePath("org.allenai.scienceparse", "productionModel.dat", 9);
    }
    public static Path getDefaultGazetteer() {
      return datastore.filePath("org.allenai.scienceparse", "gazetteer.json", 5);
    }
    public static Path getDefaultGazetteerDir() {
      return datastore.directoryPath("org.allenai.scienceparse", "kermit-gazetteers", 1);
    }
    public static Path getDefaultBibModel() {
      return datastore.filePath("org.allenai.scienceparse", "productionBibModel.dat", 7);
    }
  
    private static Parser defaultParser = null;
    public synchronized static Parser getInstance() throws Exception {
      if(defaultParser == null)
        defaultParser = new Parser();
      return defaultParser;
    }
  
    public Parser() throws Exception {
      this(getDefaultProductionModel(), getDefaultGazetteer(), getDefaultBibModel());
    }
  
    public Parser(
            final String modelFile,
            final String gazetteerFile,
            final String bibModelFile
    ) throws Exception {
      this(new File(modelFile), new File(gazetteerFile), new File(bibModelFile));
    }
  
    public Parser(
            final Path modelFile,
            final Path gazetteerFile,
            final Path bibModelFile
    ) throws Exception {
      this(modelFile.toFile(), gazetteerFile.toFile(), bibModelFile.toFile());
    }
  
    public Parser(
            final File modelFile,
            final File gazetteerFile,
            final File bibModelFile
    ) throws Exception {
      // Load main model in one thread, and the rest in another thread, to speed up startup.
      final AtomicReference<Exception> exceptionThrownByModelLoaderThread = new AtomicReference<>();
      final Thread modelLoaderThread = new Thread(new Runnable() {
        @Override
        public void run() {
          logger.info("Loading model from {}", modelFile);
          try(
              final DataInputStream modelIs =
                  new DataInputStream(new FileInputStream(modelFile));
          ) {
            model = loadModel(modelIs);
            logger.info("Loaded model from {}", modelFile);
          } catch(final Exception e) {
            exceptionThrownByModelLoaderThread.compareAndSet(null, e);
            logger.warn("Failed loading model from {}", modelFile);
          }
        }
      }, "ModelLoaderThread");
      modelLoaderThread.start();
  
      // Load non-main model stuff
      logger.info("Loading gazetteer from {}", gazetteerFile);
      logger.info("Loading bib model from {}", bibModelFile);
      try(
        final InputStream gazetteerIs = new FileInputStream(gazetteerFile);
        final DataInputStream bibModelIs = new DataInputStream(new FileInputStream(bibModelFile))
      ) {
        // Loading the gazetteer takes a long time, so we create a cached binary version of it that
        // loads very quickly. If that version is already there, we use it. Otherwise, we create it.
        val gazCacheFilename = String.format(
            "%s-%08x.gazetteerCache.bin",
            gazetteerFile.getName(),
            gazetteerFile.getCanonicalPath().hashCode());
        val gazCachePath = Paths.get(System.getProperty("java.io.tmpdir"), gazCacheFilename);
        try (
          final RandomAccessFile gazCacheFile = new RandomAccessFile(gazCachePath.toFile(), "rw");
          final FileChannel gazCacheChannel = gazCacheFile.getChannel();
          final FileLock gazCacheLock = gazCacheChannel.lock();
        ) {
          if (gazCacheChannel.size() == 0) {
            logger.info("Creating gazetteer cache at {}", gazCachePath);
            referenceExtractor =
                ExtractReferences.createAndWriteGazCache(
                    gazetteerIs,
                    bibModelIs,
                    Channels.newOutputStream(gazCacheChannel));
          } else {
            logger.info("Reading from gazetteer cache at {}", gazCachePath);
            referenceExtractor = new ExtractReferences(
                gazetteerIs,
                bibModelIs,
                Channels.newInputStream(gazCacheChannel));
          }
        }
      }
      logger.info("Loaded gazetteer from {}", gazetteerFile);
      logger.info("Loaded bib model from {}", bibModelFile);
  
      // Close out the model loader thread and make sure the results are OK.
      modelLoaderThread.join();
      if(exceptionThrownByModelLoaderThread.get() != null)
        throw exceptionThrownByModelLoaderThread.get();
      assert(model != null);
    }
  
    public Parser(
            final InputStream modelStream,
            final InputStream gazetteerStream,
            final InputStream bibModelStream
    ) throws Exception {
      final DataInputStream dis = new DataInputStream(modelStream);
      model = loadModel(dis);
      referenceExtractor =
              new ExtractReferences(
                      gazetteerStream,
                      new DataInputStream(bibModelStream));
    }
  
    public static Pair<List<BibRecord>, List<CitationRecord>> getReferences(
      final List<String> raw,
      final List<String> rawReferences,
      final ExtractReferences er
    ) {
      final Pair<List<BibRecord>, ExtractReferences.BibStractor> fnd = er.findReferences(rawReferences);
      final List<BibRecord> brs =
              fnd.getOne().stream().map(BibRecord::withNormalizedAuthors).collect(Collectors.toList());
      final ExtractReferences.BibStractor bs = fnd.getTwo();
      final List<CitationRecord> crs = ExtractReferences.findCitations(raw, brs, bs);
      return Tuples.pair(brs, crs);
    }
  
    //slow
    public static String paperToString(File f) {
      try {
        FileInputStream fis = new FileInputStream(f);
        PDFDoc doc = (new PDFExtractor()).extractFromInputStream(fis);
        fis.close();
        val seq = PDFToCRFInput.getSequence(doc);
        return PDFToCRFInput.stringAt(seq, Tuples.pair(0, seq.size()));
      } catch (Exception e) {
        return null;
      }
    }
  
    //from conll.Trainer:
    private static <T> Pair<List<T>, List<T>> splitData(List<T> original, double splitForSecond) {
      List<T> first = new ArrayList<>();
      List<T> second = new ArrayList<>();
      if (splitForSecond > 0.0) {
        Collections.shuffle(original, new Random(0L));
        int numFirst = (int) ((1.0 - splitForSecond) * original.size());
        first.addAll(original.subList(0, numFirst));
        second.addAll(original.subList(numFirst, original.size()));
      } else {
        first.addAll(original);
        // second stays empty
      }
      return Tuples.pair(first, second);
    }
  
    public static List<Pair<PaperToken, String>> getPaperLabels(
        final String paperId,
        final InputStream is,
        final LabeledData labeledData,
        PDFExtractor ext,
        int headerMax,
        boolean checkAuthors
    ) throws IOException {
      logger.debug("{}: starting", paperId);
  
      final PDFDoc doc;
      try {
        doc = ext.extractFromInputStream(is);
      } catch(final Exception e) {
        logger.warn("{} failed: {}", paperId, e.toString());
        return null;
      }
      if (doc == null) {
        return null;
      }
  
      List<PaperToken> seq = PDFToCRFInput.getSequence(doc);
      if (seq.size() == 0)
        return null;
      seq = seq.subList(0, Math.min(seq.size(), headerMax));
      List<Pair<PaperToken, String>> labeledPaper =
          PDFToCRFInput.labelMetadata(paperId, seq, labeledData);
      if (labeledPaper != null && checkAuthors) {
        final ExtractedMetadata checkEM =
                new ExtractedMetadata(
                        labeledPaper.stream().map(Pair::getOne).collect(Collectors.toList()),
                        labeledPaper.stream().map(Pair::getTwo).collect(Collectors.toList()));
        final Collection<String> expectedAuthorNames =
            labeledData.javaAuthorNames().orElse(Collections.emptyList());
        if(checkEM.authors.size() != expectedAuthorNames.size()) {
          logger.debug(
              "{}: author mismatch, discarding. Expected {}, got {}.",
              paperId,
              expectedAuthorNames,
              checkEM.authors);
          labeledPaper = null;
        }
      }
  
      return labeledPaper;
    }
  
    @Data
    private static class LabelingOutput {
      public final List<List<Pair<PaperToken, String>>> labeledData;
      public final Set<String> usedPaperIds;
    }
  
    private static LabelingOutput labelFromGroundTruth(
      final Iterator<LabeledPaper> labeledPapers,
      final int headerMax,
      final int maxFiles,
      final int minYear,
      final boolean checkAuthors,
      final Set<String> excludeIDs
    ) throws IOException {
      final PDFExtractor ext = new PDFExtractor();
  
      // input we need
      final int parallelism = Runtime.getRuntime().availableProcessors() * 2;
      final int queueSize = parallelism * 4;
      final Queue<Future<List<Pair<PaperToken, String>>>> workQueue = new ArrayDeque<>(queueSize);
  
      // stuff we need while processing
      final ExecutorService executor = Executors.newFixedThreadPool(parallelism);
      final AtomicInteger triedCount = new AtomicInteger();
      final ConcurrentMap<String, Long> paperId2startTime = new ConcurrentSkipListMap<>();
  
      // capturing results
      final ArrayList<List<Pair<PaperToken, String>>> results =
          new ArrayList<>(maxFiles > 0 ? maxFiles : 1024);
      final MutableSet<String> usedPaperIds = new UnifiedSet<String>().asSynchronized();
  
      try {
        if(maxFiles > 0)
          logger.info("Will be labeling {} papers in {} threads.", maxFiles, parallelism);
        else
          logger.info("Will be labeling all papers in {} threads.", parallelism);
  
        while (labeledPapers.hasNext() && (maxFiles <= 0 || results.size() < maxFiles)) {
          // fill up the queue
          while (labeledPapers.hasNext() && workQueue.size() < queueSize) {
            final LabeledPaper paper = labeledPapers.next();
            final LabeledData labels = paper.labels();
  
            val title = labels.javaTitle();
            val authors = labels.javaAuthors();
  
            if (excludeIDs.contains(paper.paperId()))
              continue;
  
            if(minYear > 0) {
              val year = labels.javaYear();
              if(!year.isPresent() || year.getAsInt() < minYear)
                continue;
            }
  
            final Future<List<Pair<PaperToken, String>>> future = executor.submit(() -> {
              try {
                paperId2startTime.put(paper.paperId(), System.currentTimeMillis());
                final List<Pair<PaperToken, String>> result;
                try(final InputStream is = paper.inputStream()) {
                  result = getPaperLabels(
                      paper.paperId(),
                      is,
                      labels,
                      ext,
                      headerMax,
                      checkAuthors);
                }
  
                int tried = triedCount.incrementAndGet();
                if (result != null)
                  usedPaperIds.add(paper.paperId());
                if (tried % 100 == 0) {
                  // find the document that's been in flight the longest
                  final long now = System.currentTimeMillis();
                  final Set<Map.Entry<String, Long>> entries = paperId2startTime.entrySet();
                  final Optional<Map.Entry<String, Long>> maxEntry =
                          entries.stream().max(
                                  new Comparator<Map.Entry<String, Long>>() {
                                    @Override
                                    public int compare(
                                            final Map.Entry<String, Long> left,
                                            final Map.Entry<String, Long> right
                                    ) {
                                      final long l = now - left.getValue();
                                      final long r = now - right.getValue();
                                      if (l < r)
                                        return -1;
                                      else if (l > r)
                                        return 1;
                                      else
                                        return 0;
                                    }
                                  });
  
                  final int succeeded = usedPaperIds.size();
                  if (maxEntry.isPresent()) {
                    logger.info(String.format(
                            "Tried to label %d papers, succeeded %d times (%.2f%%), %d papers in flight, oldest is %s at %.2f seconds",
                            tried,
                            succeeded,
                            succeeded * 100.0 / (double) tried,
                            entries.size(),
                            maxEntry.get().getKey(),
                            (now - maxEntry.get().getValue()) / 1000.0));
                  } else {
                    logger.info(String.format(
                            "Tried to label %d papers, succeeded %d times (%.2f%%), 0 papers in flight",
                            tried,
                            succeeded,
                            succeeded * 100.0 / (double) tried));
                  }
                }
  
                return result;
              } catch (final IOException e) {
                logger.warn(
                    "IOException {} while processing paper {}",
                    e.getMessage(),
                    paper.paperId());
                return null;
              } finally {
                paperId2startTime.remove(paper.paperId());
              }
            });
            workQueue.add(future);
          }
  
          // read from the queue
          val top = workQueue.poll();
          try {
            val result = top.get();
            if(result != null)
              results.add(result);
          } catch (final InterruptedException e) {
            logger.warn("Interrupted while processing paper", e);
          } catch (final ExecutionException e) {
            logger.warn("ExecutionException while processing paper", e);
          }
        }
  
        // read the queue empty if we have to
        while(!workQueue.isEmpty() && (maxFiles <= 0 || results.size() < maxFiles)) {
          try {
            val result = workQueue.poll().get();
            if(result != null)
              results.add(result);
          } catch (final InterruptedException e) {
            logger.warn("Interrupted while processing paper", e);
          } catch (final ExecutionException e) {
            logger.warn("ExecutionException while processing paper", e);
          }
        }
      } finally {
        executor.shutdown();
        try {
          logger.info("Done labeling papers. Waiting for threads to shut down.");
          executor.awaitTermination(10, TimeUnit.MINUTES);
        } catch(final InterruptedException e) {
          logger.warn("Interrupted while waiting for the executor to shut down. We may be leaking threads.", e);
        }
      }
  
      logger.info(String.format(
              "Tried to label %d papers, succeeded %d times (%.2f%%), all done!",
              triedCount.get(),
              usedPaperIds.size(),
              usedPaperIds.size() * 100.0 / triedCount.doubleValue()));
  
      return new LabelingOutput(results, usedPaperIds.asUnmodifiable());
    }
  
    private static UnifiedSet<String> readSet(String inFile) throws IOException {
      val out = new UnifiedSet<String>();
      BufferedReader brIn = new BufferedReader(new FileReader(inFile));
      String sLine;
      while ((sLine = brIn.readLine()) != null) {
        out.add(sLine);
      }
      brIn.close();
      return out;
    }
  
    public static void trainBibliographyCRF(
            final File bibsGroundTruthDirectory,
            final ParseOpts opts
    ) throws IOException {
      final File coraTrainFile = new File(bibsGroundTruthDirectory, "cora-citations.txt");
      final File umassTrainFile = new File(bibsGroundTruthDirectory, "umass-citations.txt");
      final File kermitTrainFile = new File(bibsGroundTruthDirectory, "kermit-citations.txt");
      trainBibliographyCRF(coraTrainFile, umassTrainFile, kermitTrainFile, opts);
    }
    
    public static void trainBibliographyCRF(File coraTrainFile, File umassTrainFile, File kermitTrainFile, ParseOpts opts) throws IOException {
      List<List<Pair<String, String>>> labeledData;
      labeledData = CRFBibRecordParser.labelFromCoraFile(coraTrainFile);
      if(umassTrainFile != null)
        labeledData.addAll(CRFBibRecordParser.labelFromUMassFile(umassTrainFile));
      if(kermitTrainFile != null)
        labeledData.addAll(CRFBibRecordParser.labelFromKermitFile(kermitTrainFile));
      ReferencesPredicateExtractor predExtractor;
      
      GazetteerFeatures gf = null;
      try {
      if(opts.gazetteerDir != null)
        gf = new GazetteerFeatures(opts.gazetteerDir);
      }
      catch (IOException e) {
        logger.error("Error importing gazetteer directory, ignoring.");
      }
      ParserLMFeatures plf = null;
      if(opts.gazetteerFile != null) {
        ParserGroundTruth gaz = new ParserGroundTruth(opts.gazetteerFile);
        plf = new ParserLMFeatures(
            gaz.papers,
            new UnifiedSet<>(),
            new File(opts.backgroundDirectory),
            opts.backgroundSamples);
        predExtractor = new ReferencesPredicateExtractor(plf);
      } else {
        predExtractor = new ReferencesPredicateExtractor();
      }
      predExtractor.setGf(gf);
      
      // Split train/test data
      logger.info("CRF training for bibs with {} threads and {} labeled examples", opts.threads, labeledData.size());
      Pair<List<List<Pair<String, String>>>, List<List<Pair<String, String>>>> trainTestPair =
        splitData(labeledData, 1.0 - opts.trainFraction);
      val trainLabeledData = trainTestPair.getOne();
      val testLabeledData = trainTestPair.getTwo();
  
      // Set up Train options
      CRFTrainer.Opts trainOpts = new CRFTrainer.Opts();
      trainOpts.optimizerOpts.maxIters = opts.iterations;
      trainOpts.numThreads = opts.threads;
      trainOpts.minExpectedFeatureCount = opts.minExpectedFeatureCount;
  
      // set up evaluation function
      final Parallel.MROpts evalMrOpts =
              Parallel.MROpts.withIdAndThreads("mr-crf-bib-train-eval", opts.threads);
      final ToDoubleFunction<CRFModel<String, String, String>> testEvalFn;
      {
        final List<List<Pair<String, String>>> testEvalData =
                testLabeledData.
                        stream().
                        map(x -> x.stream().map(Pair::swap).collect(toList())).
                        collect(toList());
        testEvalFn = (model) -> {
          final Evaluation<String> eval = Evaluation.compute(model, testEvalData, evalMrOpts);
          return eval.tokenAccuracy.accuracy();
        };
      }
  
      // set up early stopping so that we stop training after 50 down-iterations
      final TrainCriterionEval<CRFModel<String, String, String>> earlyStoppingEvaluator =
              new TrainCriterionEval<>(testEvalFn);
      earlyStoppingEvaluator.maxNumDipIters = 100;
      trainOpts.iterCallback = earlyStoppingEvaluator;
  
      // training
      CRFTrainer<String, String, String> trainer =
        new CRFTrainer<>(trainLabeledData, predExtractor, trainOpts);
      trainer.train(trainLabeledData);
      final CRFModel<String, String, String> crfModel = earlyStoppingEvaluator.bestModel;
  
      Vector weights = crfModel.weights();
      Parallel.shutdownExecutor(evalMrOpts.executorService, Long.MAX_VALUE);
  
      val dos = new DataOutputStream(new FileOutputStream(opts.modelFile));
      logger.info("Writing model to {}", opts.modelFile);
      saveModel(dos, crfModel.featureEncoder, weights, plf, gf, ExtractReferences.DATA_VERSION);
      dos.close();
    }
  
    public static void trainParser(
        final Iterator<LabeledPaper> labeledTrainingData,
        final ParseOpts opts
    ) throws IOException {
      trainParser(labeledTrainingData, opts, UnifiedSet.newSet());
    }
  
    public static void trainParser(
        final Iterator<LabeledPaper> labeledTrainingData,
        final ParseOpts opts,
        final String excludeIDsFile
    ) throws IOException {
      final UnifiedSet<String> excludedIDs;
      if (excludeIDsFile == null)
        excludedIDs = new UnifiedSet<>();
      else
        excludedIDs = readSet(excludeIDsFile);
      trainParser(labeledTrainingData, opts, excludedIDs);
    }
  
    public static void trainParser(
        Iterator<LabeledPaper> labeledTrainingData,
        ParseOpts opts,
        final UnifiedSet<String> excludeIDs
    ) throws IOException {
      final LabelingOutput labelingOutput = labelFromGroundTruth(
          labeledTrainingData,
          opts.headerMax,
          opts.documentCount,
          opts.minYear,
          opts.checkAuthors,
          excludeIDs);
      ParserLMFeatures plf = null;
      final PDFPredicateExtractor predExtractor;
      if (opts.gazetteerFile != null) {
        ParserGroundTruth gaz = new ParserGroundTruth(opts.gazetteerFile);
        UnifiedSet<String> excludedIds = new UnifiedSet<String>(labelingOutput.usedPaperIds);
        excludedIds.addAll(excludeIDs);
        plf = new ParserLMFeatures(
                gaz.papers,
                excludedIds,
                new File(opts.backgroundDirectory),
                opts.backgroundSamples);
        predExtractor = new PDFPredicateExtractor(plf);
      } else {
        predExtractor = new PDFPredicateExtractor();
      }
  
      // Split train/test data
      logger.info(
              "CRF training with {} threads and {} labeled examples",
              opts.threads,
              labelingOutput.labeledData.size());
      Pair<List<List<Pair<PaperToken, String>>>, List<List<Pair<PaperToken, String>>>> trainTestPair =
        splitData(labelingOutput.labeledData, 1.0 - opts.trainFraction);
      val trainLabeledData = trainTestPair.getOne();
      val testLabeledData = trainTestPair.getTwo();
  
      // log test and training data
      if(labeledDataLogger.isDebugEnabled()) {
        labeledDataLogger.info("Training data before:");
        logLabeledData(trainLabeledData);
  
        labeledDataLogger.info("Test data before:");
        logLabeledData(testLabeledData);
      }
  
      // Set up Train options
      CRFTrainer.Opts trainOpts = new CRFTrainer.Opts();
      trainOpts.optimizerOpts.maxIters = opts.iterations;
      trainOpts.numThreads = opts.threads;
      trainOpts.minExpectedFeatureCount = opts.minExpectedFeatureCount;
  
      // set up evaluation function
      final Parallel.MROpts evalMrOpts =
              Parallel.MROpts.withIdAndThreads("mr-crf-train-eval", opts.threads);
      final ToDoubleFunction<CRFModel<String, PaperToken, String>> testEvalFn;
      {
        final List<List<Pair<String, PaperToken>>> testEvalData =
                testLabeledData.
                        stream().
                        map(x -> x.stream().map(Pair::swap).collect(toList())).
                        collect(toList());
        testEvalFn = (model) -> {
          final Evaluation<String> eval = Evaluation.compute(model, testEvalData, evalMrOpts);
          logger.info("Test Label F-measures");
          eval.stateFMeasures.forEach((label, fMeasure) -> {
            logger.info(String.format("-- %s: p:%.3f r:%.3f f1:%.3f",
                label, fMeasure.precision(), fMeasure.recall(), fMeasure.f1()));
          });
          logger.info("");
          return eval.tokenAccuracy.accuracy();
        };
      }
  
      // set up early stopping so that we stop training after 50 down-iterations
      final TrainCriterionEval<CRFModel<String, PaperToken, String>> earlyStoppingEvaluator =
              new TrainCriterionEval<>(testEvalFn);
      earlyStoppingEvaluator.maxNumDipIters = 100;
      trainOpts.iterCallback = earlyStoppingEvaluator;
  
      // training
      final CRFTrainer<String, PaperToken, String> trainer =
        new CRFTrainer<>(trainLabeledData, predExtractor, trainOpts);
      trainer.train(trainLabeledData);
      final CRFModel<String, PaperToken, String> crfModel = earlyStoppingEvaluator.bestModel;
  
      final Vector weights = crfModel.weights();
      Parallel.shutdownExecutor(evalMrOpts.executorService, Long.MAX_VALUE);
  
      try(val dos = new DataOutputStream(new FileOutputStream(opts.modelFile))) {
        logger.info("Writing model to {}", opts.modelFile);
        saveModel(dos, crfModel.featureEncoder, weights, plf);
      }
  
      // log test and training data
      if(labeledDataLogger.isDebugEnabled()) {
        labeledDataLogger.info("Training data after:");
        logLabeledData(trainLabeledData);
  
        labeledDataLogger.info("Test data after:");
        logLabeledData(testLabeledData);
      }
    }
  
    private static void logLabeledData(final List<List<Pair<PaperToken, String>>> data) {
      for(List<Pair<PaperToken, String>> line : data) {
        final String logLine = line.stream().map(pair ->
                String.format(
                        "%s/%x/%s",
                        pair.getOne().getPdfToken() == null ?
                                "null" :
                                pair.getOne().getPdfToken().getToken(),
                        pair.getOne().hashCode(),
                        pair.getTwo())).collect(Collectors.joining(" "));
        labeledDataLogger.info(logLine);
      }
    }
  
    public static <T> void saveModel(
      final DataOutputStream dos,
      final CRFFeatureEncoder<String, T, String> fe,
      final Vector weights,
      final ParserLMFeatures plf,
      final String dataVersion
    ) throws IOException {
      saveModel(dos, fe,weights, plf, null, dataVersion);
    }
    
    public static <T> void saveModel(
      final DataOutputStream dos,
      final CRFFeatureEncoder<String, T, String> fe,
      final Vector weights,
      final ParserLMFeatures plf,
      final GazetteerFeatures gf,
      final String dataVersion
    ) throws IOException {
      dos.writeUTF(dataVersion);
      fe.stateSpace.save(dos);
      fe.nodeFeatures.save(dos);
      fe.edgeFeatures.save(dos);
      IOUtils.saveDoubles(dos, weights.toDoubles());
  
      logger.debug("Saving ParserLMFeatures");
      try(final FSTObjectOutput out = new FSTObjectOutput(dos)) {
        out.writeObject(plf);
        if (plf != null)
          plf.logState();
  
        logger.debug("Saving gazetteer features");
        out.writeObject(gf);
      }
    }
  
    public static <T> void saveModel(
      final DataOutputStream dos,
      final CRFFeatureEncoder<String, T, String> fe,
      final Vector weights,
      final ParserLMFeatures plf
    ) throws IOException {
      saveModel(dos, fe, weights, plf, Parser.DATA_VERSION);
    }
  
    @Data
    public static class ModelComponents {
      public final PDFPredicateExtractor predExtractor;
      public final CRFFeatureEncoder<String, PaperToken, String> featureEncoder;
      public final CRFWeightsEncoder<String> weightsEncoder;
      public final CRFModel<String, PaperToken, String> model;
    }
  
    public static ModelComponents loadModelComponents(
      final DataInputStream dis,
      String dataVersion
    ) throws IOException {
      IOUtils.ensureVersionMatch(dis, dataVersion);
      val stateSpace = StateSpace.load(dis);
      Indexer<String> nodeFeatures = Indexer.load(dis);
      Indexer<String> edgeFeatures = Indexer.load(dis);
      Vector weights = DenseVector.of(IOUtils.loadDoubles(dis));
  
      logger.debug("Loading ParserLMFeatures");
      final ParserLMFeatures plf;
      try(final FSTObjectInput in = new FSTObjectInput(dis)) {
        try {
          plf = (ParserLMFeatures) in.readObject();
        } catch (final ClassNotFoundException e) {
          throw new IOException("Model file contains unknown class.", e);
        }
      }
      if(plf != null && logger.isDebugEnabled())
        plf.logState();
  
      val predExtractor = new PDFPredicateExtractor(plf);
      val featureEncoder = new CRFFeatureEncoder<String, PaperToken, String>
              (predExtractor, stateSpace, nodeFeatures, edgeFeatures);
      val weightsEncoder = new CRFWeightsEncoder<String>(stateSpace, nodeFeatures.size(), edgeFeatures.size());
      val model = new CRFModel<String, PaperToken, String>(featureEncoder, weightsEncoder, weights);
      return new ModelComponents(predExtractor, featureEncoder, weightsEncoder, model);
    }
  
    public static ModelComponents loadModelComponents(
      final DataInputStream dis
    ) throws IOException {
      return loadModelComponents(dis, DATA_VERSION);
    }
  
    public static CRFModel<String, PaperToken, String> loadModel(
      final DataInputStream dis
    ) throws IOException {
      return loadModelComponents(dis).model;
    }
  
    public static void clean(ExtractedMetadata em) {
      em.title = cleanTitle(em.title);
      em.authors = trimAuthors(em.authors);
    }
  
    public static String cleanTitle(String t) {
      if (t == null || t.length() == 0)
        return t;
      // strip accents and unicode changes
      t = Normalizer.normalize(t, Normalizer.Form.NFKD);
      // kill non-character letters
      // kill xml
      t = t.replaceAll("\\&.*?\\;", "");
      if (t.endsWith("."))
        t = t.substring(0, t.length() - 1);
      return t;
    }
  
    public static String processTitle(String t) {
      // case fold and remove lead/trail space
      if(t==null || t.length()==0)
        return t;
      t = t.trim().toLowerCase();
      t = cleanTitle(t);
      // kill non-letter chars
      t = t.replaceAll("\\W", "");
      return t.replaceAll("\\s+", " ");
    }
  
    //changes extraction to remove common failure modes
    public static String processExtractedTitle(String t) {
      String out = t.replaceAll("(?<=[a-z])\\- ", ""); //continuation dash
      if (!out.endsWith("?") && !out.endsWith("\"") && !out.endsWith(")"))
        out = out.replaceFirst("\\W$", ""); //end of title punctuation if not ?, ", or )
      return out.trim();
    }
  
    //lowercases
    public static String lastName(String s) {
      String[] words = s.split(" ");
      if (words.length > 0)
        return processTitle(words[words.length - 1]);
      else
        return "";
    }
  
    public static List<String> lastNames(List<String> ss) {
      return (ss==null)?null:ss.stream().map(s -> lastName(s)).collect(Collectors.toList());
    }
  
    public static int scoreAuthors(String[] expected, List<String> guessed) {
      List<String> guessLastName = lastNames(guessed);
      List<String> expectedLastName = lastNames(Arrays.asList(expected));
      //slow:
      int truePos = 0;
      for (String s : expectedLastName) {
        if (guessLastName.contains(s))
          truePos++;
      }
      return truePos;
    }
  
    /** Fixes common author problems. This is applied to the output of both normalized and
     * unnormalized author names, and it is used in training as well. Experience shows that if you
     * make changes here, there is a good chance you'll need to retrain, even if you think the change
     * is fairly trivial. */
    public static String fixupAuthors(String s) {
      // delete trailing special characters
      String sFix = s.replaceAll("([^\\p{javaLowerCase}\\p{javaUpperCase}])+$", "");
      if (sFix.contains(","))
        sFix = sFix.substring(0, sFix.indexOf(","));
      if (sFix.endsWith("Jr"))
        sFix = sFix + ".";
      return sFix;
    }
  
    private final static int MAX_AUTHOR_LENGTH = 32;
    private final static int MIN_AUTHOR_LENGTH = 2;
    public static List<String> trimAuthors(List<String> auth) {
      return auth.stream().
              flatMap(s -> Arrays.stream(s.split("(?!,\\s*Jr),|\\band\\b"))).
              map(String::trim).
              map(Parser::fixupAuthors).
              filter(s -> !s.isEmpty()).
              filter(s -> s.length() <= MAX_AUTHOR_LENGTH).
              filter(s -> s.length() >= MIN_AUTHOR_LENGTH).
              distinct().
              collect(Collectors.toList());
    }
  
    public static class ParsingTimeout extends RuntimeException { }
    private final Timer parserKillerTimer = new Timer("Science-parse killer timer", true);
    private final MutableIntSet parseNumbersInProgress = IntSets.mutable.empty();
    private final AtomicInteger nextParseNumber = new AtomicInteger();
    public ExtractedMetadata doParseWithTimeout(final InputStream is, final long timeoutInMs) throws IOException {
      final int parseNumber = nextParseNumber.getAndIncrement();
  
      final Thread t = Thread.currentThread();
  
      final TimerTask killTaskHard = new TimerTask() {
        @Override
        public void run() {
          synchronized (parseNumbersInProgress) {
            if(parseNumbersInProgress.contains(parseNumber)) {
              logger.info("Killing parsing thread {} because it's taking too long", t.getId());
              t.stop();
              // I know this is dangerous. This is a last resort.
            }
          }
        }
      };
  
      final TimerTask killTaskSoftly = new TimerTask() {
        @Override
        public void run() {
          synchronized (parseNumbersInProgress) {
            if(parseNumbersInProgress.contains(parseNumber)) {
              logger.info("Interrupting parsing thread {} because it's taking too long", t.getId());
              t.interrupt();
            }
          }
        }
      };
  
      synchronized (parseNumbersInProgress) {
        parseNumbersInProgress.add(parseNumber);
      }
  
      final ExtractedMetadata result;
      final boolean wasInterrupted;
      parserKillerTimer.schedule(killTaskSoftly, timeoutInMs);
      parserKillerTimer.schedule(killTaskHard, 3*timeoutInMs);
      try {
        try {
          result = doParse(is);
        } catch(final ThreadDeath e) {
          throw new RuntimeException("Science-parse killer got impatient", e);
        }
      } finally {
        synchronized (parseNumbersInProgress) {
          parseNumbersInProgress.remove(parseNumber);
        }
  
        // This clears the interrupted flag, in case it happened after we were already done parsing,
        // but before we could remove the parse number. We don't want to leave this function with
        // the interrupted flag set on the thread.
        // Actually, the window of opportunity is from the last time that doParse() checks the
        // flag to the time we remove the parse number, which is quite a bit bigger.
        wasInterrupted = Thread.interrupted();
      }
  
      if(wasInterrupted)
        logger.info("Overriding interruption of parsing thread {} because it finished before we could react", t.getId());
      assert !Thread.interrupted();
  
      return result;
    }
  
    public ExtractedMetadata doParse(final InputStream is) throws IOException {
      return doParse(is, MAXHEADERWORDS);
    }
  
    /**
     * Given a body of text (e.g. an entire section of a paper) within which a citation is mentioned, extracts a
     * single-sentence context from that larger body of text. Used both here for Science Parse extraction and in
     * GrobidParser for evaluation of Grobid citation mention extraction.
     */
    public static CitationRecord extractContext(int referenceID, String context, int begin, int end) {
      int sentenceStart = begin;
      // Some citation have the form "This is my sentence.(1,2)". So if we just search backwards from
      // `begin`, we find the end of the sentence we want, rather than the beginning. Thus, we rewind
      // `sentenceStart` a little if necessary.
      while(
          sentenceStart > 0 && (
              context.charAt(sentenceStart) == '(' ||
              context.charAt(sentenceStart) == '[' ||
              context.charAt(sentenceStart) == '.' ||
              context.charAt(sentenceStart) == '⍐'
          )
      ) {
        sentenceStart -= 1;
      }
      sentenceStart = context.substring(0, sentenceStart).lastIndexOf('.') + 1; // this evaluates to 0 if '.' is not found
  
      // Trim away superscripts at the beginning of sentences.
      if(context.charAt(sentenceStart) == '⍐') {
        final int newSentenceStart = context.indexOf('⍗', sentenceStart);
        if(newSentenceStart > 0 && newSentenceStart < begin)
          sentenceStart = newSentenceStart + 1;
      }
  
      int crSentenceEnd = context.indexOf('.', end);
      if(crSentenceEnd < 0)
        crSentenceEnd = context.length();
      else
        crSentenceEnd += 1;
  
      String contextSentenceUntrimmed = context.substring(sentenceStart, crSentenceEnd);
      String contextSentence = contextSentenceUntrimmed.trim();
      sentenceStart += contextSentenceUntrimmed.indexOf(contextSentence);
      return new CitationRecord(referenceID, contextSentence, begin - sentenceStart,
              end - sentenceStart);
    }
  
    public ExtractedMetadata doParse(final InputStream is, int headerMax) throws IOException {
      final ExtractedMetadata em;
      final PDDocument pdDoc = PDDocument.load(is);
  
      //
      // Run Science-parse
      //
      {
        PDFExtractor ext = new PDFExtractor();
        final PDFDoc doc = ext.extractResultFromPDDocument(pdDoc).document;
        final PDFDoc docWithoutSuperscripts = doc.withoutSuperscripts();
  
        List<PaperToken> seq = PDFToCRFInput.getSequence(docWithoutSuperscripts);
        seq = seq.subList(0, Math.min(seq.size(), headerMax));
        seq = PDFToCRFInput.padSequence(seq);
  
        { // get title and authors from the CRF
          List<String> outSeq = model.bestGuess(seq);
          //the output tag sequence will not include the start/stop states!
          outSeq = PDFToCRFInput.padTagSequence(outSeq);
          em = new ExtractedMetadata(seq, outSeq);
          em.source = ExtractedMetadata.Source.CRF;
        }
  
        // use PDF metadata if it's there
        if (doc.meta != null) {
          if (doc.meta.title != null) {
            em.setTitle(doc.meta.title);
            em.source = ExtractedMetadata.Source.META;
          }
          if (doc.meta.createDate != null)
            em.setYearFromDate(doc.meta.createDate);
        }
  
        clean(em);
  
        if(doc.meta != null)
          em.creator = doc.meta.creator;
  
        // extract references
        try {
          final List<String> lines = PDFDocToPartitionedText.getRaw(doc);
          final List<String> rawReferences = PDFDocToPartitionedText.getRawReferences(doc);
          final Pair<List<BibRecord>, List<CitationRecord>> pair =
              getReferences(lines, rawReferences, referenceExtractor);
          em.references = new ArrayList<>(pair.getOne().size());
          for(final BibRecord record : pair.getOne())
            em.references.add(record.withoutSuperscripts());
  
          // add contexts to the mentions
          List<CitationRecord> crs = new ArrayList<>();
          for(final CitationRecord cr : pair.getTwo()) {
            final CitationRecord crWithContext =
                extractContext(cr.referenceID, cr.context, cr.startOffset, cr.endOffset);
            final int contextLength =
                crWithContext.context.length() -
                (crWithContext.endOffset - crWithContext.startOffset);
            if(contextLength >= 35) // Heuristic number
              crs.add(crWithContext);
          }
  
          // find the predominant mention style, and assign it to em.referenceMentions
          final Function<CitationRecord, Character> getMentionStyle = (final CitationRecord cr) -> {
            char firstChar = cr.context.charAt(cr.startOffset);
            if(firstChar == '(' || firstChar == '[' || firstChar == '⍐')
              return firstChar;
            else
              return '\0';
          };
          final MutableCharIntMap styleToCount = CharIntMaps.mutable.empty();
          int maxStyleCount = 0;
          char predominantStyle = '\0';
          for(final CitationRecord cr : crs) {
            final char style = getMentionStyle.apply(cr);
            final int newStyleCount = styleToCount.addToValue(style, 1);
            if(newStyleCount > maxStyleCount) {
              maxStyleCount = newStyleCount;
              predominantStyle = style;
            }
          }
  
          // Override this in a special case: If we have more than 4 in the [] style, let's take that
          // as the style. That style is unlikely to happen by accident, while the () style happens
          // all the time in other contexts.
          if(predominantStyle == '(' && styleToCount.getIfAbsent('[', 0) > 4)
            predominantStyle = '[';
  
          em.referenceMentions = new ArrayList<>(crs.size());
          for(final CitationRecord cr : crs)
            if(predominantStyle == '\0' || getMentionStyle.apply(cr) == predominantStyle)
              em.referenceMentions.add(cr.withConvertedSuperscriptTags());
        } catch (final RegexWithTimeout.RegexTimeout|Parser.ParsingTimeout e) {
          logger.warn("Timeout while extracting references. References may be incomplete or missing.");
          if (em.references == null)
            em.references = Collections.emptyList();
          if (em.referenceMentions == null)
            em.referenceMentions = Collections.emptyList();
        }
        logger.debug(em.references.size() + " refs for " + em.title);
  
        try {
          final List<String> lines = PDFDocToPartitionedText.getRaw(docWithoutSuperscripts);
          // Fix-up of lines that should not be necessary, but is
          for(int i = 0; i < lines.size(); ++i) {
            final String s = lines.get(i).replaceAll("-<lb>", "").replaceAll("<lb>", " ");
            lines.set(i, s);
          }
  
          em.abstractText =
              PDFDocToPartitionedText.getAbstract(lines, docWithoutSuperscripts).trim();
          em.abstractText = em.abstractText.replaceAll("⍐[^⍗]⍗", "");
          if (em.abstractText.isEmpty())
            em.abstractText = null;
        } catch (final RegexWithTimeout.RegexTimeout|Parser.ParsingTimeout e) {
          logger.warn("Timeout while extracting abstract. Abstract will be missing.");
          em.abstractText = null;
        }
      }
  
      //
      // Run figure extraction to get sections
      //
      try {
        final FigureExtractor fe = new FigureExtractor(false, true, true, true, true);
  
        final FigureExtractor.Document doc =
            fe.getFiguresWithText(pdDoc, scala.Option.apply(null), scala.Option.apply(null));
  
        em.sections = ScalaStreamSupport.stream(doc.sections()).map(documentSection ->
            new Section(
                OptionConverters.toJava(documentSection.titleText()).orElse(null),
                documentSection.bodyText()
            )
        ).filter(documentSection ->
            // filter out reference sections
            !(
                documentSection.getHeading() != null &&
                PDFDocToPartitionedText.referenceHeaders.contains(
                  documentSection.getHeading().trim().toLowerCase().replaceAll("\\p{Punct}*$", ""))
            )
        ).collect(Collectors.toList());
      } catch (final Exception e) {
        logger.warn(
            "Exception {} while getting sections. Section data will be missing.",
            e.getMessage());
        em.sections = null;
      }
  
      return em;
    }
  
    public static class ParseOpts {
      public String modelFile;
      public int iterations;
      public int threads;
      public int headerMax;
      public double trainFraction;
      public String gazetteerFile; //record of references
      public String gazetteerDir; //directory of entity lists (universities, person names, etc.)
      public int backgroundSamples;
      public String backgroundDirectory;
      public int minYear; //only process papers this year or later
      public boolean checkAuthors; //only bootstraps papers if all authors are found
      public int documentCount = -1; // how many documents to train on. set to -1 to train on all.
      public int minExpectedFeatureCount = 1;
    }
  }
  