Client refactoring; MergeJsonl tool;

This commit is contained in:
extio-de
2024-11-01 19:41:03 +01:00
parent 94a20bd4ed
commit 1813ab8928
2 changed files with 57 additions and 10 deletions

View File

@@ -4,9 +4,9 @@ import java.time.Duration;
import java.time.LocalDateTime; import java.time.LocalDateTime;
import java.util.ArrayList; import java.util.ArrayList;
import java.util.List; import java.util.List;
import java.util.stream.Collectors;
import java.util.Map; 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.Logger;
import org.slf4j.LoggerFactory; import org.slf4j.LoggerFactory;
@@ -48,20 +48,20 @@ public class Client {
@Value("${client.collectStatistics}") @Value("${client.collectStatistics}")
private boolean 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) { public Completion completion(final String instruction, final String question, final String fullText, final ModelCategory modelCategory) {
final var modelProfile = this.modelProfileService.getModelProfile(modelCategory.getModelProfile()); final var modelProfile = this.modelProfileService.getModelProfile(modelCategory.getModelProfile());
final PromptStrategy promptStrategy = this.promptStrategyFactory.getStrategy(modelProfile.prompt()); final PromptStrategy promptStrategy = this.promptStrategyFactory.getStrategy(modelProfile.prompt());
final List<String> texts = new ArrayList<>(); final List<String> texts = new ArrayList<>();
List<Long> tokenizedInstrAndQuestion = this.tokenizer.tokenize(instruction + question, modelProfile); final List<Long> tokenizedInstrAndQuestion = this.tokenizer.tokenize(instruction + question, modelProfile);
List<Long> tokenizedFullText = this.tokenizer.tokenize(fullText, modelProfile); final List<Long> tokenizedFullText = this.tokenizer.tokenize(fullText, modelProfile);
int inputTokens = tokenizedInstrAndQuestion.size() + tokenizedFullText.size() + this.getPromptTemplateTokenLength(modelProfile); final int inputTokens = tokenizedInstrAndQuestion.size() + tokenizedFullText.size() + this.getPromptTemplateTokenLength(modelProfile);
if (inputTokens > modelProfile.maxContextLength() - modelProfile.maxTokens()) { 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) { 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)); texts.add(this.tokenizer.detokenize(split, modelProfile));
} }
} }
@@ -167,7 +167,7 @@ public class Client {
private int getPromptTemplateTokenLength(final ModelProfile modelProfile) { private int getPromptTemplateTokenLength(final ModelProfile modelProfile) {
return this.promptTemplateTokenLengths.computeIfAbsent(modelProfile.prompt(), key -> { return this.promptTemplateTokenLengths.computeIfAbsent(modelProfile.prompt(), key -> {
final PromptStrategy promptStrategy = this.promptStrategyFactory.getStrategy(modelProfile.prompt()); final PromptStrategy promptStrategy = this.promptStrategyFactory.getStrategy(modelProfile.prompt());
var prompt = promptStrategy.start(" ", " ", " "); final var prompt = promptStrategy.start(" ", " ", " ");
return this.tokenizer.count(prompt.toString(), modelProfile); return this.tokenizer.count(prompt.toString(), modelProfile);
}); });
} }

View File

@@ -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);
}
}
}