def test_fsa_io(self): s = ''' 0 1 10 0.1 0 2 20 0.2 1 3 -1 0.3 2 3 -1 0.4 3 ''' fsa = k2.Fsa.from_str(s) tensor = k2.to_tensor(fsa) assert tensor.ndim == 2 assert tensor.dtype == torch.int32 del fsa # tensor is still accessible fsa = k2.Fsa(tensor) del tensor assert torch.allclose( fsa.scores, torch.tensor([0.1, 0.2, 0.3, 0.4], dtype=torch.float32)) assert torch.allclose( fsa.arcs.values()[:, :-1], # skip the last field `scores` torch.tensor([[0, 1, 10], [0, 2, 20], [1, 3, -1], [2, 3, -1]], dtype=torch.int32)) # now test vector of FSAs ragged_arc = _k2.fsa_to_fsa_vec(fsa.arcs) del fsa fsa_vec = k2.Fsa(ragged_arc) del ragged_arc assert fsa_vec.shape == (1, None, None) assert torch.allclose( fsa_vec.scores, torch.tensor([0.1, 0.2, 0.3, 0.4], dtype=torch.float32)) assert torch.allclose( fsa_vec.arcs.values()[:, :-1], # skip the last field `scores` torch.tensor([[0, 1, 10], [0, 2, 20], [1, 3, -1], [2, 3, -1]], dtype=torch.int32)) tensor = k2.to_tensor(fsa_vec) assert tensor.ndim == 1 assert tensor.dtype == torch.int32 del fsa_vec fsa_vec = k2.Fsa(tensor) assert torch.allclose( fsa_vec.scores, torch.tensor([0.1, 0.2, 0.3, 0.4], dtype=torch.float32)) assert torch.allclose( fsa_vec.arcs.values()[:, :-1], # skip the last field `scores` torch.tensor([[0, 1, 10], [0, 2, 20], [1, 3, -1], [2, 3, -1]], dtype=torch.int32))
def test_symbol_table_and_dot(self): isym_str = ''' <eps> 0 a 1 b 2 c 3 ''' osym_str = ''' <eps> 0 x 1 y 2 z 3 ''' isym = k2.SymbolTable.from_str(isym_str) osym = k2.SymbolTable.from_str(osym_str) rules = ''' 0 1 1 1 0.5 0 1 2 2 1.5 1 2 3 3 2.5 2 3 -1 0 3.5 3 ''' fsa = k2.Fsa(_remove_leading_spaces(rules)) fsa.set_isymbol(isym) fsa.set_osymbol(osym) dot = fsa.to_dot() dot.render('/tmp/fsa', format='pdf')
def test_acceptor_from_tensor(self): fsa_tensor = torch.tensor( [[0, 1, 2, _k2._float_as_int(-1.2)], [0, 2, 10, _k2._float_as_int(-2.2)], [1, 6, -1, _k2._float_as_int(-3.2)], [1, 3, 3, _k2._float_as_int(-4.2)], [2, 6, -1, _k2._float_as_int(-5.2)], [2, 4, 2, _k2._float_as_int(-6.2)], [3, 6, -1, _k2._float_as_int(-7.2)], [5, 0, 1, _k2._float_as_int(-8.2)]], dtype=torch.int32) fsa = k2.Fsa(fsa_tensor) expected_str = ''' 0 1 2 -1.2 0 2 10 -2.2 1 6 -1 -3.2 1 3 3 -4.2 2 6 -1 -5.2 2 4 2 -6.2 3 6 -1 -7.2 5 0 1 -8.2 6 ''' assert _remove_leading_spaces(expected_str) == _remove_leading_spaces( fsa.to_str()) arcs = fsa.arcs assert isinstance(arcs, torch.Tensor) assert arcs.dtype == torch.int32 assert arcs.device.type == 'cpu' assert arcs.shape == (8, 3), 'there should be 8 arcs' assert torch.allclose(arcs[0], torch.tensor([0, 1, 2], dtype=torch.int32)) assert torch.allclose( fsa.weights, torch.tensor([-1.2, -2.2, -3.2, -4.2, -5.2, -6.2, -7.2, -8.2], dtype=torch.float32)) fsa = fsa.to('cuda') arcs[0][0] += 10 assert arcs[0][0] == 10, 'arcs should still be accessible' arcs = fsa.arcs assert arcs.dtype == torch.int32 assert arcs.device.type == 'cuda' assert arcs.device.index == 0 assert arcs.shape == (8, 3), 'there should be 8 arcs' assert torch.allclose( arcs[1], torch.tensor([0, 2, 10], dtype=torch.int32, device=arcs.device))
def test_acceptor_from_tensor(self): fsa_tensor = torch.tensor( [[0, 1, 2, _k2._float_as_int(-1.2)], [0, 2, 10, _k2._float_as_int(-2.2)], [1, 6, -1, _k2._float_as_int(-3.2)], [1, 3, 3, _k2._float_as_int(-4.2)], [2, 6, -1, _k2._float_as_int(-5.2)], [2, 4, 2, _k2._float_as_int(-6.2)], [3, 6, -1, _k2._float_as_int(-7.2)], [5, 0, 1, _k2._float_as_int(-8.2)]], dtype=torch.int32) fsa = k2.Fsa(fsa_tensor) expected_str = ''' 0 1 2 -1.2 0 2 10 -2.2 1 6 -1 -3.2 1 3 3 -4.2 2 6 -1 -5.2 2 4 2 -6.2 3 6 -1 -7.2 5 0 1 -8.2 6 ''' assert _remove_leading_spaces(expected_str) == _remove_leading_spaces( k2.to_str(fsa)) arcs = fsa.arcs.values()[:, :-1] assert isinstance(arcs, torch.Tensor) assert arcs.dtype == torch.int32 assert arcs.device.type == 'cpu' assert arcs.shape == (8, 3), 'there should be 8 arcs' assert torch.allclose(arcs[0], torch.tensor([0, 1, 2], dtype=torch.int32)) assert torch.allclose( fsa.scores, torch.tensor([-1.2, -2.2, -3.2, -4.2, -5.2, -6.2, -7.2, -8.2], dtype=torch.float32)) fsa.scores *= -1 assert torch.allclose( fsa.scores, torch.tensor([1.2, 2.2, 3.2, 4.2, 5.2, 6.2, 7.2, 8.2], dtype=torch.float32))
def test_transducer_from_str(self): s = ''' 0 1 2 22 -1.2 0 2 10 100 -2.2 1 3 3 33 -3.2 1 6 -1 16 -4.2 2 6 -1 26 -5.2 2 4 2 22 -6.2 3 6 -1 36 -7.2 5 0 1 50 -8.2 6 ''' fsa = k2.Fsa(_remove_leading_spaces(s)) assert fsa.aux_labels.dtype == torch.int32 assert fsa.aux_labels.device.type == 'cpu' assert torch.allclose( fsa.aux_labels, torch.tensor([22, 100, 33, 16, 26, 22, 36, 50], dtype=torch.int32)) expected_str = ''' 0 1 2 22 -1.2 0 2 10 100 -2.2 1 3 3 33 -3.2 1 6 -1 16 -4.2 2 6 -1 26 -5.2 2 4 2 22 -6.2 3 6 -1 36 -7.2 5 0 1 50 -8.2 6 ''' assert _remove_leading_spaces(expected_str) == _remove_leading_spaces( fsa.to_str()) expected_str = ''' 0 1 2 22 1.2 0 2 10 100 2.2 1 3 3 33 3.2 1 6 -1 16 4.2 2 6 -1 26 5.2 2 4 2 22 6.2 3 6 -1 36 7.2 5 0 1 50 8.2 6 ''' assert _remove_leading_spaces(expected_str) == _remove_leading_spaces( fsa.to_str(negate_scores=True))
def test_transducer_from_tensor(self): for device in self.devices: fsa_tensor = torch.tensor( [[0, 1, 2, _k2.float_as_int(-1.2)], [0, 2, 10, _k2.float_as_int(-2.2)], [1, 6, -1, _k2.float_as_int(-4.2)], [1, 3, 3, _k2.float_as_int(-3.2)], [2, 6, -1, _k2.float_as_int(-5.2)], [2, 4, 2, _k2.float_as_int(-6.2)], [3, 6, -1, _k2.float_as_int(-7.2)], [5, 0, 1, _k2.float_as_int(-8.2)]], dtype=torch.int32).to(device) aux_labels_tensor = torch.tensor([22, 100, 16, 33, 26, 22, 36, 50], dtype=torch.int32).to(device) fsa = k2.Fsa(fsa_tensor, aux_labels_tensor) assert fsa.aux_labels.dtype == torch.int32 assert fsa.aux_labels.device.type == device.type assert torch.all( torch.eq( fsa.aux_labels, torch.tensor([22, 100, 16, 33, 26, 22, 36, 50], dtype=torch.int32).to(device))) assert torch.allclose( fsa.scores, torch.tensor([-1.2, -2.2, -4.2, -3.2, -5.2, -6.2, -7.2, -8.2], dtype=torch.float32, device=device)) expected_str = ''' 0 1 2 22 -1.2 0 2 10 100 -2.2 1 6 -1 16 -4.2 1 3 3 33 -3.2 2 6 -1 26 -5.2 2 4 2 22 -6.2 3 6 -1 36 -7.2 5 0 1 50 -8.2 6 ''' assert _remove_leading_spaces(expected_str) == \ _remove_leading_spaces(k2.to_str(fsa))
def test_transducer_from_tensor(self): devices = [torch.device('cpu')] if torch.cuda.is_available(): devices.append(torch.device('cuda', 0)) for device in devices: fsa_tensor = torch.tensor( [[0, 1, 2, _k2.float_as_int(-1.2)], [0, 2, 10, _k2.float_as_int(-2.2)], [1, 6, -1, _k2.float_as_int(-4.2)], [1, 3, 3, _k2.float_as_int(-3.2)], [2, 6, -1, _k2.float_as_int(-5.2)], [2, 4, 2, _k2.float_as_int(-6.2)], [3, 6, -1, _k2.float_as_int(-7.2)], [5, 0, 1, _k2.float_as_int(-8.2)]], dtype=torch.int32).to(device) aux_labels_tensor = torch.tensor([22, 100, 16, 33, 26, 22, 36, 50], dtype=torch.int32).to(device) fsa = k2.Fsa(fsa_tensor, aux_labels_tensor) assert fsa.aux_labels.dtype == torch.int32 assert fsa.aux_labels.device.type == device.type assert torch.allclose( fsa.aux_labels, torch.tensor([22, 100, 16, 33, 26, 22, 36, 50], dtype=torch.int32).to(device)) expected_str = ''' 0 1 2 22 -1.2 0 2 10 100 -2.2 1 6 -1 16 -4.2 1 3 3 33 -3.2 2 6 -1 26 -5.2 2 4 2 22 -6.2 3 6 -1 36 -7.2 5 0 1 50 -8.2 6 ''' assert _remove_leading_spaces( expected_str) == _remove_leading_spaces(k2.to_str(fsa))
def test_transducer_from_tensor(self): device_id = 0 device = torch.device('cuda', device_id) fsa_tensor = torch.tensor( [[0, 1, 2, _k2._float_as_int(-1.2)], [0, 2, 10, _k2._float_as_int(-2.2)], [1, 6, -1, _k2._float_as_int(-4.2)], [1, 3, 3, _k2._float_as_int(-3.2)], [2, 6, -1, _k2._float_as_int(-5.2)], [2, 4, 2, _k2._float_as_int(-6.2)], [3, 6, -1, _k2._float_as_int(-7.2)], [5, 0, 1, _k2._float_as_int(-8.2)]], dtype=torch.int32).to(device) aux_labels_tensor = torch.tensor([22, 100, 16, 33, 26, 22, 36, 50], dtype=torch.int32).to(device) fsa = k2.Fsa(fsa_tensor, aux_labels_tensor) assert fsa.aux_labels.dtype == torch.int32 assert fsa.aux_labels.device.type == 'cuda' assert torch.allclose( fsa.aux_labels, torch.tensor([22, 100, 16, 33, 26, 22, 36, 50], dtype=torch.int32).to(device)) expected_str = ''' 0 1 2 22 -1.2 0 2 10 100 -2.2 1 6 -1 16 -4.2 1 3 3 33 -3.2 2 6 -1 26 -5.2 2 4 2 22 -6.2 3 6 -1 36 -7.2 5 0 1 50 -8.2 6 ''' assert _remove_leading_spaces(expected_str) == _remove_leading_spaces( fsa.to('cpu').to_str()) # String IO supported on CPU only
def test_bad_case3(self): # `arc_indexes` and `arcs` in this state are not consistent arc_indexes = torch.IntTensor([0, 2, 2, 2]) arcs = torch.IntTensor([[0, 1, 0], [0, 2, 1], [1, 2, 0]]) fsa = k2.Fsa(arc_indexes, arcs) self.assertFalse(k2.is_valid(fsa))
def test_acceptor_from_str(self): s = ''' 0 1 2 -1.2 0 2 10 -2.2 1 3 3 -3.2 1 6 -1 -4.2 2 6 -1 -5.2 2 4 2 -6.2 3 6 -1 -7.2 5 0 1 -8.2 6 ''' fsa = k2.Fsa(_remove_leading_spaces(s)) expected_str = ''' 0 1 2 -1.2 0 2 10 -2.2 1 3 3 -3.2 1 6 -1 -4.2 2 6 -1 -5.2 2 4 2 -6.2 3 6 -1 -7.2 5 0 1 -8.2 6 ''' assert _remove_leading_spaces(expected_str) == _remove_leading_spaces( fsa.to_str()) expected_str = ''' 0 1 2 1.2 0 2 10 2.2 1 3 3 3.2 1 6 -1 4.2 2 6 -1 5.2 2 4 2 6.2 3 6 -1 7.2 5 0 1 8.2 6 ''' assert _remove_leading_spaces(expected_str) == _remove_leading_spaces( fsa.to_str(negate_scores=True)) arcs = fsa.arcs assert isinstance(arcs, torch.Tensor) assert arcs.dtype == torch.int32 assert arcs.device.type == 'cpu' assert arcs.shape == (8, 3), 'there should be 8 arcs' assert torch.allclose(arcs[0], torch.tensor([0, 1, 2], dtype=torch.int32)) assert torch.allclose( fsa.weights, torch.tensor([-1.2, -2.2, -3.2, -4.2, -5.2, -6.2, -7.2, -8.2], dtype=torch.float32)) fsa = fsa.to('cuda') arcs[0][0] += 10 assert arcs[0][0] == 10, 'arcs should still be accessible' arcs = fsa.arcs assert arcs.dtype == torch.int32 assert arcs.device.type == 'cuda' assert arcs.device.index == 0 assert arcs.shape == (8, 3), 'there should be 8 arcs' assert torch.allclose( arcs[1], torch.tensor([0, 2, 10], dtype=torch.int32, device=arcs.device))