From efb520fc072ce0d557c0c59ba08c806648aa3d65 Mon Sep 17 00:00:00 2001 From: extio-de Date: Tue, 29 Oct 2024 23:51:41 +0100 Subject: [PATCH] New generator: Jsonl QA style contextual prompts; JsonL: Add sliding windows; Streamed jsonl output; Memory optimization: Buffered, blocking task creation; --- .../lmdatasetprep/preparer/Text2Jsonl.java | 70 +++++++----- .../Text2Jsonl2WithContextualPrompts.java | 100 ++++++++++++++++++ .../extio/lmdatasetprep/preparer/Utils.java | 56 ++++++++-- 3 files changed, 191 insertions(+), 35 deletions(-) create mode 100644 src/main/java/de/extio/lmdatasetprep/preparer/Text2Jsonl2WithContextualPrompts.java diff --git a/src/main/java/de/extio/lmdatasetprep/preparer/Text2Jsonl.java b/src/main/java/de/extio/lmdatasetprep/preparer/Text2Jsonl.java index 026b348..8dcd4c5 100644 --- a/src/main/java/de/extio/lmdatasetprep/preparer/Text2Jsonl.java +++ b/src/main/java/de/extio/lmdatasetprep/preparer/Text2Jsonl.java @@ -1,10 +1,10 @@ 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.ArrayList; -import java.util.Collections; import java.util.List; import java.util.function.Consumer; @@ -20,6 +20,8 @@ public class Text2Jsonl implements Consumer { 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) { @@ -27,20 +29,27 @@ public class Text2Jsonl implements Consumer { return; } - final List lines = Collections.synchronizedList(new ArrayList<>()); - - Utils.transformDirectory(args[1], p -> { - return List.of(() -> { - lines.addAll(this.splitToJsonl(p, Integer.parseInt(args[3]), Integer.parseInt(args[4]))); - }); - }); - final Path out = Path.of(args[2], "dataset.jsonl"); - try { - Files.writeString(out, String.join("\n", lines)); + try (var fos = Files.newOutputStream(out)) { + Utils.transformDirectory(args[1], p -> { + return List.of(() -> { + final List splits = this.splitToJsonl(p, Integer.parseInt(args[3]), Integer.parseInt(args[4])); + synchronized (this) { + try { + for (final String split : splits) { + fos.write(split.getBytes(StandardCharsets.UTF_8)); + fos.write(NEWLINE); + } + } + catch (final IOException e) { + LOGGER.error("IO exception", e); + } + } + }); + }); } - catch (final IOException e) { - LOGGER.error("Cannot store dataset", e); + catch (final IOException e1) { + LOGGER.error("IO exception", e1); } } @@ -55,21 +64,30 @@ public class Text2Jsonl implements Consumer { throw new RuntimeException("Cannot read file", e); } + final List jsonl = new ArrayList<>(); final ObjectMapper mapper = new ObjectMapper(); - return Utils - .splitParagraphs(text, chunkNorm, chunkVar) - .stream() - .map(split -> { - try { - return mapper.writeValueAsString(new TextLine(split)); - } - catch (final JsonProcessingException e) { - LOGGER.error("Cannot convert split to json", e); - return ""; - } - }) - .toList(); + final List paragraphs = Utils.splitParagraphs(text, chunkNorm, chunkVar); + for (int i = 0; i < paragraphs.size(); i++) { + String split = paragraphs.get(i); + + if (i > 0) { // Add sliding window + final String prevSplit = paragraphs.get(i - 1); + final List prevSplitted = Utils.splitParagraphs(prevSplit, prevSplit.length() - chunkVar, chunkVar - 1); + if (prevSplitted.size() > 1) { + split = prevSplitted.getLast() + "\n\n" + split; + } + } + + try { + jsonl.add(mapper.writeValueAsString(new TextLine(split))); + } + catch (final JsonProcessingException e) { + LOGGER.error("Cannot convert split to json", e); + } + } + + return jsonl; } record TextLine(String text) { diff --git a/src/main/java/de/extio/lmdatasetprep/preparer/Text2Jsonl2WithContextualPrompts.java b/src/main/java/de/extio/lmdatasetprep/preparer/Text2Jsonl2WithContextualPrompts.java new file mode 100644 index 0000000..d182f5d --- /dev/null +++ b/src/main/java/de/extio/lmdatasetprep/preparer/Text2Jsonl2WithContextualPrompts.java @@ -0,0 +1,100 @@ +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.ArrayList; +import java.util.List; +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 com.fasterxml.jackson.databind.ObjectMapper; + +import de.extio.lmdatasetprep.client.Client; +import de.extio.lmdatasetprep.client.profile.ModelCategory; + +@Component +public class Text2Jsonl2WithContextualPrompts implements Consumer { + + private static final byte[] NEWLINE = "\n".getBytes(StandardCharsets.UTF_8); + + private static final Logger LOGGER = LoggerFactory.getLogger(Text2Jsonl2WithContextualPrompts.class); + + @Autowired + private Client client; + + @Override + public void accept(final String[] args) { + if (args.length < 6) { + LOGGER.error("Arguments are missing. "); + return; + } + final int chunksNorm = Integer.parseInt(args[3]); + final int chunksVar = Integer.parseInt(args[4]); + final int variations = Integer.parseInt(args[5]); + + final ObjectMapper mapper = new ObjectMapper(); + + final Path out = Path.of(args[2], "dataset_with_context_" + chunksNorm + ".jsonl"); + try (var fos = Files.newOutputStream(out)) { + Utils.transformDirectory(args[1], p -> { + final List tasks = new ArrayList<>(); + final List paragraphs = this.fileToParagraphs(p, chunksNorm, chunksVar); + for (final String paragraph : paragraphs) { + for (int i = 0; i < variations; i++) { + final int fI = i; + tasks.add(() -> { + + LOGGER.info("### TASK ### " + paragraphs.indexOf(paragraph) + " " + fI + " " + paragraphs.size()); + + final var completion = this.client.completion("You are an assistant with great authoring skills.", + "Enhance the given text by crafting a prompt that effectively provides additional context, enabling the language model to better understand the subject matter and generate more informative and relevant responses. This prompt should be clear, concise, and tailored to facilitate contextual learning, allowing the model to grasp the essential details and nuances of the topic. Return the prompt, 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:", + paragraph, + ModelCategory.COLD); + final var qaLine = new QaLine(completion.response().replace("Prompt:", ""), paragraph); + try { + final var jsonb = mapper + .writeValueAsString(qaLine) + .getBytes(StandardCharsets.UTF_8); + synchronized (this) { + fos.write(jsonb); + fos.write(NEWLINE); + } + } + catch (final IOException e) { + LOGGER.error("Cannot convert split to json", e); + } + }); + } + } + return tasks; + }); + } + catch (final IOException e1) { + LOGGER.error("IO exception", e1); + } + } + + List fileToParagraphs(final Path file, final int chunkNorm, final int chunkVar) { + LOGGER.info("Splitting to jsonl " + file); + + 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); + } + + record QaLine(String question, String answer) { + } + +} diff --git a/src/main/java/de/extio/lmdatasetprep/preparer/Utils.java b/src/main/java/de/extio/lmdatasetprep/preparer/Utils.java index 1e637d9..68682e2 100644 --- a/src/main/java/de/extio/lmdatasetprep/preparer/Utils.java +++ b/src/main/java/de/extio/lmdatasetprep/preparer/Utils.java @@ -7,9 +7,10 @@ import java.nio.file.Paths; import java.util.ArrayList; import java.util.Collections; import java.util.List; -import java.util.concurrent.CompletableFuture; -import java.util.concurrent.ExecutorService; -import java.util.concurrent.Executors; +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.Function; import java.util.stream.Stream; @@ -105,19 +106,56 @@ public class Utils { } static void transformDirectory(final String path, final Function> createTasks) { - try (final ExecutorService executor = Executors.newFixedThreadPool(4); - Stream stream = Files.list(Paths.get(path))) { - final List> tasks = new ArrayList<>(); + final int EXECUTORS = 4; + + try (Stream stream = Files.list(Paths.get(path))) { + 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); + } + stream .filter(p -> Files.isRegularFile(p)) .forEach(p -> { for (final Runnable r : createTasks.apply(p)) { - tasks.add(CompletableFuture.runAsync(r, executor)); + try { + tasks.put(r); + } + catch (final InterruptedException e) { + throw new RuntimeException("Interrupted", e); + } } }); - CompletableFuture.allOf(tasks.toArray(new CompletableFuture[0])).join(); + + while (!tasks.isEmpty()) { + Thread.sleep(1); + } + finished.set(true); + + for (int i = 0; i < EXECUTORS; i++) { + threads.get(i).join(); + } } - catch (final IOException e) { + catch (final IOException | InterruptedException e) { LOGGER.error("Error", e); } }