Client: Collect and show more statistics; MergeJsonl: Shuffle merged json; Rewrite: Fix progress indicator;

src/main/java/de/extio/lmdatasetprep/client/Client.java

@@ -10,6 +10,7 @@ import java.util.stream.Collectors;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import org.springframework.beans.factory.DisposableBean;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.beans.factory.annotation.Value;
 import org.springframework.http.HttpMethod;
@@ -26,7 +27,7 @@ import de.extio.lmdatasetprep.client.prompt.PromptStrategy;
 import de.extio.lmdatasetprep.client.prompt.PromptStrategyFactory;
 
 @Component
-public class Client {
+public class Client implements DisposableBean {
 	
 	private static final Logger LOGGER = LoggerFactory.getLogger(Client.class);
 	
@@ -50,6 +51,8 @@ public class Client {
 	
 	private final Map<String, Integer> promptTemplateTokenLengths = new ConcurrentHashMap<>();
 	
+	private final CompletionStatistics totalStatistics = new CompletionStatistics();
+	
 	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());
@@ -129,7 +132,11 @@ public class Client {
 			}
 			
 			if (this.collectStatistics) {
-				statistics.add(Duration.between(start, LocalDateTime.now()), this.tokenizer.count(promptStr, modelProfile), this.tokenizer.count(content, modelProfile));
+				final var dur = Duration.between(start, LocalDateTime.now());
+				final var in = this.tokenizer.count(promptStr, modelProfile);
+				final var out = this.tokenizer.count(content, modelProfile);
+				statistics.add(dur, in, out);
+				this.totalStatistics.add(dur, in, out);
 			}
 			
 			answer.append(content);
@@ -172,17 +179,24 @@ public class Client {
 		});
 	}
 	
+	@Override
+	public void destroy() throws Exception {
+		if (this.collectStatistics) {
+			LOGGER.info("Client Statistics: {}", this.totalStatistics);
+		}
+	}
+	
 	static class CompletionStatistics {
 		
-		int requests;
+		volatile int requests;
 		
-		Duration duration = Duration.ofMillis(0l);
+		volatile Duration duration = Duration.ofMillis(0l);
 		
-		long inTokens;
+		volatile long inTokens;
 		
-		long outTokens;
+		volatile long outTokens;
 		
-		void add(final Duration duration, final int inTokens, final int outTokens) {
+		synchronized void add(final Duration duration, final int inTokens, final int outTokens) {
 			LOGGER.debug("Request duration: " + duration + "; in tokens: " + inTokens + "; out tokens: " + outTokens);
 			
 			this.requests++;
@@ -190,6 +204,22 @@ public class Client {
 			this.inTokens += inTokens;
 			this.outTokens += outTokens;
 		}
+		
+		@Override
+		public String toString() {
+			final StringBuilder builder = new StringBuilder();
+			builder.append("CompletionStatistics [requests=");
+			builder.append(this.requests);
+			builder.append(", duration=");
+			builder.append(this.duration);
+			builder.append(", inTokens=");
+			builder.append(this.inTokens);
+			builder.append(", outTokens=");
+			builder.append(this.outTokens);
+			builder.append("]");
+			return builder.toString();
+		}
+		
 	}
 	
 }

src/main/java/de/extio/lmdatasetprep/preparer/MergeJsonl.java

@@ -1,8 +1,11 @@
 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;
 
@@ -10,9 +13,13 @@ import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.springframework.stereotype.Component;
 
+import de.extio.lmdatasetprep.XorShift128Random;
+
 @Component
 public class MergeJsonl implements Consumer<String[]> {
 	
+	private static final byte[] NEWLINE = "\n".getBytes(StandardCharsets.UTF_8);
+	
 	private static final Logger LOGGER = LoggerFactory.getLogger(MergeJsonl.class);
 	
 	@Override
@@ -22,22 +29,27 @@ public class MergeJsonl implements Consumer<String[]> {
 			return;
 		}
 		
+		final List<String> lines = Collections.synchronizedList(new ArrayList<>());
+		Utils.transformDirectory(args[1], p -> {
+			return List.of(() -> {
+				try {
+					lines.addAll(Files.readAllLines(p));
+				}
+				catch (final IOException e) {
+					throw new RuntimeException("Cannot read file", e);
+				}
+				
+			});
+		});
+		
+		Collections.shuffle(lines, new XorShift128Random());
+		
 		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);
-					}
-					
-				});
-			});
+			for (final String line : lines) {
+				fos.write(line.getBytes(StandardCharsets.UTF_8));
+				fos.write(NEWLINE);
+			}
 		}
 		catch (final IOException e1) {
 			LOGGER.error("IO exception", e1);

src/main/java/de/extio/lmdatasetprep/preparer/Rewrite.java

@@ -115,7 +115,7 @@ public class Rewrite implements Consumer<String[]> {
 		final List<String> splits = Utils.splitParagraphs(text, 1250, 350, false);
 		
 		for (final String split : splits) {
-			LOGGER.info("Split " + splits.indexOf(split) + "/" + splits.size());
+			LOGGER.info("Split " + (splits.indexOf(split) + 1) + "/" + splits.size());
 			final var completion = this.client.completion("You are a helpful assistant with great authoring skills.",
 					prompt,
 					split,
@@ -123,5 +123,4 @@ public class Rewrite implements Consumer<String[]> {
 			consumer.accept(completion.response());
 		}
 	}
-	
 }