Add skipSpecialTokens option to Tokenizer.decode (#148)

huggingface · Dec 26, 2024 · 44e2c04 · 44e2c04
1 parent fc95ce1
commit 44e2c04
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 3 deletions.
diff --git a/Sources/Tokenizers/Tokenizer.swift b/Sources/Tokenizers/Tokenizer.swift
@@ -112,6 +112,7 @@ public protocol Tokenizer {
 
     /// Decode
     func decode(tokens: [Int]) -> String
+    func decode(tokens: [Int], skipSpecialTokens: Bool) -> String
 
     func convertTokenToId(_ token: String) -> Int?
     func convertTokensToIds(_ tokens: [String]) -> [Int?]
@@ -150,6 +151,10 @@ public extension Tokenizer {
     func callAsFunction(_ text: String, addSpecialTokens: Bool = true) -> [Int] {
         encode(text: text, addSpecialTokens: addSpecialTokens)
     }
+
+    func decode(tokens: [Int]) -> String {
+        decode(tokens: tokens, skipSpecialTokens: false)
+    }
 
     func convertTokensToIds(_ tokens: [String]) -> [Int?] {
         return tokens.map { convertTokenToId($0) }
@@ -315,10 +320,17 @@ public class PreTrainedTokenizer: Tokenizer {
         return encode(text: text, addSpecialTokens: true)
     }
 
-    /// Decode
-    public func decode(tokens: [Int]) -> String {
+    public func decode(tokens: [Int], skipSpecialTokens: Bool = false) -> String {
         // IDs to tokens
-        let tokenStrings = tokens.compactMap { model.convertIdToToken($0) }
+        let tokenStrings: [String]
+        if skipSpecialTokens {
+            let specialTokenIDs = Set(specialTokens.values)
+            tokenStrings = tokens
+                .filter { !specialTokenIDs.contains($0) }
+                .compactMap { model.convertIdToToken($0) }
+        } else {
+            tokenStrings = tokens.compactMap { model.convertIdToToken($0) }
+        }
         let decoded = decodeTokens(tokenStrings)
         // At this point we should have a single String
         return cleanUp(text: decoded.joined(separator: ""))

diff --git a/Tests/TokenizersTests/TokenizerTests.swift b/Tests/TokenizersTests/TokenizerTests.swift
@@ -218,6 +218,10 @@ class TokenizerTester {
                 tokenizer.decode(tokens: edgeCase.encoded.input_ids),
                 edgeCase.decoded_with_special
             )
+            XCTAssertEqual(
+                tokenizer.decode(tokens: edgeCase.encoded.input_ids, skipSpecialTokens: true),
+                edgeCase.decoded_without_special
+            )
         }
     }