1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
|
// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
package com.yahoo.language.huggingface;
import com.yahoo.api.annotations.Beta;
import com.yahoo.component.AbstractComponent;
import com.yahoo.component.annotation.Inject;
import com.yahoo.language.Language;
import com.yahoo.language.process.Embedder;
import com.yahoo.language.process.Segmenter;
import com.yahoo.language.tools.Embed;
import com.yahoo.tensor.Tensor;
import com.yahoo.tensor.TensorType;
import java.nio.file.Path;
import java.util.Arrays;
import java.util.Collection;
import java.util.EnumMap;
import java.util.List;
import java.util.Map;
import static com.yahoo.yolean.Exceptions.uncheck;
/**
* {@link Embedder}/{@link Segmenter} using Deep Java Library's HuggingFace Tokenizer.
*
* @author bjorncs
*/
@Beta
public class HuggingFaceTokenizer extends AbstractComponent implements Embedder, Segmenter, AutoCloseable {
private final Map<Language, ai.djl.huggingface.tokenizers.HuggingFaceTokenizer> models = new EnumMap<>(Language.class);
@Inject public HuggingFaceTokenizer(HuggingFaceTokenizerConfig cfg) { this(new Builder(cfg)); }
private HuggingFaceTokenizer(Builder b) {
var original = Thread.currentThread().getContextClassLoader();
Thread.currentThread().setContextClassLoader(HuggingFaceTokenizer.class.getClassLoader());
try {
b.models.forEach((language, path) -> {
models.put(language,
uncheck(() -> ai.djl.huggingface.tokenizers.HuggingFaceTokenizer.builder()
.optTokenizerPath(path)
.optAddSpecialTokens(b.addSpecialTokens != null ? b.addSpecialTokens : true)
.build()));
});
} finally {
Thread.currentThread().setContextClassLoader(original);
}
}
@Override
public List<Integer> embed(String text, Context ctx) {
var encoding = resolve(ctx.getLanguage()).encode(text);
return Arrays.stream(encoding.getIds()).mapToInt(Math::toIntExact).boxed().toList();
}
@Override
public Tensor embed(String text, Context ctx, TensorType type) {
return Embed.asTensor(text, this, ctx, type);
}
@Override
public List<String> segment(String input, Language language) {
return List.of(resolve(language).encode(input).getTokens());
}
@Override
public String decode(List<Integer> tokens, Context ctx) {
return resolve(ctx.getLanguage()).decode(toArray(tokens));
}
public Encoding encode(String text) { return encode(text, Language.UNKNOWN); }
public Encoding encode(String text, Language language) { return Encoding.from(resolve(language).encode(text)); }
public String decode(List<Long> tokens) { return decode(tokens, Language.UNKNOWN); }
public String decode(List<Long> tokens, Language language) { return resolve(language).decode(toArray(tokens)); }
@Override public void close() { models.forEach((__, model) -> model.close()); }
@Override public void deconstruct() { close(); }
private ai.djl.huggingface.tokenizers.HuggingFaceTokenizer resolve(Language language) {
// Disregard language if there is default model
if (models.size() == 1 && models.containsKey(Language.UNKNOWN)) return models.get(Language.UNKNOWN);
if (models.containsKey(language)) return models.get(language);
throw new IllegalArgumentException("No model for language " + language);
}
private static long[] toArray(Collection<? extends Number> c) { return c.stream().mapToLong(Number::longValue).toArray(); }
public static final class Builder {
private final Map<Language, Path> models = new EnumMap<>(Language.class);
private Boolean addSpecialTokens;
public Builder() {}
public Builder(HuggingFaceTokenizerConfig cfg) {
for (var model : cfg.model())
addModel(Language.fromLanguageTag(model.language()), model.path());
addSpecialTokens(cfg.addSpecialTokens());
}
public Builder addModel(Language lang, Path path) { models.put(lang, path); return this; }
public Builder addDefaultModel(Path path) { return addModel(Language.UNKNOWN, path); }
public Builder addSpecialTokens(boolean enabled) { addSpecialTokens = enabled; return this; }
public HuggingFaceTokenizer build() { return new HuggingFaceTokenizer(this); }
}
}
|
