From 77637436f2850035064e0f6cad16aa415a808f56 Mon Sep 17 00:00:00 2001
From: extio-de <github@extio.de>
Date: Sat, 25 Jan 2025 17:36:05 +0100
Subject: [PATCH] MergeJsonl: Separate outputs per filter

---
 .../extio/lmdatasetprep/tools/MergeJsonl.java | 67 ++++++++++---------
 1 file changed, 36 insertions(+), 31 deletions(-)

diff --git a/src/main/java/de/extio/lmdatasetprep/tools/MergeJsonl.java b/src/main/java/de/extio/lmdatasetprep/tools/MergeJsonl.java
index 12741f7..fd141e1 100644
--- a/src/main/java/de/extio/lmdatasetprep/tools/MergeJsonl.java
+++ b/src/main/java/de/extio/lmdatasetprep/tools/MergeJsonl.java
@@ -10,6 +10,7 @@ import java.util.List;
 import java.util.Properties;
 import java.util.concurrent.ThreadLocalRandom;
 
+import org.apache.commons.lang3.StringUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.springframework.stereotype.Component;
@@ -25,44 +26,48 @@ public class MergeJsonl implements DatasetTool {
 	
 	@Override
 	public void accept(final Properties properties) {
-		final List<String> lines = Collections.synchronizedList(new ArrayList<>());
-		for (int i = 0; i < 99; i++) {
-			final String source = properties.getProperty("mergeJsonl.source." + i, "");
-			if (source.isBlank()) {
-				break;
-			}
-			Execution.transform(source, p -> {
-				final var filter = properties.getProperty("mergeJsonl.filter", "");
-				if (!filter.isEmpty() && !p.file().getFileName().toString().contains(filter)) {
-					return List.of();
+		final var filters = properties.getProperty("mergeJsonl.filter", "*");
+		for (final var filter : StringUtils.split(filters, ',')) {
+			final List<String> lines = Collections.synchronizedList(new ArrayList<>());
+			for (int i = 0; i < 99; i++) {
+				final String source = properties.getProperty("mergeJsonl.source." + i, "");
+				if (source.isBlank()) {
+					break;
 				}
-				
-				return List.of(() -> {
-					try {
-						lines.addAll(Files.readAllLines(p.file()));
-					}
-					catch (final IOException e) {
-						throw new RuntimeException("Cannot read file", e);
+				Execution.transform(source, p -> {
+					if (!filter.equals("*") && !p.file().getFileName().toString().contains(filter)) {
+						return List.of();
 					}
 					
+					return List.of(() -> {
+						try {
+							lines.addAll(Files.readAllLines(p.file()));
+						}
+						catch (final IOException e) {
+							throw new RuntimeException("Cannot read file", e);
+						}
+						
+					});
 				});
+			}
+			
+			Collections.shuffle(lines, ThreadLocalRandom.current());
+			
+			final Path out = Path.of("dataset-merged" + (filter.equals("*") ? "" : filter) + ".jsonl");
+			
+			Execution.streamOut(out, "mergeJsonl.destination", properties, fos -> {
+				try {
+					for (final String line : lines) {
+						fos.write(line.getBytes(StandardCharsets.UTF_8));
+						fos.write(NEWLINE);
+					}
+				}
+				catch (final IOException e1) {
+					LOGGER.error("IO exception", e1);
+				}
 			});
 		}
 		
-		Collections.shuffle(lines, ThreadLocalRandom.current());
-		
-		final Path out = Path.of("dataset-merged.jsonl");
-		Execution.streamOut(out, "mergeJsonl.destination", properties, fos -> {
-			try {
-				for (final String line : lines) {
-					fos.write(line.getBytes(StandardCharsets.UTF_8));
-					fos.write(NEWLINE);
-				}
-			}
-			catch (final IOException e1) {
-				LOGGER.error("IO exception", e1);
-			}
-		});
 	}
 	
 }