1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
|
import torch
from torch import nn
from text_recognizer.networks.transducer import load_transducer_loss, Transducer
import unittest
class TestTransducer(unittest.TestCase):
def test_viterbi(self):
T = 5
N = 4
B = 2
# fmt: off
emissions1 = torch.tensor((
0, 4, 0, 1,
0, 2, 1, 1,
0, 0, 0, 2,
0, 0, 0, 2,
8, 0, 0, 2,
),
dtype=torch.float,
).view(T, N)
emissions2 = torch.tensor((
0, 2, 1, 7,
0, 2, 9, 1,
0, 0, 0, 2,
0, 0, 5, 2,
1, 0, 0, 2,
),
dtype=torch.float,
).view(T, N)
# fmt: on
# Test without blank:
labels = [[1, 3, 0], [3, 2, 3, 2, 3]]
transducer = Transducer(
tokens=["a", "b", "c", "d"],
graphemes_to_idx={"a": 0, "b": 1, "c": 2, "d": 3},
blank="none",
)
emissions = torch.stack([emissions1, emissions2], dim=0)
predictions = transducer.viterbi(emissions)
self.assertEqual([p.tolist() for p in predictions], labels)
# Test with blank without repeats:
labels = [[1, 0], [2, 2]]
transducer = Transducer(
tokens=["a", "b", "c"],
graphemes_to_idx={"a": 0, "b": 1, "c": 2},
blank="optional",
allow_repeats=False,
)
emissions = torch.stack([emissions1, emissions2], dim=0)
predictions = transducer.viterbi(emissions)
self.assertEqual([p.tolist() for p in predictions], labels)
if __name__ == "__main__":
unittest.main()
|
