diff --git a/src/main/java/de/extio/lmdatasetprep/client/Client.java b/src/main/java/de/extio/lmdatasetprep/client/Client.java index 0887079..259adac 100755 --- a/src/main/java/de/extio/lmdatasetprep/client/Client.java +++ b/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 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(); + } + } } diff --git a/src/main/java/de/extio/lmdatasetprep/preparer/MergeJsonl.java b/src/main/java/de/extio/lmdatasetprep/preparer/MergeJsonl.java index 3b9fa29..5d75b64 100644 --- a/src/main/java/de/extio/lmdatasetprep/preparer/MergeJsonl.java +++ b/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 { + 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 { return; } + final List 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); diff --git a/src/main/java/de/extio/lmdatasetprep/preparer/Rewrite.java b/src/main/java/de/extio/lmdatasetprep/preparer/Rewrite.java index fb95982..2598a23 100644 --- a/src/main/java/de/extio/lmdatasetprep/preparer/Rewrite.java +++ b/src/main/java/de/extio/lmdatasetprep/preparer/Rewrite.java @@ -115,7 +115,7 @@ public class Rewrite implements Consumer { final List 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 { consumer.accept(completion.response()); } } - } diff --git a/src/main/resources/application.properties b/src/main/resources/application.properties index 254c28f..5baa306 100755 --- a/src/main/resources/application.properties +++ b/src/main/resources/application.properties @@ -5,4 +5,4 @@ logging.level.de.extio=info profile.hot=llama3-hot profile.cold=llama3-cold -client.collectStatistics=false +client.collectStatistics=true