Reset shoudlEarlyStop as an instance property of TextDecoder and up…

…date tests
argmaxinc · May 25, 2024 · 3aa94e8 · 3aa94e8
1 parent 302fc8c
commit 3aa94e8
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/Sources/WhisperKit/Core/TextDecoder.swift b/Sources/WhisperKit/Core/TextDecoder.swift
@@ -588,7 +588,7 @@ open class TextDecoder: TextDecoding, WhisperMLModel {
         Logging.debug("Running main loop for a maximum of \(loopCount) iterations, starting at index \(prefilledIndex)")
         var hasAlignment = false
         var isFirstTokenLogProbTooLow = false
-        var shouldEarlyStop = false
+        shouldEarlyStop = false
         for tokenIndex in prefilledIndex..<loopCount {
             let loopStart = Date()
 

diff --git a/Tests/WhisperKitTests/UnitTests.swift b/Tests/WhisperKitTests/UnitTests.swift
@@ -327,7 +327,7 @@ final class UnitTests: XCTestCase {
         XCTAssertEqual(decodingTimePerTokenWithWait, decodingTimePerToken, accuracy: decodingTimePerToken * 0.75, "Decoding predictions per token should not be significantly slower with waiting")
 
         // Assert that more tokens are returned in the callback with waiting
-        XCTAssertGreaterThanOrEqual(tokenCountWithWait, tokenCount, "More tokens should be returned in the callback with waiting")
+        XCTAssertGreaterThan(tokenCountWithWait, tokenCount, "More tokens should be returned in the callback with waiting")
     }
 
     // MARK: - Tokenizer Tests