From 1813ab89285f3a87ffd895ed993a95daa4c286bc Mon Sep 17 00:00:00 2001
From: extio-de <github@extio.de>
Date: Fri, 1 Nov 2024 19:41:03 +0100
Subject: [PATCH] Client refactoring; MergeJsonl tool;

---
 .../de/extio/lmdatasetprep/client/Client.java | 20 ++++----
 .../lmdatasetprep/preparer/MergeJsonl.java    | 47 +++++++++++++++++++
 2 files changed, 57 insertions(+), 10 deletions(-)
 create mode 100644 src/main/java/de/extio/lmdatasetprep/preparer/MergeJsonl.java

diff --git a/src/main/java/de/extio/lmdatasetprep/client/Client.java b/src/main/java/de/extio/lmdatasetprep/client/Client.java
index 9d42df4..0887079 100755
--- a/src/main/java/de/extio/lmdatasetprep/client/Client.java
+++ b/src/main/java/de/extio/lmdatasetprep/client/Client.java
@@ -4,9 +4,9 @@ import java.time.Duration;
 import java.time.LocalDateTime;
 import java.util.ArrayList;
 import java.util.List;
-import java.util.stream.Collectors;
 import java.util.Map;
-import java.util.HashMap;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.stream.Collectors;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -48,20 +48,20 @@ public class Client {
 	@Value("${client.collectStatistics}")
 	private boolean collectStatistics;
 	
-	private Map<String, Integer> promptTemplateTokenLengths = new HashMap<>();
+	private final Map<String, Integer> promptTemplateTokenLengths = new ConcurrentHashMap<>();
 	
 	public Completion completion(final String instruction, final String question, final String fullText, final ModelCategory modelCategory) {
 		final var modelProfile = this.modelProfileService.getModelProfile(modelCategory.getModelProfile());
 		final PromptStrategy promptStrategy = this.promptStrategyFactory.getStrategy(modelProfile.prompt());
-
+		
 		final List<String> texts = new ArrayList<>();
-		List<Long> tokenizedInstrAndQuestion = this.tokenizer.tokenize(instruction + question, modelProfile);
-		List<Long> tokenizedFullText = this.tokenizer.tokenize(fullText, modelProfile);
-		int inputTokens = tokenizedInstrAndQuestion.size() + tokenizedFullText.size() + this.getPromptTemplateTokenLength(modelProfile);
+		final List<Long> tokenizedInstrAndQuestion = this.tokenizer.tokenize(instruction + question, modelProfile);
+		final List<Long> tokenizedFullText = this.tokenizer.tokenize(fullText, modelProfile);
+		final int inputTokens = tokenizedInstrAndQuestion.size() + tokenizedFullText.size() + this.getPromptTemplateTokenLength(modelProfile);
 		if (inputTokens > modelProfile.maxContextLength() - modelProfile.maxTokens()) {
-			int steps = modelProfile.maxContextLength() - modelProfile.maxTokens() - tokenizedInstrAndQuestion.size() - this.getPromptTemplateTokenLength(modelProfile);
+			final int steps = modelProfile.maxContextLength() - modelProfile.maxTokens() - tokenizedInstrAndQuestion.size() - this.getPromptTemplateTokenLength(modelProfile);
 			for (int i = 0; i < tokenizedFullText.size(); i += steps) {
-				List<Long> split = tokenizedFullText.subList(i, Math.min(i + steps, tokenizedFullText.size()));
+				final List<Long> split = tokenizedFullText.subList(i, Math.min(i + steps, tokenizedFullText.size()));
 				texts.add(this.tokenizer.detokenize(split, modelProfile));
 			}
 		}
@@ -167,7 +167,7 @@ public class Client {
 	private int getPromptTemplateTokenLength(final ModelProfile modelProfile) {
 		return this.promptTemplateTokenLengths.computeIfAbsent(modelProfile.prompt(), key -> {
 			final PromptStrategy promptStrategy = this.promptStrategyFactory.getStrategy(modelProfile.prompt());
-			var prompt = promptStrategy.start(" ", " ", " ");
+			final var prompt = promptStrategy.start(" ", " ", " ");
 			return this.tokenizer.count(prompt.toString(), modelProfile);
 		});
 	}
diff --git a/src/main/java/de/extio/lmdatasetprep/preparer/MergeJsonl.java b/src/main/java/de/extio/lmdatasetprep/preparer/MergeJsonl.java
new file mode 100644
index 0000000..3b9fa29
--- /dev/null
+++ b/src/main/java/de/extio/lmdatasetprep/preparer/MergeJsonl.java
@@ -0,0 +1,47 @@
+package de.extio.lmdatasetprep.preparer;
+
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.List;
+import java.util.function.Consumer;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.springframework.stereotype.Component;
+
+@Component
+public class MergeJsonl implements Consumer<String[]> {
+	
+	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. <Path to jsonl files> <Path to output dir>");
+			return;
+		}
+		
+		final Path out = Path.of(args[2], "dataset-merged.jsonl");
+		try (var fos = Files.newOutputStream(out)) {
+			Utils.transformDirectory(args[1], p -> {
+				return List.of(() -> {
+					try {
+						final byte[] b = Files.readAllBytes(p);
+						synchronized (fos) {
+							fos.write(b);
+						}
+					}
+					catch (final IOException e) {
+						throw new RuntimeException("Cannot read file", e);
+					}
+					
+				});
+			});
+		}
+		catch (final IOException e1) {
+			LOGGER.error("IO exception", e1);
+		}
+	}
+	
+}