support training with transformer parameter

kermitt2 · Nov 21, 2022 · b9d92c6 · b9d92c6
1 parent cb39dac
commit b9d92c6
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 0 deletions.
diff --git a/grobid-core/src/main/java/org/grobid/core/jni/DeLFTModel.java b/grobid-core/src/main/java/org/grobid/core/jni/DeLFTModel.java
@@ -257,6 +257,11 @@ public void run() {
                     localArgs += ", batch_size="+
                         GrobidProperties.getInstance().getDelftTrainingBatchSize(this.modelName);
 
+                if (GrobidProperties.getInstance().getDelftTranformer(modelName) != null) {
+                    localArgs += ", transformer="+
+                        GrobidProperties.getInstance().getDelftTranformer(modelName);
+                }
+
                 // init model to be trained
                 if (architecture == null)
                     jep.eval("model = Sequence('"+this.modelName+
@@ -310,6 +315,10 @@ public static void train(String modelName, File trainingData, File outputModel,
                 command.add("--architecture");
                 command.add(architecture);
             }
+            if (GrobidProperties.getInstance().getDelftTranformer(modelName) != null) {
+                command.add("--transformer");
+                command.add(GrobidProperties.getInstance().getDelftTranformer(modelName));
+            }
             if (GrobidProperties.getInstance().useELMo(modelName) && modelName.toLowerCase().indexOf("bert") == -1) {
                 command.add("--use-ELMo");
             }

diff --git a/grobid-core/src/main/java/org/grobid/core/utilities/GrobidProperties.java b/grobid-core/src/main/java/org/grobid/core/utilities/GrobidProperties.java
@@ -791,6 +791,20 @@ public static String getDelftEmbeddingsName(final String modelName) {
         return param.delft.embeddings_name;
     }
 
+    public static String getDelftTranformer(final String modelName) {
+        ModelParameters param = modelMap.get(modelName);
+        if (param == null) {
+            LOGGER.debug("No configuration parameter defined for model " + modelName);
+            return null;
+        }
+        DelftModelParameters delftParam = param.delft;
+        if (delftParam == null) {
+            LOGGER.debug("No configuration parameter defined for DeLFT engine for model " + modelName);
+            return null;
+        }
+        return param.delft.transformer;
+    }
+
     /**
     *  Return -1 if not set in the configuration and the default DeLFT value will be used in this case.
     */