Forge
search
csharp72c0d4fb
1using System.Collections.Frozen;
2using System.Text;
3
4namespace HybridCodebaseIndex.Core.Embeddings;
5
6/// <summary>
7/// Minimal BERT-style tokenizer for sentence-transformers WordPiece vocab.txt.
8/// Intentionally small: enough for local embeddings without extra dependencies.
9/// </summary>
10internal sealed class WordPieceTokenizer
11{
12 private readonly FrozenDictionary<string, int> _vocab;
13 private readonly bool _doLowerCase;
14 private readonly int _unkId;
15 private readonly int _clsId;
16 private readonly int _sepId;
17 private readonly int _padId;
18
19 public WordPieceTokenizer(IReadOnlyDictionary<string, int> vocab, bool doLowerCase)
20 {
21 _vocab = vocab.ToFrozenDictionary(StringComparer.Ordinal);
22 _doLowerCase = doLowerCase;
23
24 _unkId = GetRequiredId("[UNK]");
25 _clsId = GetRequiredId("[CLS]");
26 _sepId = GetRequiredId("[SEP]");
27 _padId = _vocab.TryGetValue("[PAD]", out var pad) ? pad : 0;
28 }
29
30 public static WordPieceTokenizer FromVocabFile(string vocabPath, bool doLowerCase)
31 {
32 if (string.IsNullOrWhiteSpace(vocabPath))
33 throw new ArgumentException("embedding_vocab_path is required for embedding_provider=onnx with WordPiece models.", nameof(vocabPath));
34 if (!File.Exists(vocabPath))
35 throw new FileNotFoundException("vocab.txt not found.", vocabPath);
36
37 var dict = new Dictionary<string, int>(capacity: 32_000, comparer: StringComparer.Ordinal);
38 var i = 0;
39 foreach (var line in File.ReadLines(vocabPath))
40 {
41 var token = line.TrimEnd('\r', '\n');
42 if (token.Length == 0)
43 continue;
44 dict[token] = i++;
45 }
46
47 return new WordPieceTokenizer(dict, doLowerCase);
48 }
49
50 public (long[] inputIds, long[] attentionMask, long[] tokenTypeIds) Encode(string text, int seqLen)
51 {
52 seqLen = Math.Clamp(seqLen, 8, 512);
53
54 var ids = new long[seqLen];
55 var attn = new long[seqLen];
56 var typeIds = new long[seqLen]; // all zeros for single-sequence
57
58 ids[0] = _clsId;
59 attn[0] = 1;
60
61 var pos = 1;
62 foreach (var token in BasicTokenize(text))
63 {
64 if (pos >= seqLen - 1)
65 break;
66
67 foreach (var subId in WordPieceTokenizeToIds(token))
68 {
69 if (pos >= seqLen - 1)
70 break;
71 ids[pos] = subId;
72 attn[pos] = 1;
73 pos++;
74 }
75 }
76
77 if (pos < seqLen)
78 {
79 ids[pos] = _sepId;
80 attn[pos] = 1;
81 pos++;
82 }
83
84 for (var i = pos; i < seqLen; i++)
85 ids[i] = _padId;
86
87 return (ids, attn, typeIds);
88 }
89
90 private IEnumerable<string> BasicTokenize(string text)
91 {
92 if (string.IsNullOrWhiteSpace(text))
93 yield break;
94
95 var s = _doLowerCase ? text.ToLowerInvariant() : text;
96
97 var sb = new StringBuilder(capacity: 64);
98 foreach (var ch in s)
99 {
100 if (char.IsWhiteSpace(ch))
101 {
102 if (sb.Length > 0)
103 {
104 yield return sb.ToString();
105 sb.Clear();
106 }
107 continue;
108 }
109
110 if (IsPunctuation(ch))
111 {
112 if (sb.Length > 0)
113 {
114 yield return sb.ToString();
115 sb.Clear();
116 }
117 yield return ch.ToString();
118 continue;
119 }
120
121 sb.Append(ch);
122 }
123 if (sb.Length > 0)
124 yield return sb.ToString();
125 }
126
127 private IEnumerable<int> WordPieceTokenizeToIds(string token)
128 {
129 if (string.IsNullOrEmpty(token))
130 yield break;
131
132 // fast-path: exact match
133 if (_vocab.TryGetValue(token, out var exact))
134 {
135 yield return exact;
136 yield break;
137 }
138
139 const int maxChars = 200;
140 if (token.Length > maxChars)
141 {
142 yield return _unkId;
143 yield break;
144 }
145
146 var start = 0;
147 var isBad = false;
148 var subTokens = new List<int>(capacity: 8);
149
150 while (start < token.Length)
151 {
152 var end = token.Length;
153 int? curId = null;
154
155 while (start < end)
156 {
157 var substr = token[start..end];
158 if (start > 0)
159 substr = "##" + substr;
160
161 if (_vocab.TryGetValue(substr, out var id))
162 {
163 curId = id;
164 break;
165 }
166 end--;
167 }
168
169 if (curId is null)
170 {
171 isBad = true;
172 break;
173 }
174
175 subTokens.Add(curId.Value);
176 start = end;
177 }
178
179 if (isBad)
180 {
181 yield return _unkId;
182 yield break;
183 }
184
185 foreach (var id in subTokens)
186 yield return id;
187 }
188
189 private int GetRequiredId(string token)
190 {
191 if (_vocab.TryGetValue(token, out var id))
192 return id;
193 throw new InvalidOperationException($"Tokenizer vocab is missing required token '{token}'.");
194 }
195
196 private static bool IsPunctuation(char ch)
197 => char.IsPunctuation(ch) || ch is '«' or '»' or '…';
198}
199
200
View only · write via MCP/CIDE