add algorithm option for wapiti

kermitt2 · Feb 3, 2021 · ac4dc2c · ac4dc2c
1 parent bfc10f7
commit ac4dc2c
Show file tree

Hide file tree

Showing 6 changed files with 41 additions and 10 deletions.
diff --git a/grobid-trainer/src/main/java/org/grobid/trainer/AbstractTrainer.java b/grobid-trainer/src/main/java/org/grobid/trainer/AbstractTrainer.java
@@ -49,6 +49,7 @@ public abstract class AbstractTrainer implements Trainer {
     protected double epsilon = 0.0; // size of the interval for stopping criterion
     protected int window = 0; // similar to CRF++
     protected int nbMaxIterations = 0; // maximum number of iterations in training
+    protected String algorithm = ""; // algorithm
 
     protected GrobidModel model;
     private File trainDataPath;
@@ -94,6 +95,9 @@ public void train() {
         if (nbMaxIterations != 0)
             trainer.setNbMaxIterations(nbMaxIterations);
 
+        if (StringUtils.isNotBlank(algorithm))
+            trainer.setAlgorithm(algorithm);
+
         File dirModelPath = new File(GrobidProperties.getModelPath(model).getAbsolutePath()).getParentFile();
         if (!dirModelPath.exists()) {
             LOGGER.warn("Cannot find the destination directory " + dirModelPath.getAbsolutePath() + " for the model " + model.getModelName() + ". Creating it.");
@@ -151,6 +155,8 @@ public String splitTrainEvaluate(Double split) {
             trainer.setWindow(window);
         if (nbMaxIterations != 0)
             trainer.setNbMaxIterations(nbMaxIterations);
+        if (StringUtils.isNotBlank(algorithm))
+            trainer.setAlgorithm(algorithm);
 
         File dirModelPath = new File(GrobidProperties.getModelPath(model).getAbsolutePath()).getParentFile();
         if (!dirModelPath.exists()) {
@@ -197,6 +203,8 @@ public String nFoldEvaluate(int numFolds, boolean includeRawResults) {
             trainer.setWindow(window);
         if (nbMaxIterations != 0)
             trainer.setNbMaxIterations(nbMaxIterations);
+        if (StringUtils.isNotBlank(algorithm))
+            trainer.setAlgorithm(algorithm);
 
         //We dump the model in the tmp directory
         File tmpDirectory = new File(GrobidProperties.getTempPath().getAbsolutePath());

diff --git a/grobid-trainer/src/main/java/org/grobid/trainer/CRFPPGenericTrainer.java b/grobid-trainer/src/main/java/org/grobid/trainer/CRFPPGenericTrainer.java
@@ -26,6 +26,7 @@ public class CRFPPGenericTrainer implements GenericTrainer {
 	protected double epsilon = 0.00001; // default size of the interval for stopping criterion
 	protected int window = 20; // default similar to CRF++
     protected int nbMaxIterations = 6000;
+    protected String algorithm = "crf-l2";
 
     public CRFPPGenericTrainer() {
         crfppTrainer = new CRFPPTrainer();
@@ -70,7 +71,12 @@ public int getWindow() {
     public void setNbMaxIterations(int interations) {
         this.nbMaxIterations = interations;
     }
-
+
+    @Override
+    public void setAlgorithm(String algorithm) {
+        this.algorithm = algorithm;
+    }
+
     @Override
     public int getNbMaxIterations() {
         return nbMaxIterations;

diff --git a/grobid-trainer/src/main/java/org/grobid/trainer/DeLFTTrainer.java b/grobid-trainer/src/main/java/org/grobid/trainer/DeLFTTrainer.java
@@ -50,7 +50,11 @@ public int getWindow() {
     @Override
     public void setNbMaxIterations(int interations) {
     }
-
+
+    @Override
+    public void setAlgorithm(String algorithm) {
+    }
+
     @Override
     public int getNbMaxIterations() {
         return 0;

diff --git a/grobid-trainer/src/main/java/org/grobid/trainer/DummyTrainer.java b/grobid-trainer/src/main/java/org/grobid/trainer/DummyTrainer.java
@@ -45,6 +45,9 @@ public int getNbMaxIterations() {
 
     @Override
     public void setNbMaxIterations(int iterations) {
+    }
 
+    @Override
+    public void setAlgorithm(String algorithm) {
     }
 }
diff --git a/grobid-trainer/src/main/java/org/grobid/trainer/GenericTrainer.java b/grobid-trainer/src/main/java/org/grobid/trainer/GenericTrainer.java
@@ -11,10 +11,11 @@
 public interface GenericTrainer {
     void train(File template, File trainingData, File outputModel, int numThreads, GrobidModel model);
     String getName();
-	public void setEpsilon(double epsilon);
-	public void setWindow(int window);
-	public double getEpsilon();
-	public int getWindow();
-	public int getNbMaxIterations();
-	public void setNbMaxIterations(int iterations);
+	void setEpsilon(double epsilon);
+	void setWindow(int window);
+	double getEpsilon();
+	int getWindow();
+	int getNbMaxIterations();
+	void setNbMaxIterations(int iterations);
+    void setAlgorithm(String algorithm);
 }
diff --git a/grobid-trainer/src/main/java/org/grobid/trainer/WapitiTrainer.java b/grobid-trainer/src/main/java/org/grobid/trainer/WapitiTrainer.java
@@ -20,15 +20,19 @@ public class WapitiTrainer implements GenericTrainer {
 	protected double epsilon = 0.00001; // default size of the interval for stopping criterion
 	protected int window = 20; // default similar to CRF++
     protected int nbMaxIterations = 2000; // by default maximum of training iterations
+    protected String algorithm = "l-bfgs"; // algorithm to be used, values: l-bfgs (default), sgd-l1, bcd, rprop, rprop+, rprop-
+
 
     @Override
     public void train(File template, File trainingData, File outputModel, int numThreads, GrobidModel model) {
 		System.out.println("\tepsilon: " + epsilon);
 		System.out.println("\twindow: " + window);
         System.out.println("\tnb max iterations: " + nbMaxIterations);
 		System.out.println("\tnb threads: " + numThreads);
+		System.out.println("\talgorithm: " + algorithm);
+
         WapitiModel.train(template, trainingData, outputModel, "--nthread " + numThreads +
-//       		" --algo sgd-l1" +
+       		" --algo " + algorithm +
 			" -e " + BigDecimal.valueOf(epsilon).toPlainString() +
 			" -w " + window +
 			" -i " + nbMaxIterations
@@ -64,7 +68,12 @@ public int getWindow() {
     public void setNbMaxIterations(int interations) {
         this.nbMaxIterations = interations;
     }
-
+
+    @Override
+    public void setAlgorithm(String algorithm) {
+        this.algorithm = algorithm;
+    }
+
     @Override
     public int getNbMaxIterations() {
         return nbMaxIterations;