From dfc760608f8730c0f02f7441e00f64271236e7e7 Mon Sep 17 00:00:00 2001
From: Luca Foppiano <luca@foppiano.org>
Date: Sat, 21 Sep 2024 10:50:33 +0200
Subject: [PATCH] add segmentation light with references

---
 build.gradle                                                    | 2 ++
 .../src/main/java/org/grobid/trainer/SegmentationTrainer.java   | 2 ++
 2 files changed, 4 insertions(+)

diff --git a/build.gradle b/build.gradle
index 77b0772ea4..7b2192428f 100644
--- 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
index 321de5e804..cb055cdb06 100644
--- 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 {