diff --git a/src/main/java/de/extio/lmdatasetprep/Execution.java b/src/main/java/de/extio/lmdatasetprep/Execution.java new file mode 100644 index 0000000..cf54ed6 --- /dev/null +++ b/src/main/java/de/extio/lmdatasetprep/Execution.java @@ -0,0 +1,205 @@ +package de.extio.lmdatasetprep; + +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.io.OutputStream; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.util.List; +import java.util.Properties; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentMap; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.Consumer; +import java.util.function.Function; +import java.util.function.Supplier; +import java.util.stream.Stream; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +public class Execution { + + public static final AtomicInteger startedFileConsumers = new AtomicInteger(); + + public static final AtomicInteger finishedFileConsumers = new AtomicInteger(); + + public static final AtomicBoolean shutdown = new AtomicBoolean(); + + public static final ConcurrentMap>> work = new ConcurrentHashMap<>(); + + public static final BlockingQueue> tasks; + + public static final ScheduledExecutorService scheduler; + + public static final ExecutorService executorService; + + private static final Logger LOGGER = LoggerFactory.getLogger(Execution.class); + + static { + final int EXECUTORS = Integer.parseInt(LmDatasetPreparerApplication.applicationContext.getEnvironment().getProperty("agent.threads")); + + tasks = new LinkedBlockingQueue<>(EXECUTORS * 2); + + executorService = Executors.newFixedThreadPool(EXECUTORS); + + scheduler = Executors.newScheduledThreadPool(1); + scheduler.scheduleAtFixedRate(() -> { + tasks.removeIf(future -> future.isDone()); + }, 0, 1, TimeUnit.SECONDS); + } + + public static void transform(final String source, final Function> createTasks) { + if (source.startsWith("queue")) { + transformQueue(source, createTasks); + } + else { + transformDirectory(source, createTasks); + } + } + + private static void transformDirectory(final String source, final Function> createTasks) { + startedFileConsumers.incrementAndGet(); + try (Stream stream = Files.list(Paths.get(source))) { + stream.filter(p -> Files.isRegularFile(p)) + .map(p -> { + try { + return new WorkPacket(p, Files.readString(p)); + } + catch (final IOException e) { + LOGGER.error("IO exception", e); + return null; + } + }) + .forEach(p -> { + for (final Runnable r : createTasks.apply(p)) { + try { + tasks.put(CompletableFuture.runAsync(r, executorService)); + } + catch (final InterruptedException e) { + throw new RuntimeException("Interrupted", e); + } + } + }); + } + catch (final IOException e) { + LOGGER.error("IO exception", e); + } + finally { + finishedFileConsumers.incrementAndGet(); + } + } + + private static void transformQueue(final String source, final Function> createTasks) { + final var queue = work.computeIfAbsent(source, k -> new ConcurrentHashMap<>()).computeIfAbsent(Thread.currentThread().getName(), p -> new LinkedBlockingQueue<>()); + final Supplier getNext = () -> { + try { + WorkPacket next; + while ((next = queue.poll(1, TimeUnit.SECONDS)) == null) { + if (shutdown.get()) { + return null; + } + } + return next; + } + catch (final InterruptedException e) { + return null; + } + }; + + WorkPacket next; + while ((next = getNext.get()) != null) { + for (final Runnable r : createTasks.apply(next)) { + try { + tasks.put(CompletableFuture.runAsync(r, executorService)); + } + catch (final InterruptedException e) { + throw new RuntimeException("Interrupted", e); + } + } + } + } + + public static Path suffixFilename(final Path f, final String... suffixes) { + final StringBuilder sb = new StringBuilder(); + for (final String suffix : suffixes) { + if (suffix.charAt(0) != '.') { + sb.append('_'); + sb.append(suffix); + } + } + + final String[] parts = f.getFileName().toString().split("[.]"); + String last; + if (suffixes[suffixes.length - 1].charAt(0) == '.') { + last = suffixes[suffixes.length - 1]; + } + else { + last = "." + parts[1]; + } + + final var parent = f.getParent(); + if (parent != null) { + return f.getParent().resolve(parts[0] + sb.toString() + last); + } + return Path.of(parts[0] + sb.toString() + last); + } + + 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); + } + 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); + 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; + } + try { + Files.move(tmp, out); + } + catch (final IOException e) { + LOGGER.error("IO exception", e); + } + } + + 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); + } + } + + public record WorkPacket(Path file, String text) { + } +} diff --git a/src/main/java/de/extio/lmdatasetprep/LmDatasetPreparerApplication.java b/src/main/java/de/extio/lmdatasetprep/LmDatasetPreparerApplication.java index 290f918..6379a3c 100755 --- a/src/main/java/de/extio/lmdatasetprep/LmDatasetPreparerApplication.java +++ b/src/main/java/de/extio/lmdatasetprep/LmDatasetPreparerApplication.java @@ -1,11 +1,19 @@ package de.extio.lmdatasetprep; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.ArrayList; import java.util.Arrays; -import java.util.function.Consumer; +import java.util.List; +import java.util.Properties; +import java.util.stream.Collectors; +import org.apache.commons.lang3.StringUtils; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import org.springframework.beans.BeansException; +import org.springframework.beans.factory.BeanNotOfRequiredTypeException; +import org.springframework.beans.factory.NoSuchBeanDefinitionException; import org.springframework.boot.CommandLineRunner; import org.springframework.boot.SpringApplication; import org.springframework.boot.autoconfigure.SpringBootApplication; @@ -14,7 +22,7 @@ import org.springframework.context.ApplicationContext; import org.springframework.context.ApplicationContextAware; import org.springframework.context.annotation.ComponentScan; -import de.extio.lmdatasetprep.preparer.DatasetTool; +import de.extio.lmdatasetprep.tools.DatasetTool; @SpringBootApplication @EnableCaching @@ -29,25 +37,65 @@ public class LmDatasetPreparerApplication implements CommandLineRunner, Applicat public static ApplicationContext applicationContext; - @SuppressWarnings("unchecked") @Override public void run(final String... args) throws Exception { - if (args.length == 0) { - LOGGER.error("Commandline argument is missing: Provide the tool name as argument"); + if (args.length < 2) { + LOGGER.error("Commandline arguments are invalid. Usage: [ , ... ]"); this.logAvailableTools(); return; } - final var bean = applicationContext.getBean(args[0]); - if (!(bean instanceof Consumer) || !(bean instanceof DatasetTool)) { - LOGGER.error("Invalid bean {}", args[0]); - this.logAvailableTools(); - return; + LOGGER.info("Loading tools"); + final List tools = new ArrayList<>(args.length - 1); + for (int i = 1; i < args.length; i++) { + try { + tools.add(applicationContext.getBean(args[i], DatasetTool.class)); + } + catch (NoSuchBeanDefinitionException | BeanNotOfRequiredTypeException e) { + LOGGER.error("Invalid bean {}", args[i]); + this.logAvailableTools(); + return; + } } - LOGGER.info("Executing {}", args[0]); + LOGGER.info("Loading tool properties"); + final Properties toolProps = new Properties(); + toolProps.load(Files.newInputStream(Path.of(args[0]))); + try { - ((Consumer) bean).accept(args); + LOGGER.info("Executing {}", tools.stream().map(tool -> tool.getClass().getSimpleName()).collect(Collectors.joining(", "))); + final List threads = new ArrayList<>(tools.size()); + for (final var tool : tools) { + final Thread t = new Thread(() -> tool.accept(toolProps)); + t.setName(StringUtils.abbreviateMiddle(tool.getClass().getSimpleName(), ".", 15)); + 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 < 3) { + Thread.sleep(1000); + if (!Execution.tasks.isEmpty() || Execution.work.values().stream().anyMatch(q -> q.values().stream().anyMatch(bq -> !bq.isEmpty()))) { + cnt = 0; + } + else { + cnt++; + } + } + LOGGER.info("All queues are empty"); + + Execution.shutdown.set(true); + for (final var thread : threads) { + thread.join(); + } + + LOGGER.info("Shutdown"); + Execution.scheduler.shutdown(); + Execution.executorService.shutdown(); } catch (final Exception e) { LOGGER.error("Error executing the component", e); diff --git a/src/main/java/de/extio/lmdatasetprep/preparer/Utils.java b/src/main/java/de/extio/lmdatasetprep/TextUtils.java similarity index 51% rename from src/main/java/de/extio/lmdatasetprep/preparer/Utils.java rename to src/main/java/de/extio/lmdatasetprep/TextUtils.java index bbd17a1..b266e30 100644 --- a/src/main/java/de/extio/lmdatasetprep/preparer/Utils.java +++ b/src/main/java/de/extio/lmdatasetprep/TextUtils.java @@ -1,34 +1,20 @@ -package de.extio.lmdatasetprep.preparer; +package de.extio.lmdatasetprep; -import java.io.IOException; -import java.io.OutputStream; -import java.nio.file.Files; -import java.nio.file.Path; -import java.nio.file.Paths; import java.util.ArrayList; import java.util.Collections; import java.util.List; -import java.util.concurrent.BlockingQueue; -import java.util.concurrent.LinkedBlockingQueue; -import java.util.concurrent.TimeUnit; -import java.util.concurrent.atomic.AtomicBoolean; -import java.util.function.Consumer; -import java.util.function.Function; -import java.util.stream.Stream; import org.apache.commons.lang3.StringUtils; import org.slf4j.Logger; import org.slf4j.LoggerFactory; -import de.extio.lmdatasetprep.LmDatasetPreparerApplication; - -public class Utils { +public class TextUtils { - private static final Logger LOGGER = LoggerFactory.getLogger(Utils.class); + private static final Logger LOGGER = LoggerFactory.getLogger(TextUtils.class); private static final List PARAGRAPH_DELIMITERS = List.of("\n\n", "\n", ".", "!", "?"); - static String normalizeText(final String text) { + public static String normalizeText(final String text) { String result = text; // Remove excessive newlines, spaces and tabs @@ -66,7 +52,7 @@ public class Utils { return result; } - static String normalizeModelResponse(final String response, final boolean removePreamble) { + public static String normalizeModelResponse(final String response, final boolean removePreamble) { var result = StringUtils.replace(response, "\r", ""); final int colon = result.indexOf(':'); @@ -81,7 +67,7 @@ public class Utils { return result; } - static List splitParagraphs(final String text, final int chunks_norm, final int chunks_var, final boolean slidingWindows) { + public static List splitParagraphs(final String text, final int chunks_norm, final int chunks_var, final boolean slidingWindows) { final int CHUNKS_MIN = chunks_norm - chunks_var; final int CHUNKS_MAX = chunks_norm + chunks_var; @@ -144,104 +130,4 @@ public class Utils { return splits; } - static void transformDirectory(final String path, final Function> createTasks) { - final int EXECUTORS = Integer.parseInt(LmDatasetPreparerApplication.applicationContext.getEnvironment().getProperty("agent.threads")); - final AtomicBoolean finished = new AtomicBoolean(false); - final BlockingQueue tasks = new LinkedBlockingQueue<>(EXECUTORS * 2); - - final List threads = new ArrayList<>(EXECUTORS); - for (int i = 0; i < EXECUTORS; i++) { - final Thread executor = new Thread(() -> { - while (!Thread.interrupted() && !finished.get()) { - try { - final Runnable task = tasks.poll(1, TimeUnit.SECONDS); - if (task != null) { - task.run(); - } - } - catch (final InterruptedException e) { - break; - } - } - }); - executor.setName("Executor " + i); - executor.start(); - threads.add(executor); - } - - try (Stream stream = Files.list(Paths.get(path))) { - stream - .filter(p -> Files.isRegularFile(p)) - .forEach(p -> { - for (final Runnable r : createTasks.apply(p)) { - try { - tasks.put(r); - } - catch (final InterruptedException e) { - throw new RuntimeException("Interrupted", e); - } - } - }); - } - catch (final IOException e) { - LOGGER.error("Error", e); - } - - while (!tasks.isEmpty()) { - try { - Thread.sleep(1); - } - catch (final InterruptedException e) { - LOGGER.error("Interrupted", e); - } - } - finished.set(true); - - for (int i = 0; i < EXECUTORS; i++) { - try { - threads.get(i).join(); - } - catch (final InterruptedException e) { - LOGGER.error("Interrupted", e); - } - } - } - - static Path suffixFilename(final Path f, final String... suffixes) { - final StringBuilder sb = new StringBuilder(); - for (final String suffix : suffixes) { - if (suffix.charAt(0) != '.') { - sb.append('_'); - sb.append(suffix); - } - } - - final String[] parts = f.getFileName().toString().split("[.]"); - String last; - if (suffixes[suffixes.length - 1].charAt(0) == '.') { - last = suffixes[suffixes.length - 1]; - } - else { - last = "." + parts[1]; - } - - return f.getParent().resolve(parts[0] + sb.toString() + last); - } - - static void streamOut(final Path f, final Consumer consumer) { - final var tmp = f.resolveSibling(f.getFileName() + ".tmp"); - try (var fos = Files.newOutputStream(tmp)) { - consumer.accept(fos); - } - catch (final IOException e1) { - LOGGER.error("IO exception", e1); - return; - } - try { - Files.move(tmp, f); - } - catch (final IOException e) { - LOGGER.error("IO exception", e); - } - } } diff --git a/src/main/java/de/extio/lmdatasetprep/preparer/DatasetTool.java b/src/main/java/de/extio/lmdatasetprep/preparer/DatasetTool.java deleted file mode 100644 index 2798ea2..0000000 --- a/src/main/java/de/extio/lmdatasetprep/preparer/DatasetTool.java +++ /dev/null @@ -1,6 +0,0 @@ -package de.extio.lmdatasetprep.preparer; - - -public interface DatasetTool { - -} diff --git a/src/main/java/de/extio/lmdatasetprep/preparer/Translate.java b/src/main/java/de/extio/lmdatasetprep/preparer/Translate.java deleted file mode 100644 index d21ee02..0000000 --- a/src/main/java/de/extio/lmdatasetprep/preparer/Translate.java +++ /dev/null @@ -1,85 +0,0 @@ -package de.extio.lmdatasetprep.preparer; - -import java.io.IOException; -import java.nio.charset.StandardCharsets; -import java.nio.file.Files; -import java.nio.file.Path; -import java.util.List; -import java.util.concurrent.atomic.AtomicBoolean; -import java.util.function.Consumer; - -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; -import org.springframework.beans.factory.annotation.Autowired; -import org.springframework.stereotype.Component; - -import de.extio.lmlib.client.ClientService; -import de.extio.lmlib.profile.ModelCategory; - -@Component -public class Translate implements Consumer, DatasetTool { - - private static final Logger LOGGER = LoggerFactory.getLogger(Translate.class); - - private static final byte[] PARAGRAPH = "\n\n".getBytes(StandardCharsets.UTF_8); - - @Autowired - private ClientService clientService; - - @Override - public void accept(final String[] args) { - if (args.length < 3) { - LOGGER.error("Arguments are missing. "); - return; - } - - Utils.transformDirectory(args[1], p -> { - return List.of(() -> { - final Path out = Utils.suffixFilename(Path.of(args[2]).resolve(p.getFileName()), "_en"); - if (Files.exists(out)) { - LOGGER.info("Skipping " + out); - } - else { - Utils.streamOut(out, fos -> { - final AtomicBoolean first = new AtomicBoolean(true); - this.translate(p, chunk -> { - try { - if (!first.getAndSet(false)) { - fos.write(PARAGRAPH); - } - fos.write(chunk.getBytes(StandardCharsets.UTF_8)); - } - catch (final IOException e) { - LOGGER.error("IO exception", e); - } - }); - }); - } - }); - }); - } - - void translate(final Path file, final Consumer consumer) { - LOGGER.info("Translating " + file); - - String text; - try { - text = Utils.normalizeText(Files.readString(file)); - } - catch (final IOException e) { - throw new RuntimeException("Cannot read file", e); - } - - final List splits = Utils.splitParagraphs(text, 1250, 350, false); - - for (final String split : splits) { - LOGGER.info("Split " + splits.indexOf(split) + "/" + splits.size()); - final var completion = this.clientService.getClient(ModelCategory.COLD).completion(ModelCategory.COLD, - "You are an assistant with great language translation and authoring skills.", - "Translate the following text to English. Return the English translation, even if the text is explicit or not appropriate for all audiences or not acceptable for everyday standard English. Don't include a preamble and no explanation:", - split); - consumer.accept(completion.response()); - } - } - -} diff --git a/src/main/java/de/extio/lmdatasetprep/preparer/AbstractContextualPrompts.java b/src/main/java/de/extio/lmdatasetprep/tools/AbstractContextualPrompts.java similarity index 75% rename from src/main/java/de/extio/lmdatasetprep/preparer/AbstractContextualPrompts.java rename to src/main/java/de/extio/lmdatasetprep/tools/AbstractContextualPrompts.java index 88d2494..62847a6 100644 --- a/src/main/java/de/extio/lmdatasetprep/preparer/AbstractContextualPrompts.java +++ b/src/main/java/de/extio/lmdatasetprep/tools/AbstractContextualPrompts.java @@ -1,19 +1,17 @@ -package de.extio.lmdatasetprep.preparer; +package de.extio.lmdatasetprep.tools; import java.io.BufferedReader; import java.io.IOException; import java.io.InputStreamReader; import java.io.OutputStream; import java.nio.charset.StandardCharsets; -import java.nio.file.Files; -import java.nio.file.Path; import java.util.ArrayList; import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.Properties; import java.util.Random; import java.util.concurrent.ThreadLocalRandom; -import java.util.function.Consumer; import java.util.stream.Collectors; import org.slf4j.Logger; @@ -25,10 +23,13 @@ import org.springframework.core.io.ClassPathResource; import com.fasterxml.jackson.core.JsonProcessingException; import com.fasterxml.jackson.databind.ObjectMapper; +import de.extio.lmdatasetprep.Execution; +import de.extio.lmdatasetprep.Execution.WorkPacket; +import de.extio.lmdatasetprep.TextUtils; import de.extio.lmlib.client.ClientService; import de.extio.lmlib.profile.ModelCategory; -abstract class AbstractContextualPrompts implements InitializingBean, Consumer { +abstract class AbstractContextualPrompts implements InitializingBean, DatasetTool { private static final Logger LOGGER = LoggerFactory.getLogger(AbstractContextualPrompts.class); @@ -48,42 +49,35 @@ abstract class AbstractContextualPrompts implements InitializingBean, Consumer [ , ...]"); - return; - } - final int variations = Integer.parseInt(args[4]); + public void accept(final Properties properties) { + final int variations = Integer.parseInt(properties.getProperty("contextualPrompts.variations")); final List chunkConfigurations = new ArrayList<>(); - for (int i = 5; i < args.length; i += 2) { - final var cfg = new ChunkCfg(Integer.parseInt(args[i]), Integer.parseInt(args[i + 1])); + for (int i = 0; i < 99; i++) { + final String chunkNorm = properties.getProperty("contextualPrompts.chunkNorm." + i, ""); + if (chunkNorm.isBlank()) { + break; + } + final String chunkVar = properties.getProperty("contextualPrompts.chunkVar." + i, ""); + final var cfg = new ChunkCfg(Integer.parseInt(chunkNorm), Integer.parseInt(chunkVar)); chunkConfigurations.add(cfg); } - Utils.transformDirectory(args[1], this.createTasks(args, chunkConfigurations, variations)::createTasks); + Execution.transform(properties.getProperty("contextualPrompts.source"), this.createTasks(properties, chunkConfigurations, variations)::createTasks); } @FunctionalInterface static interface CreateTasks { - List createTasks(Path path); + List createTasks(WorkPacket packet); } - protected abstract CreateTasks createTasks(String[] args, List chunkConfigurations, int variations); + protected abstract CreateTasks createTasks(Properties properties, List chunkConfigurations, int variations); - protected List fileToParagraphs(final Path file, final int chunkNorm, final int chunkVar) { - LOGGER.info("Splitting to jsonl " + file + " -> " + chunkNorm + "±" + chunkVar); - - String text; - try { - text = Utils.normalizeText(Files.readString(file)); - } - catch (final IOException e) { - throw new RuntimeException("Cannot read file", e); - } - - return Utils.splitParagraphs(text, chunkNorm, chunkVar, false); + protected List fileToParagraphs(final WorkPacket packet, final int chunkNorm, final int chunkVar) { + LOGGER.info("Splitting to jsonl " + packet.file() + " -> " + chunkNorm + "±" + chunkVar); + final String text = TextUtils.normalizeText(packet.text()); + return TextUtils.splitParagraphs(text, chunkNorm, chunkVar, false); } protected Map createCharacterNameMapping(final String paragraph) { @@ -141,7 +135,7 @@ abstract class AbstractContextualPrompts implements InitializingBean, Consumer { + +} diff --git a/src/main/java/de/extio/lmdatasetprep/preparer/MergeJsonl.java b/src/main/java/de/extio/lmdatasetprep/tools/MergeJsonl.java similarity index 52% rename from src/main/java/de/extio/lmdatasetprep/preparer/MergeJsonl.java rename to src/main/java/de/extio/lmdatasetprep/tools/MergeJsonl.java index 193e448..1c1d1e7 100644 --- a/src/main/java/de/extio/lmdatasetprep/preparer/MergeJsonl.java +++ b/src/main/java/de/extio/lmdatasetprep/tools/MergeJsonl.java @@ -1,4 +1,4 @@ -package de.extio.lmdatasetprep.preparer; +package de.extio.lmdatasetprep.tools; import java.io.IOException; import java.nio.charset.StandardCharsets; @@ -7,33 +7,34 @@ import java.nio.file.Path; import java.util.ArrayList; import java.util.Collections; import java.util.List; +import java.util.Properties; import java.util.concurrent.ThreadLocalRandom; -import java.util.function.Consumer; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import org.springframework.stereotype.Component; +import de.extio.lmdatasetprep.Execution; + @Component -public class MergeJsonl implements Consumer, DatasetTool { +public class MergeJsonl implements DatasetTool { private static final byte[] NEWLINE = "\n".getBytes(StandardCharsets.UTF_8); private static final Logger LOGGER = LoggerFactory.getLogger(MergeJsonl.class); @Override - public void accept(final String[] args) { - if (args.length < 3) { - LOGGER.error("Arguments are missing. [ , ... ]"); - return; - } - + public void accept(final Properties properties) { final List lines = Collections.synchronizedList(new ArrayList<>()); - for (int i = 2; i < args.length; i++) { - Utils.transformDirectory(args[i], p -> { + for (int i = 0; i < 99; i++) { + final String source = properties.getProperty("mergeJsonl.source." + i, ""); + if (source.isBlank()) { + break; + } + Execution.transform(source, p -> { return List.of(() -> { try { - lines.addAll(Files.readAllLines(p)); + lines.addAll(Files.readAllLines(p.file())); } catch (final IOException e) { throw new RuntimeException("Cannot read file", e); @@ -45,16 +46,18 @@ public class MergeJsonl implements Consumer, DatasetTool { Collections.shuffle(lines, ThreadLocalRandom.current()); - final Path out = Path.of(args[1], "dataset-merged.jsonl"); - try (var fos = Files.newOutputStream(out)) { - for (final String line : lines) { - fos.write(line.getBytes(StandardCharsets.UTF_8)); - fos.write(NEWLINE); + 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); - } + catch (final IOException e1) { + LOGGER.error("IO exception", e1); + } + }); } } diff --git a/src/main/java/de/extio/lmdatasetprep/preparer/Rewrite.java b/src/main/java/de/extio/lmdatasetprep/tools/Rewrite.java similarity index 71% rename from src/main/java/de/extio/lmdatasetprep/preparer/Rewrite.java rename to src/main/java/de/extio/lmdatasetprep/tools/Rewrite.java index 72913e8..1d45ccc 100644 --- a/src/main/java/de/extio/lmdatasetprep/preparer/Rewrite.java +++ b/src/main/java/de/extio/lmdatasetprep/tools/Rewrite.java @@ -1,11 +1,11 @@ -package de.extio.lmdatasetprep.preparer; +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.List; +import java.util.Properties; import java.util.concurrent.ThreadLocalRandom; import java.util.concurrent.atomic.AtomicBoolean; import java.util.function.Consumer; @@ -15,11 +15,14 @@ import org.slf4j.LoggerFactory; import org.springframework.beans.factory.annotation.Autowired; import org.springframework.stereotype.Component; +import de.extio.lmdatasetprep.Execution; +import de.extio.lmdatasetprep.Execution.WorkPacket; +import de.extio.lmdatasetprep.TextUtils; import de.extio.lmlib.client.ClientService; import de.extio.lmlib.profile.ModelCategory; @Component -public class Rewrite implements Consumer, DatasetTool { +public class Rewrite implements DatasetTool { private static final Logger LOGGER = LoggerFactory.getLogger(Rewrite.class); @@ -50,40 +53,32 @@ public class Rewrite implements Consumer, DatasetTool { private ClientService clientService; @Override - public void accept(final String[] args) { - if (args.length < 5) { - LOGGER.error("Arguments are missing. "); - return; - } - - Utils.transformDirectory(args[1], - f -> { + public void accept(final Properties properties) { + Execution.transform(properties.getProperty("rewrite.source"), + packet -> { final List tasks = new ArrayList<>(); - tasks.add(() -> this.rewriteFile(args, f, 0, "impr", ModelCategory.COLD, IMPROVEMENT_PROMPT)); - for (int i = 0; i < Integer.parseInt(args[3]); i++) { + tasks.add(() -> this.rewriteFile(properties, packet, 0, "impr", ModelCategory.COLD, IMPROVEMENT_PROMPT)); + for (int i = 0; i < Integer.parseInt(properties.getProperty("rewrite.cnt")); i++) { final int fi = i; - tasks.add(() -> this.rewriteFile(args, f, fi, "enh", ModelCategory.COLD, String.format(ENHANCE_PROMPT, ENHANCEMENTS.get(ThreadLocalRandom.current().nextInt(ENHANCEMENTS.size()))))); + tasks.add(() -> this.rewriteFile(properties, packet, fi, "enh", ModelCategory.COLD, String.format(ENHANCE_PROMPT, ENHANCEMENTS.get(ThreadLocalRandom.current().nextInt(ENHANCEMENTS.size()))))); } return tasks; }); } - private void rewriteFile(final String[] args, final Path f, final int i, final String identifier, final ModelCategory modelCategory, final String prompt) { - final Path out = Utils.suffixFilename(Path.of(args[2]).resolve(f.getFileName()), + private void rewriteFile(final Properties properties, final WorkPacket packet, final int i, final String identifier, final ModelCategory modelCategory, final String prompt) { + final Path out = Execution.suffixFilename(packet.file(), "rewr", - args[4], + properties.getProperty("rewrite.model"), identifier, modelCategory.toString().toLowerCase(), String.valueOf(i)); - if (Files.exists(out)) { - LOGGER.info("Skipping " + out); - return; - } - LOGGER.info("Rewriting " + f.getFileName() + " " + identifier + " " + modelCategory.name() + " " + i); - Utils.streamOut(out, fos -> { + LOGGER.info("Rewriting " + packet.file().getFileName() + " " + identifier + " " + modelCategory.name() + " " + i); + + Execution.streamOut(out, "rewrite.destination", properties, fos -> { final AtomicBoolean first = new AtomicBoolean(true); - this.rewrite(f, prompt, modelCategory, chunk -> { + this.rewrite(packet, prompt, modelCategory, chunk -> { try { if (!first.getAndSet(false)) { fos.write(PARAGRAPH); @@ -97,17 +92,9 @@ public class Rewrite implements Consumer, DatasetTool { }); } - void rewrite(final Path file, final String prompt, final ModelCategory modelCategory, final Consumer consumer) { - String text; - try { - text = Utils.normalizeText(Files.readString(file)); - } - catch (final IOException e) { - throw new RuntimeException("Cannot read file", e); - } - - final List splits = Utils.splitParagraphs(text, 1250, 350, false); - + void rewrite(final WorkPacket packet, final String prompt, final ModelCategory modelCategory, final Consumer consumer) { + final String text = TextUtils.normalizeText(packet.text()); + final List splits = TextUtils.splitParagraphs(text, 1250, 350, false); for (final String split : splits) { LOGGER.info("Split " + (splits.indexOf(split) + 1) + "/" + splits.size()); final var completion = this.clientService.getClient(modelCategory).completion(modelCategory, diff --git a/src/main/java/de/extio/lmdatasetprep/preparer/Text2Jsonl.java b/src/main/java/de/extio/lmdatasetprep/tools/Text2Jsonl.java similarity index 57% rename from src/main/java/de/extio/lmdatasetprep/preparer/Text2Jsonl.java rename to src/main/java/de/extio/lmdatasetprep/tools/Text2Jsonl.java index 951a9bf..2de8923 100644 --- a/src/main/java/de/extio/lmdatasetprep/preparer/Text2Jsonl.java +++ b/src/main/java/de/extio/lmdatasetprep/tools/Text2Jsonl.java @@ -1,12 +1,11 @@ -package de.extio.lmdatasetprep.preparer; +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.List; -import java.util.function.Consumer; +import java.util.Properties; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -15,25 +14,24 @@ import org.springframework.stereotype.Component; import com.fasterxml.jackson.core.JsonProcessingException; import com.fasterxml.jackson.databind.ObjectMapper; +import de.extio.lmdatasetprep.Execution; +import de.extio.lmdatasetprep.Execution.WorkPacket; +import de.extio.lmdatasetprep.TextUtils; + @Component -public class Text2Jsonl implements Consumer, DatasetTool { +public class Text2Jsonl implements DatasetTool { private static final Logger LOGGER = LoggerFactory.getLogger(Text2Jsonl.class); private static final byte[] NEWLINE = "\n".getBytes(StandardCharsets.UTF_8); @Override - public void accept(final String[] args) { - if (args.length < 5) { - LOGGER.error("Arguments are missing. "); - return; - } - - final Path out = Path.of(args[2], "dataset.jsonl"); - try (var fos = Files.newOutputStream(out)) { - Utils.transformDirectory(args[1], p -> { + public void accept(final Properties properties) { + final Path out = Path.of("dataset.jsonl"); + Execution.streamOut(out, "text2Jsonl.destination", properties, fos -> { + Execution.transform(properties.getProperty("text2Jsonl.source"), p -> { return List.of(() -> { - final List splits = this.splitToJsonl(p, Integer.parseInt(args[3]), Integer.parseInt(args[4])); + final List splits = this.splitToJsonl(p, Integer.parseInt(properties.getProperty("text2Jsonl.chunkNorm")), Integer.parseInt(properties.getProperty("text2Jsonl.chunkVar"))); synchronized (this) { try { for (final String split : splits) { @@ -47,27 +45,18 @@ public class Text2Jsonl implements Consumer, DatasetTool { } }); }); - } - catch (final IOException e1) { - LOGGER.error("IO exception", e1); - } + }); } - List splitToJsonl(final Path file, final int chunkNorm, final int chunkVar) { - LOGGER.info("Splitting to jsonl " + file); + List splitToJsonl(final WorkPacket packet, final int chunkNorm, final int chunkVar) { + LOGGER.info("Splitting to jsonl " + packet.file()); - String text; - try { - text = Utils.normalizeText(Files.readString(file)); - } - catch (final IOException e) { - throw new RuntimeException("Cannot read file", e); - } + final String text = TextUtils.normalizeText(packet.text()); final List jsonl = new ArrayList<>(); final ObjectMapper mapper = new ObjectMapper(); - final List paragraphs = Utils.splitParagraphs(text, chunkNorm, chunkVar, true); + final List paragraphs = TextUtils.splitParagraphs(text, chunkNorm, chunkVar, true); for (int i = 0; i < paragraphs.size(); i++) { final String split = paragraphs.get(i); try { diff --git a/src/main/java/de/extio/lmdatasetprep/preparer/Text2Jsonl2HistoryInstructCompletion.java b/src/main/java/de/extio/lmdatasetprep/tools/Text2Jsonl2HistoryInstructCompletion.java similarity index 80% rename from src/main/java/de/extio/lmdatasetprep/preparer/Text2Jsonl2HistoryInstructCompletion.java rename to src/main/java/de/extio/lmdatasetprep/tools/Text2Jsonl2HistoryInstructCompletion.java index 69fb2cb..9f3d0d2 100644 --- a/src/main/java/de/extio/lmdatasetprep/preparer/Text2Jsonl2HistoryInstructCompletion.java +++ b/src/main/java/de/extio/lmdatasetprep/tools/Text2Jsonl2HistoryInstructCompletion.java @@ -1,9 +1,9 @@ -package de.extio.lmdatasetprep.preparer; +package de.extio.lmdatasetprep.tools; -import java.nio.file.Files; 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; @@ -11,13 +11,15 @@ import org.springframework.stereotype.Component; import com.fasterxml.jackson.databind.ObjectMapper; +import de.extio.lmdatasetprep.Execution; + @Component -public class Text2Jsonl2HistoryInstructCompletion extends AbstractContextualPrompts implements DatasetTool { +public class Text2Jsonl2HistoryInstructCompletion extends AbstractContextualPrompts { private static final Logger LOGGER = LoggerFactory.getLogger(Text2Jsonl2HistoryInstructCompletion.class); @Override - protected CreateTasks createTasks(final String[] args, final List chunkConfigurations, final int variations) { + protected CreateTasks createTasks(final Properties properties, final List chunkConfigurations, final int variations) { return p -> { final ObjectMapper mapper = new ObjectMapper(); final List tasks = new ArrayList<>(); @@ -26,19 +28,15 @@ public class Text2Jsonl2HistoryInstructCompletion extends AbstractContextualProm final List paragraphs = this.fileToParagraphs(p, chunkCfg.chunksNorm(), chunkCfg.chunksVar()); for (int i = 0; i < variations; i++) { - final Path out = Utils.suffixFilename(Path.of(args[2]).resolve(p.getFileName()), + final Path out = Execution.suffixFilename(p.file().getFileName(), "histinstrcompl", - args[3], + properties.getProperty("contextualPrompts.model"), String.valueOf(chunkCfg.chunksNorm()), String.valueOf(i), ".jsonl"); - if (Files.exists(out)) { - LOGGER.info("Skipping {}", out); - continue; - } tasks.add(() -> { - Utils.streamOut(out, fos -> { + Execution.streamOut(out, "contextualPrompts.destination", properties, fos -> { for (int j = 1; j < paragraphs.size(); j++) { LOGGER.info("Paragraph " + (j + 1) + "/" + paragraphs.size()); diff --git a/src/main/java/de/extio/lmdatasetprep/preparer/Text2Jsonl2WithContextualPrompts.java b/src/main/java/de/extio/lmdatasetprep/tools/Text2Jsonl2WithContextualPrompts.java similarity index 84% rename from src/main/java/de/extio/lmdatasetprep/preparer/Text2Jsonl2WithContextualPrompts.java rename to src/main/java/de/extio/lmdatasetprep/tools/Text2Jsonl2WithContextualPrompts.java index 31b9e8d..16473a7 100644 --- a/src/main/java/de/extio/lmdatasetprep/preparer/Text2Jsonl2WithContextualPrompts.java +++ b/src/main/java/de/extio/lmdatasetprep/tools/Text2Jsonl2WithContextualPrompts.java @@ -1,10 +1,9 @@ -package de.extio.lmdatasetprep.preparer; +package de.extio.lmdatasetprep.tools; -import java.io.IOException; -import java.nio.file.Files; 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; @@ -12,13 +11,15 @@ import org.springframework.stereotype.Component; import com.fasterxml.jackson.databind.ObjectMapper; +import de.extio.lmdatasetprep.Execution; + @Component -public class Text2Jsonl2WithContextualPrompts extends AbstractContextualPrompts implements DatasetTool { +public class Text2Jsonl2WithContextualPrompts extends AbstractContextualPrompts { private static final Logger LOGGER = LoggerFactory.getLogger(Text2Jsonl2WithContextualPrompts.class); @Override - protected CreateTasks createTasks(final String[] args, final List chunkConfigurations, final int variations) { + protected CreateTasks createTasks(final Properties properties, final List chunkConfigurations, final int variations) { return p -> { final ObjectMapper mapper = new ObjectMapper(); final List tasks = new ArrayList<>(); @@ -27,19 +28,15 @@ public class Text2Jsonl2WithContextualPrompts extends AbstractContextualPrompts final List paragraphs = this.fileToParagraphs(p, chunkCfg.chunksNorm(), chunkCfg.chunksVar()); for (int i = 0; i < variations; i++) { - final Path out = Utils.suffixFilename(Path.of(args[2]).resolve(p.getFileName()), + final Path out = Execution.suffixFilename(p.file().getFileName(), "contextualds", - args[3], + properties.getProperty("contextualPrompts.model"), String.valueOf(chunkCfg.chunksNorm()), String.valueOf(i), ".jsonl"); - if (Files.exists(out)) { - LOGGER.info("Skipping {}", out); - continue; - } tasks.add(() -> { - Utils.streamOut(out, fos -> { + Execution.streamOut(out, "contextualPrompts.destination", properties, fos -> { for (int j = 0; j < paragraphs.size(); j++) { LOGGER.info("Paragraph " + (j + 1) + "/" + paragraphs.size()); diff --git a/src/main/java/de/extio/lmdatasetprep/tools/Translate.java b/src/main/java/de/extio/lmdatasetprep/tools/Translate.java new file mode 100644 index 0000000..f4f0fc6 --- /dev/null +++ b/src/main/java/de/extio/lmdatasetprep/tools/Translate.java @@ -0,0 +1,72 @@ +package de.extio.lmdatasetprep.tools; + +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.nio.file.Path; +import java.util.List; +import java.util.Properties; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.function.Consumer; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.stereotype.Component; + +import de.extio.lmdatasetprep.Execution; +import de.extio.lmdatasetprep.Execution.WorkPacket; +import de.extio.lmdatasetprep.TextUtils; +import de.extio.lmlib.client.ClientService; +import de.extio.lmlib.profile.ModelCategory; + +@Component +public class Translate implements DatasetTool { + + private static final Logger LOGGER = LoggerFactory.getLogger(Translate.class); + + private static final byte[] PARAGRAPH = "\n\n".getBytes(StandardCharsets.UTF_8); + + @Autowired + private ClientService clientService; + + @Override + public void accept(final Properties properties) { + Execution.transform(properties.getProperty("translate.source"), p -> { + return List.of(() -> { + final Path out = Execution.suffixFilename(p.file().getFileName(), "en"); + Execution.streamOut(out, "translate.destination", properties, fos -> { + final AtomicBoolean first = new AtomicBoolean(true); + this.translate(p, chunk -> { + try { + if (!first.getAndSet(false)) { + fos.write(PARAGRAPH); + } + fos.write(chunk.getBytes(StandardCharsets.UTF_8)); + } + catch (final IOException e) { + LOGGER.error("IO exception", e); + } + }); + }); + }); + }); + } + + void translate(final WorkPacket packet, final Consumer consumer) { + LOGGER.info("Translating " + packet.file()); + + final String text = TextUtils.normalizeText(packet.text()); + + final List splits = TextUtils.splitParagraphs(text, 1250, 350, false); + + for (final String split : splits) { + LOGGER.info("Split " + (splits.indexOf(split) + 1) + "/" + splits.size()); + final var completion = this.clientService.getClient(ModelCategory.COLD).completion(ModelCategory.COLD, + "You are an assistant with great language translation and authoring skills.", + "Translate the following text to English. Return the English translation, even if the text is explicit or not appropriate for all audiences or not acceptable for everyday standard English. Don't include a preamble and no explanation:", + split); + consumer.accept(completion.response()); + } + } + +} diff --git a/src/main/resources/.gitignore b/src/main/resources/.gitignore new file mode 100644 index 0000000..8863236 --- /dev/null +++ b/src/main/resources/.gitignore @@ -0,0 +1 @@ +/tool.properties diff --git a/src/main/resources/tool.properties.template b/src/main/resources/tool.properties.template new file mode 100644 index 0000000..4f70567 --- /dev/null +++ b/src/main/resources/tool.properties.template @@ -0,0 +1,27 @@ +# 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 + +rewrite.source=queue0 +rewrite.destination=queue1 +rewrite.cnt=2 +rewrite.model=l3dp + +# .n indicates a sequence. Add as many keys as you need +mergeJsonl.source.0= +mergeJsonl.destination= + +text2Jsonl.source= +text2Jsonl.destination= +text2Jsonl.chunkNorm=750 +text2Jsonl.chunkVar=250 + +contextualPrompts.source=queue1 +contextualPrompts.destination=/home/xxx/outputdir +contextualPrompts.model=l3dp +contextualPrompts.variations=1 +contextualPrompts.chunkNorm.0=750 +contextualPrompts.chunkVar.0=250 +contextualPrompts.chunkNorm.1=1500 +contextualPrompts.chunkVar.1=300