1 /*
2 * Copyright (c) 2026, Oracle and/or its affiliates. All rights reserved.
3 * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
4 *
5 * This code is free software; you can redistribute it and/or modify it
6 * under the terms of the GNU General Public License version 2 only, as
7 * published by the Free Software Foundation.
8 *
9 * This code is distributed in the hope that it will be useful, but WITHOUT
10 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
11 * FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License
12 * version 2 for more details (a copy is included in the LICENSE file that
13 * accompanied this code).
14 *
15 * You should have received a copy of the GNU General Public License version
16 * 2 along with this work; if not, write to the Free Software Foundation,
17 * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
18 *
19 * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
20 * or visit www.oracle.com if you need additional information or have any
21 * questions.
22 */
23 package oracle.code.onnx.bert;
24
25 import oracle.code.onnx.Tensor;
26
27 import java.io.IOException;
28 import java.lang.foreign.Arena;
29 import java.net.URISyntaxException;
30 import java.nio.file.Files;
31 import java.nio.file.Path;
32 import java.util.*;
33
34 public class BertTokenizer {
35 record Batch(Tensor<Long> inputIds, Tensor<Long> attentionMask, Tensor<Long> tokenTypeIds) {}
36
37 private static final Long CLS_TOKEN_ID = 101L;
38 private static final int MAX_TOKENS = 512; // max_position_embeddings
39 private static final long SEP_TOKEN_ID = 102L;
40 private static final long UNK_TOKEN_ID = 100L;
41
42 private static final Set<Integer> PUNCTUATION_CATEGORIES = Set.of((int)Character.CONNECTOR_PUNCTUATION,
43 (int)Character.DASH_PUNCTUATION,
44 (int)Character.START_PUNCTUATION,
45 (int)Character.END_PUNCTUATION,
46 (int)Character.INITIAL_QUOTE_PUNCTUATION,
47 (int)Character.FINAL_QUOTE_PUNCTUATION,
48 (int)Character.OTHER_PUNCTUATION);
49
50 Map<String, Long> vocab;
51
52 public BertTokenizer() throws IOException, URISyntaxException {
53 Path vocabLocation = Path.of(Objects.requireNonNull(BertTokenizer.class.getResource("vocab.txt")).toURI());
54 List<String> lines = Files.readAllLines(vocabLocation);
55 vocab = new HashMap<>(lines.size() * 2);
56 for (int i = 0; i < lines.size(); i++) {
57 vocab.put(lines.get(i), (long) i);
58 }
59 }
60
61 Batch tokenize(Arena arena, String[] sentences) {
62 List<long[]> tokenized = new ArrayList<>(sentences.length);
63 int sequenceLength = 0;
64 for (String sentence : sentences) {
65 long[] ids = encode(sentence);
66 tokenized.add(ids);
67 sequenceLength = Math.max(sequenceLength, ids.length);
68 }
69
70 long[] shape = {sentences.length, sequenceLength};
71 long[] inputIds = new long[sentences.length * sequenceLength];
72 long[] attentionMask = new long[inputIds.length];
73 long[] tokenTypeIds = new long[inputIds.length];
74 for (int row = 0; row < tokenized.size(); row++) {
75 long[] ids = tokenized.get(row);
76 int offset = row * sequenceLength;
77 System.arraycopy(ids, 0, inputIds, offset, ids.length);
78 Arrays.fill(attentionMask, offset, offset + ids.length, 1L);
79 }
80
81 return new Batch(
82 Tensor.ofShape(arena, shape, inputIds),
83 Tensor.ofShape(arena, shape, attentionMask),
84 Tensor.ofShape(arena, shape, tokenTypeIds));
85 }
86
87 private long[] encode(String sentence) {
88 List<Long> ids = new ArrayList<>();
89 ids.add(CLS_TOKEN_ID);
90 for (String token : basicTokens(sentence)) {
91 for (long id : wordPiece(token)) {
92 if (ids.size() == MAX_TOKENS - 1) {
93 ids.add(SEP_TOKEN_ID);
94 return ids.stream().mapToLong(Long::longValue).toArray();
95 }
96 ids.add(id);
97 }
98 }
99 ids.add(SEP_TOKEN_ID);
100 return ids.stream().mapToLong(Long::longValue).toArray();
101 }
102
103 private static List<String> basicTokens(String sentence) {
104 String normalized = sentence.toLowerCase(Locale.ENGLISH);
105 List<String> tokens = new ArrayList<>();
106 StringBuilder current = new StringBuilder();
107 for (int i = 0; i < normalized.length(); i++) {
108 char c = normalized.charAt(i);
109 if (Character.isWhitespace(c) || Character.isISOControl(c)) {
110 flush(current, tokens);
111 } else if (PUNCTUATION_CATEGORIES.contains(Character.getType(c))) {
112 flush(current, tokens);
113 tokens.add(String.valueOf(c));
114 } else {
115 current.append(c);
116 }
117 }
118 flush(current, tokens);
119 return tokens;
120 }
121
122 private static void flush(StringBuilder current, List<String> tokens) {
123 if (!current.isEmpty()) {
124 tokens.add(current.toString());
125 current.setLength(0);
126 }
127 }
128
129 private long[] wordPiece(String token) {
130 Long direct = vocab.get(token);
131 if (direct != null) {
132 return new long[] {direct};
133 }
134
135 List<Long> pieces = new ArrayList<>();
136 int start = 0;
137 while (start < token.length()) {
138 int end = token.length();
139 Long piece = null;
140 while (start < end) {
141 String candidate = start == 0 ? token.substring(start, end) : "##" + token.substring(start, end);
142 piece = vocab.get(candidate);
143 if (piece != null) {
144 break;
145 }
146 end--;
147 }
148 if (piece == null) {
149 return new long[] {UNK_TOKEN_ID};
150 }
151 pieces.add(piece);
152 start = end;
153 }
154 return pieces.stream().mapToLong(Long::longValue).toArray();
155 }
156 }