diff --git a/grobid-trainer/src/main/java/org/grobid/trainer/FulltextTrainer.java b/grobid-trainer/src/main/java/org/grobid/trainer/FulltextTrainer.java index 91925769cd..ca2c2c2efb 100755 --- a/grobid-trainer/src/main/java/org/grobid/trainer/FulltextTrainer.java +++ b/grobid-trainer/src/main/java/org/grobid/trainer/FulltextTrainer.java @@ -22,7 +22,7 @@ public FulltextTrainer() { } public FulltextTrainer(GrobidModels.Flavor modelFlavor) { - super(GrobidModels.getModelFlavor(GrobidModels.SEGMENTATION, modelFlavor)); + super(GrobidModels.getModelFlavor(GrobidModels.FULLTEXT, modelFlavor)); flavor = modelFlavor; } @@ -251,16 +251,28 @@ public boolean accept(File dir, String name) { return totalExamples; } - /** - * Command line execution. - * - * @param args Command line arguments. - * @throws Exception - */ + public static void main(String[] args) throws Exception { - GrobidProperties.getInstance(); - AbstractTrainer.runTraining(new FulltextTrainer()); - System.out.println(AbstractTrainer.runEvaluation(new FulltextTrainer())); + // if we have a parameter, it gives the flavor refinement to consider + GrobidModels.Flavor theFlavor = null; + if (args.length > 0) { + String flavor = args[0]; + theFlavor = GrobidModels.Flavor.fromLabel(flavor); + if (theFlavor == null) { + System.out.println("Warning, the flavor is not recognized, " + + "must one one of [article/light, article/light-ref, sdo/ietf], " + + "defaulting training with no flavor..."); + } + } + + GrobidProperties.getInstance(); + if (theFlavor == null) { + AbstractTrainer.runTraining(new FulltextTrainer()); + System.out.println(AbstractTrainer.runEvaluation(new FulltextTrainer())); + } else { + AbstractTrainer.runTraining(new FulltextTrainer(theFlavor)); + System.out.println(AbstractTrainer.runEvaluation(new FulltextTrainer(theFlavor))); + } System.exit(0); } -} \ No newline at end of file +} \ No newline at end of file