update support to trainer

kermitt2 · Jan 6, 2025 · f32f90e · f32f90e
1 parent 702407e
commit f32f90e
Show file tree

Hide file tree

Showing 4 changed files with 19 additions and 19 deletions.
diff --git a/grobid-trainer/src/main/java/org/grobid/trainer/HeaderTrainer.java b/grobid-trainer/src/main/java/org/grobid/trainer/HeaderTrainer.java
@@ -296,19 +296,16 @@ public boolean accept(File dir, String name) {
      * Command line execution.
      *
      * @param args Command line arguments.
-     * @throws Exception 
      */
     public static void main(String[] args) throws Exception {
         // if we have a parameter, it gives the flavor refinement to consider
         Flavor theFlavor = null;
         if (args.length > 0) {
             String flavor = args[0];
-            if (flavor.toLowerCase().equals("light")) {
-                theFlavor = Flavor.LIGHT;
-            } else if (flavor.toLowerCase().equals("ietf")) {
-                theFlavor = Flavor.IETF;
-            } else {
-                System.out.println("Warning, the flavor is not recognized, must one one of [3gpp,ietf], defaulting training to no collection...");
+            theFlavor = GrobidModels.Flavor.fromLabel(flavor);
+            if (theFlavor == null) {
+                System.out.println("Warning, the flavor is not recognized, " +
+                    "must one one of " + Flavor.getLabels() + ", defaulting training to no collection...");
             }
         }
 

diff --git a/grobid-trainer/src/main/java/org/grobid/trainer/SegmentationTrainer.java b/grobid-trainer/src/main/java/org/grobid/trainer/SegmentationTrainer.java
@@ -348,12 +348,11 @@ public static void main(String[] args) throws Exception {
         Flavor theFlavor = null;
         if (args.length > 0) {
             String flavor = args[0];
-            if (flavor.toLowerCase().equals("light")) {
-                theFlavor = Flavor.LIGHT;
-            } else if (flavor.toLowerCase().equals("ietf")) {
-                theFlavor = Flavor.IETF;
-            } else {
-                System.out.println("Warning, the flavor is not recognized, must one one of [3gpp,ietf], defaulting training to no collection...");
+            theFlavor = Flavor.fromLabel(flavor);
+            if (theFlavor == null) {
+                System.out.println("Warning, the flavor is not recognized, " +
+                    "must one one of "+ Flavor.getLabels() +", " +
+                    "defaulting training with no flavor...");
             }
         }
 

diff --git a/grobid-trainer/src/main/java/org/grobid/trainer/TrainerRunner.java b/grobid-trainer/src/main/java/org/grobid/trainer/TrainerRunner.java
@@ -17,8 +17,12 @@
  */
 public class TrainerRunner {
 
-    private static final List<String> models = Arrays.asList("affiliation", "chemical", "date", "citation", 
-        "ebook", "fulltext", "header", "header-ietf", "name-citation", "name-header", "patent", "segmentation",
+    private static final List<String> models = Arrays.asList(
+        "affiliation", "chemical", "date", "citation",
+        "ebook", "fulltext", "header",
+        "header-ietf",
+        "name-citation", "name-header", "patent",
+        "segmentation",
         "segmentation-ietf");
     private static final List<String> options = Arrays.asList("0 - train", "1 - evaluate", "2 - split, train and evaluate", "3 - n-fold evaluation");
 

diff --git a/grobid-trainer/src/main/java/org/grobid/trainer/sax/TEIHeaderSaxParser.java b/grobid-trainer/src/main/java/org/grobid/trainer/sax/TEIHeaderSaxParser.java
@@ -40,7 +40,7 @@ public class TEIHeaderSaxParser extends DefaultHandler {
     private List<String> ignoredTags = Arrays.asList("location", "version", "web", "degree", "page", "title", "phone", "publisher");  
 
     public TEIHeaderSaxParser() {
-        labeled = new ArrayList<String>();
+        labeled = new ArrayList<>();
     }
 
     public void characters(char[] buffer, int start, int length) {
@@ -65,9 +65,9 @@ public ArrayList<String> getLabeledResult() {
         return labeled;
     }
 
-    public void endElement(java.lang.String uri,
-                           java.lang.String localName,
-                           java.lang.String qName) throws SAXException {
+    public void endElement(String uri,
+                           String localName,
+                           String qName) throws SAXException {
         if (endTags.contains(qName)) {
             writeData();
             accumulator.setLength(0);