exo/lm-dataset-preparer
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Update Client.java

Browse Source
This commit is contained in:
stephan-birkl
2024-11-01 17:43:40 +01:00
committed by GitHub
parent e40e648557
commit 3b6a0ee74a
1 changed files with 7 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
src/main/java/de/extio/lmdatasetprep/client/Client.java
View File

@@ -5,6 +5,8 @@ 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 org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -46,7 +48,7 @@ public class Client {
@Value("${client.collectStatistics}")
private boolean collectStatistics;
private Integer promptTemplateTokenLength;
private Map<String, Integer> promptTemplateTokenLengths = new HashMap<>();
public Completion completion(final String instruction, final String question, final String fullText, final ModelCategory modelCategory) {
final var modelProfile = this.modelProfileService.getModelProfile(modelCategory.getModelProfile());
@@ -59,7 +61,7 @@ public class Client {
if (inputTokens > modelProfile.maxContextLength() - modelProfile.maxTokens()) {
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, i + steps);
List<Long> split = tokenizedFullText.subList(i, Math.min(i + steps, tokenizedFullText.size()));
texts.add(this.tokenizer.detokenize(split, modelProfile));
}
}
@@ -163,12 +165,11 @@ public class Client {
}
private int getPromptTemplateTokenLength(final ModelProfile modelProfile) {
if (this.promptTemplateTokenLength == null) {
return this.promptTemplateTokenLengths.computeIfAbsent(modelProfile.prompt(), key -> {
final PromptStrategy promptStrategy = this.promptStrategyFactory.getStrategy(modelProfile.prompt());
var prompt = promptStrategy.start(" ", " ", " ");
this.promptTemplateTokenLength = this.tokenizer.count(prompt.toString(), modelProfile);
}
return this.promptTemplateTokenLength;
return this.tokenizer.count(prompt.toString(), modelProfile);
});
}
static class CompletionStatistics {