91 lines
2.9 KiB
Java
Executable File
91 lines
2.9 KiB
Java
Executable File
package de.extio.lmlib.token;
|
|
|
|
import java.io.InputStream;
|
|
import java.nio.charset.StandardCharsets;
|
|
import java.util.List;
|
|
|
|
import org.slf4j.Logger;
|
|
import org.slf4j.LoggerFactory;
|
|
import org.springframework.http.HttpMethod;
|
|
import org.springframework.web.client.RestClient;
|
|
|
|
import de.extio.lmlib.profile.ModelProfile;
|
|
|
|
final class LlamaServerTokenizer implements Tokenizer {
|
|
|
|
static final String NAME = "llamaServer";
|
|
|
|
private static final Logger LOGGER = LoggerFactory.getLogger(LlamaServerTokenizer.class);
|
|
|
|
private final RestClient.Builder restClientBuilder;
|
|
|
|
public LlamaServerTokenizer(final RestClient.Builder restClientBuilder) {
|
|
this.restClientBuilder = restClientBuilder;
|
|
}
|
|
|
|
@Override
|
|
public String getName() {
|
|
return NAME;
|
|
}
|
|
|
|
@Override
|
|
public List<Long> tokenize(final String txt, final ModelProfile modelProfile) {
|
|
LOGGER.debug("Requesting tokenizing");
|
|
final var request = new TokenizeRequest(txt);
|
|
final var restClient = this.restClientBuilder.baseUrl(modelProfile.url()).build();
|
|
final var response = restClient
|
|
.method(HttpMethod.POST)
|
|
.uri("/tokenize")
|
|
.body(request)
|
|
.retrieve()
|
|
.onStatus(status -> {
|
|
if (status.getStatusCode().isError()) {
|
|
LOGGER.error("Tokenization request failed with status code {} {}: {}", status.getStatusCode(), status.getStatusText(), new String(status.getBody().readAllBytes(), StandardCharsets.UTF_8));
|
|
throw new RuntimeException("Tokenization request failed with status code " + status.getStatusCode() + " " + status.getStatusText());
|
|
}
|
|
return false;
|
|
})
|
|
.body(TokenizeResponse.class);
|
|
return response.tokens();
|
|
}
|
|
|
|
@Override
|
|
public int count(final String txt, final ModelProfile modelProfile) {
|
|
return this.tokenize(txt, modelProfile).size();
|
|
}
|
|
|
|
@Override
|
|
public String detokenize(final List<Long> tokens, final ModelProfile modelProfile) {
|
|
LOGGER.debug("Requesting detokenizing");
|
|
final var request = new DetokenizeRequest(tokens);
|
|
final var restClient = this.restClientBuilder.baseUrl(modelProfile.url()).build();
|
|
final var response = restClient
|
|
.method(HttpMethod.POST)
|
|
.uri("/detokenize")
|
|
.body(request)
|
|
.retrieve()
|
|
.onStatus(status -> {
|
|
if (status.getStatusCode().isError()) {
|
|
LOGGER.error("Detokenization request failed with status code {} {}: {}", status.getStatusCode(), status.getStatusText(), new String(status.getBody().readAllBytes(), StandardCharsets.UTF_8));
|
|
throw new RuntimeException("Detokenization request failed with status code " + status.getStatusCode() + " " + status.getStatusText());
|
|
}
|
|
return false;
|
|
})
|
|
.body(DetokenizeResponse.class);
|
|
return response.content();
|
|
}
|
|
|
|
private record TokenizeRequest(String content) {
|
|
}
|
|
|
|
private record TokenizeResponse(List<Long> tokens) {
|
|
}
|
|
|
|
private record DetokenizeRequest(List<Long> tokens) {
|
|
}
|
|
|
|
private record DetokenizeResponse(String content) {
|
|
}
|
|
|
|
}
|