add segmentation light with references

kermitt2 · Sep 21, 2024 · dfc7606 · dfc7606
1 parent 2232b2c
commit dfc7606
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 0 deletions.
diff --git a/build.gradle b/build.gradle
@@ -520,9 +520,11 @@ project(":grobid-trainer") {
     def complexTrainerTasks = [
         "train_header"                : ["org.grobid.trainer.HeaderTrainer", ""],
         "train_header_light"           : ["org.grobid.trainer.HeaderTrainer", "light"],
+        "train_header_light_ref"           : ["org.grobid.trainer.HeaderTrainer", "light-ref"],
         "train_header_ietf"           : ["org.grobid.trainer.HeaderTrainer", "ietf"],
         "train_segmentation"          : ["org.grobid.trainer.SegmentationTrainer", ""],
         "train_segmentation_light"     : ["org.grobid.trainer.SegmentationTrainer", "light"],
+        "train_segmentation_light_ref"     : ["org.grobid.trainer.SegmentationTrainer", "light-ref"],
         "train_segmentation_ietf"     : ["org.grobid.trainer.SegmentationTrainer", "ietf"]
     ]
 

diff --git a/grobid-trainer/src/main/java/org/grobid/trainer/SegmentationTrainer.java b/grobid-trainer/src/main/java/org/grobid/trainer/SegmentationTrainer.java
@@ -250,6 +250,8 @@ public static void main(String[] args) throws Exception {
             String flavor = args[0];
             if (flavor.equalsIgnoreCase("light")) {
                 theFlavor = Flavor.ARTICLE_LIGHT;
+            } else if (flavor.equalsIgnoreCase("light-re")) {
+                theFlavor = Flavor.ARTICLE_LIGHT_WITH_REFERENCES;
             } else if (flavor.equalsIgnoreCase("ietf")) {
                 theFlavor = Flavor.IETF;
             } else {