diff --git a/pom.xml b/pom.xml
index d31a663..cd71695 100755
--- a/pom.xml
+++ b/pom.xml
@@ -49,10 +49,15 @@
caffeine
- org.apache.commons
- commons-lang3
+ org.apache.commons
+ commons-lang3
-
+
+ commons-io
+ commons-io
+ 2.17.0
+
+
de.extio
lm-lib
diff --git a/src/main/java/de/extio/lmdatasetprep/Execution.java b/src/main/java/de/extio/lmdatasetprep/Execution.java
index b4d3733..03c182a 100644
--- a/src/main/java/de/extio/lmdatasetprep/Execution.java
+++ b/src/main/java/de/extio/lmdatasetprep/Execution.java
@@ -7,6 +7,7 @@ import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
+import java.util.ArrayList;
import java.util.List;
import java.util.Properties;
import java.util.concurrent.BlockingQueue;
@@ -24,6 +25,7 @@ import java.util.function.Function;
import java.util.function.Supplier;
import java.util.stream.Stream;
+import org.apache.commons.io.output.TeeOutputStream;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -159,50 +161,85 @@ public class Execution {
}
public static void streamOut(final Path f, final String propertyPrefix, final Properties properties, final Consumer consumer) {
- final var destination = properties.getProperty(propertyPrefix);
- if (destination.startsWith("queue")) {
- streamOutToQueue(f, destination, consumer);
+ final var streamFuncs = new ArrayList();
+ for (int i = 0; i < 2; i++) {
+ final var destination = properties.getProperty(propertyPrefix + "." + i);
+ if (destination != null) {
+ if (destination.startsWith("queue")) {
+ streamFuncs.add(setupStreamOutToQueue(f, destination));
+ }
+ else {
+ streamFuncs.add(setupStreamOutToFile(f, destination));
+ }
+ }
}
- else {
- streamOutToFile(f, destination, consumer);
- }
-
- }
-
- private static void streamOutToFile(final Path f, final String destination, final Consumer consumer) {
- final var out = Path.of(destination).resolve(f);
- if (Files.exists(out)) {
- LOGGER.info("File already exists: {} ", out);
+ if (streamFuncs.isEmpty()) {
return;
}
- final var tmp = out.resolveSibling(out.getFileName() + ".tmp");
- try (var fos = Files.newOutputStream(tmp)) {
- consumer.accept(fos);
- }
- catch (final IOException e1) {
- LOGGER.error("IO exception", e1);
- return;
- }
+ final var streams = streamFuncs.stream().map(streamFunc -> streamFunc.open().get()).toList();
+ final var outStream = streams.size() == 1 ? streams.getFirst() : new TeeOutputStream(streams.get(0), streams.get(1));
try {
- Files.move(tmp, out);
+ consumer.accept(outStream);
}
- catch (final IOException e) {
- LOGGER.error("IO exception", e);
+ finally {
+ try {
+ outStream.close();
+ }
+ catch (final IOException e) {
+ LOGGER.error("IO exception", e);
+ }
}
+
+ for (int i = 0; i < streamFuncs.size(); i++) {
+ streamFuncs.get(i).afterClose().accept(streams.get(i));
+ }
+
}
- private static void streamOutToQueue(final Path f, final String destination, final Consumer consumer) {
- try (var bos = new ByteArrayOutputStream()) {
- consumer.accept(bos);
- final var text = bos.toString(StandardCharsets.UTF_8);
- work.computeIfAbsent(destination, k -> new ConcurrentHashMap<>()).values().forEach(queue -> queue.add(new WorkPacket(f, text)));
- }
- catch (final IOException e) {
- LOGGER.error("IO exception", e);
- }
+ private static OutStream setupStreamOutToFile(final Path f, final String destination) {
+ return new OutStream(() -> {
+ final var out = Path.of(destination).resolve(f);
+ if (Files.exists(out)) {
+ LOGGER.info("File already exists: {} ", out);
+ return null;
+ }
+
+ final var tmp = out.resolveSibling(out.getFileName() + ".tmp");
+ try {
+ return Files.newOutputStream(tmp);
+ }
+ catch (final IOException e) {
+ LOGGER.error("IO exception", e);
+ return null;
+ }
+ },
+ (stream) -> {
+ try {
+ final var out = Path.of(destination).resolve(f);
+ final var tmp = out.resolveSibling(out.getFileName() + ".tmp");
+ Files.move(tmp, out);
+ }
+ catch (final IOException e) {
+ LOGGER.error("IO exception", e);
+ }
+ });
+ }
+
+ private static OutStream setupStreamOutToQueue(final Path f, final String destination) {
+ return new OutStream(() -> {
+ return new ByteArrayOutputStream();
+ },
+ (stream) -> {
+ final var text = ((ByteArrayOutputStream) stream).toString(StandardCharsets.UTF_8);
+ work.computeIfAbsent(destination, k -> new ConcurrentHashMap<>()).values().forEach(queue -> queue.add(new WorkPacket(f, text)));
+ });
}
public record WorkPacket(Path file, String text) {
}
+
+ private record OutStream(Supplier open, Consumer afterClose) {
+
+ }
}
diff --git a/src/main/resources/tool.properties.template b/src/main/resources/tool.properties.template
index f0a75df..0b4e831 100644
--- a/src/main/resources/tool.properties.template
+++ b/src/main/resources/tool.properties.template
@@ -2,27 +2,28 @@ translate.category=SMALL
# Source and destination can be either a path...
translate.source=/home/xxx/folder
# ... or a queue. Queues are prefixed with "queue", after the prefix comes the queue name
-translate.destination=queue0
+translate.destination.0=queue0
rewrite.category=SMALL
rewrite.source=queue0
-rewrite.destination=queue1
+rewrite.destination.0=queue1
rewrite.improve=1
rewrite.enhance=2
rewrite.model=l3dp
# .n indicates that the key supports sequences. Add as many keys as you need
mergeJsonl.source.0=
-mergeJsonl.destination=
+# Currently we only support at most 2 destinations. This might change in the future
+mergeJsonl.destination.0=
text2Jsonl.source=
-text2Jsonl.destination=
+text2Jsonl.destination.0=
text2Jsonl.chunkNorm=750
text2Jsonl.chunkVar=250
contextualPrompts.category=SMALL
contextualPrompts.source=queue1
-contextualPrompts.destination=/home/xxx/outputdir
+contextualPrompts.destination.0=/home/xxx/outputdir
contextualPrompts.model=l3dp
contextualPrompts.variations=1
contextualPrompts.chunkNorm.0=750
@@ -31,7 +32,7 @@ contextualPrompts.chunkNorm.1=1500
contextualPrompts.chunkVar.1=300
create.source=
-create.destination=
+create.destination.0=
create.category=SMALL
create.model=l3dp
create.plotLimit=6000