9
9
import Foundation
10
10
import Hub
11
11
12
- class BertTokenizer {
12
+ public class BertTokenizer {
13
13
private let basicTokenizer = BasicTokenizer ( )
14
14
private let wordpieceTokenizer : WordpieceTokenizer
15
15
private let maxLen = 512
@@ -18,12 +18,12 @@ class BertTokenizer {
18
18
private let vocab : [ String : Int ]
19
19
private let ids_to_tokens : [ Int : String ]
20
20
21
- var bosToken : String ?
22
- var bosTokenId : Int ?
23
- var eosToken : String ?
24
- var eosTokenId : Int ?
21
+ public var bosToken : String ?
22
+ public var bosTokenId : Int ?
23
+ public var eosToken : String ?
24
+ public var eosTokenId : Int ?
25
25
26
- init ( vocab: [ String : Int ] ,
26
+ public init ( vocab: [ String : Int ] ,
27
27
merges: [ String ] ? ,
28
28
tokenizeChineseChars: Bool = true ,
29
29
bosToken: String ? = nil ,
@@ -39,7 +39,7 @@ class BertTokenizer {
39
39
self . eosTokenId = eosToken == nil ? nil : vocab [ eosToken!]
40
40
}
41
41
42
- required convenience init ( tokenizerConfig: Config , tokenizerData: Config , addedTokens: [ String : Int ] ) throws {
42
+ public required convenience init ( tokenizerConfig: Config , tokenizerData: Config , addedTokens: [ String : Int ] ) throws {
43
43
guard let vocab = tokenizerData. model? . vocab? . dictionary as? [ String : Int ] else {
44
44
throw TokenizerError . missingVocab
45
45
}
@@ -51,7 +51,7 @@ class BertTokenizer {
51
51
}
52
52
53
53
54
- func tokenize( text: String ) -> [ String ] {
54
+ public func tokenize( text: String ) -> [ String ] {
55
55
let text = tokenizeChineseCharsIfNeed ( text)
56
56
var tokens : [ String ] = [ ]
57
57
for token in basicTokenizer. tokenize ( text: text) {
@@ -128,21 +128,21 @@ class BertTokenizer {
128
128
129
129
130
130
extension BertTokenizer : PreTrainedTokenizerModel {
131
- var unknownToken : String ? { wordpieceTokenizer. unkToken }
132
- var unknownTokenId : Int ? { vocab [ unknownToken!] }
133
-
131
+ public var unknownToken : String ? { wordpieceTokenizer. unkToken }
132
+ public var unknownTokenId : Int ? { vocab [ unknownToken!] }
133
+
134
134
func encode( text: String ) -> [ Int ] { tokenizeToIds ( text: text) }
135
135
136
136
func decode( tokens: [ Int ] ) -> String {
137
137
let tokens = unTokenize ( tokens: tokens)
138
138
return convertWordpieceToBasicTokenList ( tokens)
139
139
}
140
140
141
- func convertTokenToId( _ token: String ) -> Int ? {
141
+ public func convertTokenToId( _ token: String ) -> Int ? {
142
142
return vocab [ token] ?? unknownTokenId
143
143
}
144
144
145
- func convertIdToToken( _ id: Int ) -> String ? {
145
+ public func convertIdToToken( _ id: Int ) -> String ? {
146
146
return ids_to_tokens [ id]
147
147
}
148
148
}
0 commit comments