Repeat label path (#2)

* fixes #1 - only create sequences when p > 0; produces better hp paths.

src/lib.rs

@@ -63,7 +63,7 @@ fn beam_search_<D: Data<Elem = f32>>(
                 if b == beam_prevs[beam as usize].0 {
                     new_probs.push((beam, base_prob * pr[b], 0.0));
                     let mut new_beam = beam_forward[beam as usize][b - 1];
-                    if new_beam == -1 {
+                    if new_beam == -1 && n_prob > 0.0 {
                         new_beam = beam_prevs.len() as i32;
                         beam_prevs.push((b, beam, fidx));
                         beam_forward[beam as usize][b - 1] = new_beam;

tests/test_decode.py

@@ -58,6 +58,40 @@ def test_beam_search_path(self):
         np.testing.assert_array_equal(emit, path)
         self.assertEqual(len(seq), len(path))
 
+    def test_repeat_sequence_path(self):
+        """ simple beam search path test with a repeated sequence """
+        w = 20
+        x = np.zeros((w, len(self.alphabet)), np.float32)
+        x[:, 0] = 0.5  # set stay prob
+
+        expected_path = [6, 13, 18]
+        for idx in expected_path:
+            x[idx, 0] = 0.0
+            x[idx, 1] = 1.0
+
+        seq, path = beam_search(x, self.alphabet, self.beam_size, self.beam_cut_threshold)
+
+        self.assertEqual(seq, 'AAA')
+        self.assertEqual(len(seq), len(path))
+        self.assertEqual(path, expected_path)
+
+    def test_repeat_sequence_path_with_spread(self):
+        """ simple beam search path test with a repeated sequence with probabilities spread"""
+        w = 20
+        x = np.zeros((w, len(self.alphabet)), np.float32)
+        x[:, 0] = 0.5  # set stay prob
+
+        expected_path = [6, 13, 18]
+        for idx in expected_path:
+            x[idx-1:idx + 1, 0] = 0.0
+            x[idx-1:idx + 1, 1] = 1.0
+
+        seq, path = beam_search(x, self.alphabet, self.beam_size, self.beam_cut_threshold)
+
+        self.assertEqual(seq, 'AAA')
+        self.assertEqual(len(seq), len(path))
+        self.assertEqual(path, expected_path)
+
 
 if __name__ == '__main__':
     main()