Handle nsnumber score values (#94)

huggingface · Apr 28, 2024 · fc65432 · fc65432
1 parent 0a606f5
commit fc65432
Showing 1 changed file with 15 additions and 2 deletions.
diff --git a/Sources/Tokenizers/UnigramTokenizer.swift b/Sources/Tokenizers/UnigramTokenizer.swift
@@ -6,6 +6,7 @@
 //  Copyright © 2024 Hugging Face. All rights reserved.
 //
 
+import Foundation
 import Hub
 
 class UnigramTokenizer: PreTrainedTokenizerModel {
@@ -37,8 +38,20 @@ class UnigramTokenizer: PreTrainedTokenizerModel {
         }
 
         vocab = try configVocab.map { piece in
-            guard let token = piece.first as? String else { throw TokenizerError.malformedVocab }
-            guard let score = piece.last as? Float else { throw TokenizerError.malformedVocab }
+            guard let token = piece.first as? String,
+                  let scoreValue = piece.last else {
+                throw TokenizerError.malformedVocab
+            }
+
+            let score: Float
+            if let floatScore = scoreValue as? Float {
+                score = floatScore
+            } else if let numberScore = scoreValue as? NSNumber {
+                score = numberScore.floatValue
+            } else {
+                throw TokenizerError.malformedVocab
+            }
+
             return SentencePieceToken(token: token, score: score)
         }