fix for training data generation with updated affiliation-address parser

kermitt2 · Dec 21, 2023 · e235c88 · e235c88
1 parent a9b6826
commit e235c88
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/grobid-core/src/main/java/org/grobid/core/engines/FullTextParser.java b/grobid-core/src/main/java/org/grobid/core/engines/FullTextParser.java
@@ -1364,9 +1364,15 @@ public Document createTraining(File inputFile,
                     }
 
                     // buffer for the affiliation+address block
-                    List<LayoutToken> tokenizationsAffiliation = resHeader.getLayoutTokens(TaggingLabels.HEADER_AFFILIATION);
+
+                    List<List<LayoutToken>> tokenizationsAffiliation = resHeader.getAffiliationAddresslabeledTokens();
+                    //List<LayoutToken> tokenizationsAffiliation = resHeader.getLayoutTokens(TaggingLabels.HEADER_AFFILIATION);
+                    List<LayoutToken> tokenizationAffiliation = new ArrayList<>();
+                    for (List<LayoutToken> tokenization : tokenizationsAffiliation) {
+                        tokenizationAffiliation.addAll(tokenization);
+                    }
                     StringBuilder bufferAffiliation =
-                            parsers.getAffiliationAddressParser().trainingExtraction(tokenizationsAffiliation);
+                            parsers.getAffiliationAddressParser().trainingExtraction(tokenizationAffiliation);
 
                     // buffer for the date block
                     StringBuilder bufferDate = null;