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
108
109
110
111
112
113
|
// Copyright Vespa.ai. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
package com.yahoo.language.wordpiece;
import com.yahoo.collections.Tuple2;
import com.yahoo.language.Language;
import com.yahoo.language.process.StemMode;
import com.yahoo.language.process.Token;
import com.yahoo.language.process.Tokenizer;
import java.io.BufferedReader;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStreamReader;
import java.nio.charset.StandardCharsets;
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.NavigableMap;
import java.util.TreeMap;
/**
* A WordPiece embedder "model" - just a vocabulary of strings with a fixed id (index).
*
* Adapted from
* https://github.com/eclipse/deeplearning4j/blob/master/deeplearning4j/deeplearning4j-nlp-parent/deeplearning4j-nlp/src/main/java/org/deeplearning4j/text/tokenization/tokenizer/BertWordPieceTokenizer.java
* licensed under the Apache License, Version 2.0
*
* @author bergum
* @author bratseth
*/
class Model {
private final String subwordPrefix;
private final Path source;
private final Language language;
private final NavigableMap<String, Integer> vocabulary;
private final Map<Integer, String> tokenId2Token;
Model(String subwordPrefix, Language language, Path path) {
this.subwordPrefix = subwordPrefix;
this.source = path;
this.language = language;
this.vocabulary = new TreeMap<>(Collections.reverseOrder());
this.tokenId2Token = new HashMap<>();
try (BufferedReader reader = new BufferedReader(new InputStreamReader(new FileInputStream(path.toFile()),
StandardCharsets.UTF_8))) {
String token;
int i = 0;
while ((token = reader.readLine()) != null) {
this.vocabulary.put(token, i);
this.tokenId2Token.put(i, token);
i++;
}
}
catch (IOException e) {
throw new IllegalArgumentException("Could not read a WordPiece model from " + path, e);
}
}
Language language() { return language; }
List<Integer> embed(String text, Tokenizer tokenizer) {
List<Integer> ids = new ArrayList<>();
text = text.toLowerCase();
for (Token t : tokenizer.tokenize(text, language, StemMode.NONE, true)) {
String originalToken = t.getTokenString();
String candidate = originalToken;
int count = 0;
while (candidate.length() > 0 && !candidate.equals(subwordPrefix)) {
Tuple2<String, Integer> entry = findLongestSubstring(candidate);
if (entry == null) break;
ids.add(entry.second);
candidate = subwordPrefix + candidate.substring(entry.first.length());
if (count++ > originalToken.length()) break;
}
}
return ids;
}
List<String> segment(String text, Tokenizer tokenizer) {
return embed(text, tokenizer).stream().map(tokenId -> tokenId2Token.get(tokenId)).toList();
}
private Tuple2<String, Integer> findLongestSubstring(String candidate) {
NavigableMap<String, Integer> tailMap = this.vocabulary.tailMap(candidate, true);
if (tailMap.isEmpty())
return null;
String longestSubstring = tailMap.firstKey();
Integer id = tailMap.firstEntry().getValue();
int subStringLength = Math.min(candidate.length(), longestSubstring.length());
while (!candidate.startsWith(longestSubstring)) {
subStringLength--;
tailMap = tailMap.tailMap(candidate.substring(0, subStringLength), true);
if (tailMap.isEmpty())
return null;
longestSubstring = tailMap.firstKey();
id = tailMap.firstEntry().getValue();
}
return new Tuple2<>(longestSubstring, id);
}
@Override
public String toString() {
return "WordPiece model for " + language + ": '" + source + "'";
}
}
|
