diff --git a/src/main/java/de/extio/lmdatasetprep/Execution.java b/src/main/java/de/extio/lmdatasetprep/Execution.java index 97dd578..ef946ca 100644 --- a/src/main/java/de/extio/lmdatasetprep/Execution.java +++ b/src/main/java/de/extio/lmdatasetprep/Execution.java @@ -14,8 +14,10 @@ import java.util.Properties; import java.util.concurrent.BlockingQueue; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentMap; +import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; +import java.util.concurrent.Future; import java.util.concurrent.LinkedBlockingQueue; import java.util.concurrent.Semaphore; import java.util.concurrent.TimeUnit; @@ -45,6 +47,8 @@ public class Execution { private static final Semaphore taskSemaphore; + private static final ThreadLocal>> localOutWorkers = ThreadLocal.withInitial(ArrayList::new); + private static final Logger LOGGER = LoggerFactory.getLogger(Execution.class); static { @@ -120,7 +124,7 @@ public class Execution { try { taskSemaphore.acquire(); tasksRunning.incrementAndGet(); - executorService.execute(() -> { + localOutWorkers.get().add(executorService.submit(() -> { try { r.run(); } @@ -128,7 +132,7 @@ public class Execution { tasksRunning.decrementAndGet(); taskSemaphore.release(); } - }); + })); } catch (final InterruptedException e) { throw new RuntimeException("Interrupted", e); @@ -182,17 +186,15 @@ public class Execution { return; } - final var outStream = new ProxyOutputStream(streams); - try { + localOutWorkers.get().clear(); + try (final var outStream = new ProxyOutputStream(streams)) { consumer.accept(outStream); + for (final var w : localOutWorkers.get()) { + w.get(); + } } - finally { - try { - outStream.close(); - } - catch (final IOException e) { - LOGGER.error("IO exception", e); - } + catch (InterruptedException | ExecutionException | IOException e) { + LOGGER.error("Error while streaming out", e); } for (int i = 0; i < streamFuncs.size(); i++) { diff --git a/src/main/java/de/extio/lmdatasetprep/LmDatasetPreparerApplication.java b/src/main/java/de/extio/lmdatasetprep/LmDatasetPreparerApplication.java index 633661a..d7dad11 100755 --- a/src/main/java/de/extio/lmdatasetprep/LmDatasetPreparerApplication.java +++ b/src/main/java/de/extio/lmdatasetprep/LmDatasetPreparerApplication.java @@ -72,22 +72,19 @@ public class LmDatasetPreparerApplication implements CommandLineRunner, Applicat t.start(); } - do { - Thread.sleep(1000); - } while (Execution.startedFileConsumers.get() == 0 || Execution.startedFileConsumers.get() > Execution.finishedFileConsumers.get()); - LOGGER.info("All file consumers are finished"); - int cnt = 0; while (cnt < 5) { Thread.sleep(1000); - if (Execution.tasksRunning.get() != 0 || Execution.work.values().stream().anyMatch(q -> q.values().stream().anyMatch(bq -> !bq.isEmpty()))) { + if (Execution.startedFileConsumers.get() == 0 || Execution.startedFileConsumers.get() > Execution.finishedFileConsumers.get() || + Execution.tasksRunning.get() != 0 || Execution.work.values().stream().anyMatch(q -> q.values().stream().anyMatch(bq -> !bq.isEmpty()))) { + cnt = 0; } else { cnt++; } } - LOGGER.info("All queues are empty"); + LOGGER.info("All file consumers are finished, All queues are empty"); Execution.shutdown.set(true); for (final var thread : threads) { diff --git a/src/main/java/de/extio/lmdatasetprep/tools/AbstractContextualPrompts.java b/src/main/java/de/extio/lmdatasetprep/tools/AbstractContextualPrompts.java index 684ee46..280d37d 100644 --- a/src/main/java/de/extio/lmdatasetprep/tools/AbstractContextualPrompts.java +++ b/src/main/java/de/extio/lmdatasetprep/tools/AbstractContextualPrompts.java @@ -189,4 +189,7 @@ abstract class AbstractContextualPrompts implements InitializingBean, DatasetToo protected record ConversationsLine(List conversation) { } + protected record TextLine(String text) { + } + } diff --git a/src/main/java/de/extio/lmdatasetprep/tools/GenMergedText.java b/src/main/java/de/extio/lmdatasetprep/tools/GenMergedText.java new file mode 100644 index 0000000..15d8c76 --- /dev/null +++ b/src/main/java/de/extio/lmdatasetprep/tools/GenMergedText.java @@ -0,0 +1,57 @@ +package de.extio.lmdatasetprep.tools; + +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.List; +import java.util.Properties; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import org.springframework.stereotype.Component; + +import de.extio.lmdatasetprep.Execution; + +@Component +public class GenMergedText extends AbstractContextualPrompts { + + private static final Logger LOGGER = LoggerFactory.getLogger(GenMergedText.class); + + @Override + protected CreateTasks createTasks(final Properties properties, final List chunkConfigurations, final int variations) { + return p -> { + final List tasks = new ArrayList<>(); + + for (final ChunkCfg chunkCfg : chunkConfigurations) { + for (int i = 0; i < variations; i++) { + final int fi = i; + + tasks.add(() -> { + final List paragraphs = this.fileToParagraphs(p, chunkCfg.chunksNorm(), chunkCfg.chunksVar()); + + final Path out = Execution.suffixFilename(p.file().getFileName(), + "text", + properties.getProperty("contextualPrompts.model"), + String.valueOf(chunkCfg.chunksNorm()), + String.valueOf(fi), + ".jsonl"); + + Execution.streamOut(out, "contextualPrompts.destination", properties, fos -> { + for (int j = 0; j < paragraphs.size(); j++) { + LOGGER.info("Paragraph " + (j + 1) + "/" + paragraphs.size()); + + final var mapping = this.createCharacterNameMapping(paragraphs.get(j), properties); + final var paragraph = this.renameCharacters(paragraphs.get(j), mapping); + final var textLine = new TextLine(paragraph); + + this.writeJsonLine(fos, textLine); + } + }); + }); + } + } + + return tasks; + }; + } + +} diff --git a/src/main/java/de/extio/lmdatasetprep/tools/MergeJsonl.java b/src/main/java/de/extio/lmdatasetprep/tools/MergeJsonl.java index fd141e1..a03f4a9 100644 --- a/src/main/java/de/extio/lmdatasetprep/tools/MergeJsonl.java +++ b/src/main/java/de/extio/lmdatasetprep/tools/MergeJsonl.java @@ -2,9 +2,9 @@ package de.extio.lmdatasetprep.tools; import java.io.IOException; import java.nio.charset.StandardCharsets; -import java.nio.file.Files; import java.nio.file.Path; import java.util.ArrayList; +import java.util.Arrays; import java.util.Collections; import java.util.List; import java.util.Properties; @@ -26,7 +26,7 @@ public class MergeJsonl implements DatasetTool { @Override public void accept(final Properties properties) { - final var filters = properties.getProperty("mergeJsonl.filter", "*"); + final var filters = StringUtils.defaultIfEmpty(properties.getProperty("mergeJsonl.filter", "*"), "*"); for (final var filter : StringUtils.split(filters, ',')) { final List lines = Collections.synchronizedList(new ArrayList<>()); for (int i = 0; i < 99; i++) { @@ -40,13 +40,7 @@ public class MergeJsonl implements DatasetTool { } return List.of(() -> { - try { - lines.addAll(Files.readAllLines(p.file())); - } - catch (final IOException e) { - throw new RuntimeException("Cannot read file", e); - } - + lines.addAll(Arrays.asList(StringUtils.split(p.text(), '\n'))); }); }); } diff --git a/src/main/java/de/extio/lmdatasetprep/tools/Text2Jsonl.java b/src/main/java/de/extio/lmdatasetprep/tools/Text2Jsonl.java index 2de8923..63f52cb 100644 --- a/src/main/java/de/extio/lmdatasetprep/tools/Text2Jsonl.java +++ b/src/main/java/de/extio/lmdatasetprep/tools/Text2Jsonl.java @@ -27,7 +27,7 @@ public class Text2Jsonl implements DatasetTool { @Override public void accept(final Properties properties) { - final Path out = Path.of("dataset.jsonl"); + final Path out = Path.of("dataset-text.jsonl"); Execution.streamOut(out, "text2Jsonl.destination", properties, fos -> { Execution.transform(properties.getProperty("text2Jsonl.source"), p -> { return List.of(() -> {