Fix condition in DTW backtrace (#1136)

OpenNMT · Mar 17, 2023 · a96ef19 · a96ef19
1 parent 401a82c
commit a96ef19
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 6 deletions.
diff --git a/include/ctranslate2/models/whisper.h b/include/ctranslate2/models/whisper.h
@@ -72,7 +72,7 @@ namespace ctranslate2 {
     };
 
     struct WhisperAlignmentResult {
-      std::vector<std::pair<size_t, size_t>> alignments;
+      std::vector<std::pair<dim_t, dim_t>> alignments;
       std::vector<float> text_token_probs;
     };
 

diff --git a/src/dtw.cc b/src/dtw.cc
@@ -5,7 +5,7 @@
 
 namespace ctranslate2 {
 
-  static std::vector<std::pair<size_t, size_t>> backtrace(StorageView trace) {
+  static std::vector<std::pair<dim_t, dim_t>> backtrace(StorageView trace) {
     dim_t i = trace.dim(0) - 1;
     dim_t j = trace.dim(1) - 1;
 
@@ -14,9 +14,9 @@ namespace ctranslate2 {
     for (dim_t k = 0; k < trace.dim(0); ++k)
       trace.at<int32_t>({k, 0}) = 1;
 
-    std::vector<std::pair<size_t, size_t>> result;
+    std::vector<std::pair<dim_t, dim_t>> result;
 
-    while (i > 0 && j > 0) {
+    while (i > 0 || j > 0) {
       result.emplace_back(i - 1, j - 1);
 
       const auto t = trace.at<int32_t>({i, j});
@@ -38,7 +38,7 @@ namespace ctranslate2 {
     return result;
   }
 
-  std::vector<std::pair<size_t, size_t>> negative_dtw(const StorageView& x) {
+  std::vector<std::pair<dim_t, dim_t>> negative_dtw(const StorageView& x) {
     const dim_t n = x.dim(0);
     const dim_t m = x.dim(1);
 

diff --git a/src/dtw.h b/src/dtw.h
@@ -5,6 +5,6 @@
 namespace ctranslate2 {
 
   // Dynamic time wrapping function, but x values are negated.
-  std::vector<std::pair<size_t, size_t>> negative_dtw(const StorageView& x);
+  std::vector<std::pair<dim_t, dim_t>> negative_dtw(const StorageView& x);
 
 }