def _intersect_device( a_fsas: k2.Fsa, b_fsas: k2.Fsa, b_to_a_map: torch.Tensor, sorted_match_a: bool, batch_size: int = 500, ): """Wrap k2.intersect_device This is a wrapper of k2.intersect_device and its purpose is to split b_fsas into several batches and process each batch separately to avoid CUDA OOM error. The arguments and return value of this function are the same as k2.intersect_device. NOTE: You can decrease batch_size in case of CUDA out of memory error. """ num_fsas = b_fsas.shape[0] if num_fsas <= batch_size: return k2.intersect_device( a_fsas, b_fsas, b_to_a_map=b_to_a_map, sorted_match_a=sorted_match_a ) num_batches = int(math.ceil(float(num_fsas) / batch_size)) splits = [] for i in range(num_batches): start = i * batch_size end = min(start + batch_size, num_fsas) splits.append((start, end)) ans = [] for start, end in splits: indexes = torch.arange(start, end).to(b_to_a_map) fsas = k2.index_fsa(b_fsas, indexes) b_to_a = k2.index_select(b_to_a_map, indexes) path_lats = k2.intersect_device( a_fsas, fsas, b_to_a_map=b_to_a, sorted_match_a=sorted_match_a ) ans.append(path_lats) return k2.cat(ans)
def intersect(self, lats: Fsa) -> 'Nbest': '''Intersect this Nbest object with a lattice and get 1-best path from the resulting FsaVec. Caution: We assume FSAs in `self.fsa` don't have epsilon self-loops. We also assume `self.fsa.labels` and `lats.labels` are token IDs. Args: lats: An FsaVec. It can be the return value of :func:`whole_lattice_rescoring`. Returns: Return a new Nbest. This new Nbest shares the same shape with `self`, while its `fsa` is the 1-best path from intersecting `self.fsa` and `lats. ''' assert self.fsa.device == lats.device, \ f'{self.fsa.device} vs {lats.device}' assert len(lats.shape) == 3, f'{lats.shape}' assert lats.arcs.dim0() == self.shape.dim0(), \ f'{lats.arcs.dim0()} vs {self.shape.dim0()}' lats = k2.arc_sort(lats) # no-op if lats is already arc sorted fsas_with_epsilon_loops = k2.add_epsilon_self_loops(self.fsa) path_to_seq_map = self.shape.row_ids(1) ans_lats = k2.intersect_device(a_fsas=lats, b_fsas=fsas_with_epsilon_loops, b_to_a_map=path_to_seq_map, sorted_match_a=True) one_best = k2.shortest_path(ans_lats, use_double_scores=True) one_best = k2.remove_epsilon(one_best) return Nbest(fsa=one_best, shape=self.shape)
def test(self): devices = [torch.device('cpu')] if torch.cuda.is_available(): devices.append(torch.device('cuda')) for device in devices: for use_identity_map, sorted_match_a in [(True, True), (False, True), (True, False), (False, False)]: # recognizes (0|1)(0|2) s1 = ''' 0 1 0 0.1 0 1 1 0.2 1 2 0 0.4 1 2 2 0.3 2 3 -1 0.5 3 ''' # recognizes 02* s2 = ''' 0 1 0 1 1 1 2 2 1 2 -1 3 2 ''' # recognizes 1*0 s3 = ''' 0 0 1 10 0 1 0 20 1 2 -1 30 2 ''' a_fsa = k2.Fsa.from_str(s1).to(device) b_fsa_1 = k2.Fsa.from_str(s2).to(device) b_fsa_2 = k2.Fsa.from_str(s3).to(device) a_fsa.requires_grad_(True) b_fsa_1.requires_grad_(True) b_fsa_2.requires_grad_(True) b_fsas = k2.create_fsa_vec([b_fsa_1, b_fsa_2]) if use_identity_map: a_fsas = k2.create_fsa_vec([a_fsa, a_fsa]) b_to_a_map = torch.tensor([0, 1], dtype=torch.int32).to(device) else: a_fsas = k2.create_fsa_vec([a_fsa]) b_to_a_map = torch.tensor([0, 0], dtype=torch.int32).to(device) c_fsas = k2.intersect_device(a_fsas, b_fsas, b_to_a_map, sorted_match_a) assert c_fsas.shape == (2, None, None) c_fsas = k2.connect(c_fsas.to('cpu')) # c_fsas[0] recognizes: 02 # c_fsas[1] recognizes: 10 actual_str_0 = k2.to_str(c_fsas[0]) expected_str_0 = '\n'.join( ['0 1 0 1.1', '1 2 2 2.3', '2 3 -1 3.5', '3']) assert actual_str_0.strip() == expected_str_0 actual_str_1 = k2.to_str(c_fsas[1]) expected_str_1 = '\n'.join( ['0 1 1 10.2', '1 2 0 20.4', '2 3 -1 30.5', '3']) assert actual_str_1.strip() == expected_str_1 loss = c_fsas.scores.sum() (-loss).backward() assert torch.allclose( a_fsa.grad, torch.tensor([-1, -1, -1, -1, -2]).to(a_fsa.grad)) assert torch.allclose( b_fsa_1.grad, torch.tensor([-1, -1, -1]).to(b_fsa_1.grad)) assert torch.allclose( b_fsa_2.grad, torch.tensor([-1, -1, -1]).to(b_fsa_2.grad))
def rescore_with_whole_lattice(lats: k2.Fsa, G_with_epsilon_loops: k2.Fsa) -> k2.Fsa: '''Use whole lattice to rescore. Args: lats: An FsaVec It can be the output of `k2.intersect_dense_pruned`. G_with_epsilon_loops: An FsaVec representing the language model (LM). Note that it is an FsaVec, but it contains only one Fsa. ''' assert len(lats.shape) == 3 assert hasattr(lats, 'lm_scores') assert G_with_epsilon_loops.shape == (1, None, None) device = lats.device lats.scores = lats.scores - lats.lm_scores # Now, lats.scores contains only am_scores # inverted_lats has word IDs as labels. # Its aux_labels are phone IDs, which is a ragged tensor k2.RaggedInt inverted_lats = k2.invert(lats) num_seqs = lats.shape[0] inverted_lats_with_epsilon_loops = k2.add_epsilon_self_loops(inverted_lats) b_to_a_map = torch.zeros(num_seqs, device=device, dtype=torch.int32) try: rescoring_lats = k2.intersect_device(G_with_epsilon_loops, inverted_lats_with_epsilon_loops, b_to_a_map, sorted_match_a=True) except RuntimeError as e: print(f'Caught exception:\n{e}\n') print(f'Number of FSAs: {inverted_lats.shape[0]}') print('num_arcs before pruning: ', inverted_lats_with_epsilon_loops.arcs.num_elements()) # NOTE(fangjun): The choice of the threshold 0.01 is arbitrary here # to avoid OOM. We may need to fine tune it. inverted_lats = k2.prune_on_arc_post(inverted_lats, 0.001, True) inverted_lats_with_epsilon_loops = k2.add_epsilon_self_loops( inverted_lats) print('num_arcs after pruning: ', inverted_lats_with_epsilon_loops.arcs.num_elements()) rescoring_lats = k2.intersect_device(G_with_epsilon_loops, inverted_lats_with_epsilon_loops, b_to_a_map, sorted_match_a=True) rescoring_lats = k2.top_sort(k2.connect( rescoring_lats.to('cpu'))).to(device) inverted_rescoring_lats = k2.invert(rescoring_lats) # inverted rescoring_lats has phone IDs as labels # and word IDs as aux_labels. inverted_rescoring_lats = k2.remove_epsilon_self_loops( inverted_rescoring_lats) best_paths = k2.shortest_path(inverted_rescoring_lats, use_double_scores=True) return best_paths
def rescore_with_whole_lattice(lats: k2.Fsa, G_with_epsilon_loops: k2.Fsa, lm_scale_list: List[float] ) -> Dict[str, k2.Fsa]: '''Use whole lattice to rescore. Args: lats: An FsaVec It can be the output of `k2.intersect_dense_pruned`. G_with_epsilon_loops: An FsaVec representing the language model (LM). Note that it is an FsaVec, but it contains only one Fsa. lm_scale_list: A list containing lm_scale values. Returns: A dict of FsaVec, whose key is a lm_scale and the value represents the best decoding path for each sequence in the lattice. ''' assert len(lats.shape) == 3 assert hasattr(lats, 'lm_scores') assert G_with_epsilon_loops.shape == (1, None, None) device = lats.device lats.scores = lats.scores - lats.lm_scores # We will use lm_scores from G, so remove lats.lm_scores here del lats.lm_scores assert hasattr(lats, 'lm_scores') is False # lats.scores = scores / lm_scale # Now, lats.scores contains only am_scores # inverted_lats has word IDs as labels. # Its aux_labels are phone IDs, which is a ragged tensor k2.RaggedInt inverted_lats = k2.invert(lats) num_seqs = lats.shape[0] b_to_a_map = torch.zeros(num_seqs, device=device, dtype=torch.int32) try: rescoring_lats = k2.intersect_device(G_with_epsilon_loops, inverted_lats, b_to_a_map, sorted_match_a=True) except RuntimeError as e: print(f'Caught exception:\n{e}\n') print(f'Number of FSAs: {inverted_lats.shape[0]}') print('num_arcs before pruning: ', inverted_lats.arcs.num_elements()) # NOTE(fangjun): The choice of the threshold 0.01 is arbitrary here # to avoid OOM. We may need to fine tune it. inverted_lats = k2.prune_on_arc_post(inverted_lats, 0.001, True) print('num_arcs after pruning: ', inverted_lats.arcs.num_elements()) rescoring_lats = k2.intersect_device(G_with_epsilon_loops, inverted_lats, b_to_a_map, sorted_match_a=True) rescoring_lats = k2.top_sort(k2.connect(rescoring_lats.to('cpu')).to(device)) # inv_lats has phone IDs as labels # and word IDs as aux_labels. inv_lats = k2.invert(rescoring_lats) ans = dict() # # The following implements # scores = (scores - lm_scores)/lm_scale + lm_scores # = scores/lm_scale + lm_scores*(1 - 1/lm_scale) # saved_scores = inv_lats.scores.clone() for lm_scale in lm_scale_list: am_scores = saved_scores - inv_lats.lm_scores am_scores /= lm_scale inv_lats.scores = am_scores + inv_lats.lm_scores best_paths = k2.shortest_path(inv_lats, use_double_scores=True) key = f'lm_scale_{lm_scale}' ans[key] = best_paths return ans
def nbest_decoding(lats: k2.Fsa, num_paths: int): ''' (Ideas of this function are from Dan) It implements something like CTC prefix beam search using n-best lists The basic idea is to first extra n-best paths from the given lattice, build a word seqs from these paths, and compute the total scores of these sequences in the log-semiring. The one with the max score is used as the decoding output. ''' # First, extract `num_paths` paths for each sequence. # paths is a k2.RaggedInt with axes [seq][path][arc_pos] paths = k2.random_paths(lats, num_paths=num_paths, use_double_scores=True) # word_seqs is a k2.RaggedInt sharing the same shape as `paths` # but it contains word IDs. Note that it also contains 0s and -1s. # The last entry in each sublist is -1. word_seqs = k2.index(lats.aux_labels, paths) # Note: the above operation supports also the case when # lats.aux_labels is a ragged tensor. In that case, # `remove_axis=True` is used inside the pybind11 binding code, # so the resulting `word_seqs` still has 3 axes, like `paths`. # The 3 axes are [seq][path][word] # Remove epsilons and -1 from word_seqs word_seqs = k2.ragged.remove_values_leq(word_seqs, 0) # Remove repeated sequences to avoid redundant computation later. # # Since k2.ragged.unique_sequences will reorder paths within a seq, # `new2old` is a 1-D torch.Tensor mapping from the output path index # to the input path index. # new2old.numel() == unique_word_seqs.num_elements() unique_word_seqs, _, new2old = k2.ragged.unique_sequences( word_seqs, need_num_repeats=False, need_new2old_indexes=True) # Note: unique_word_seqs still has the same axes as word_seqs seq_to_path_shape = k2.ragged.get_layer(unique_word_seqs.shape(), 0) # path_to_seq_map is a 1-D torch.Tensor. # path_to_seq_map[i] is the seq to which the i-th path # belongs. path_to_seq_map = seq_to_path_shape.row_ids(1) # Remove the seq axis. # Now unique_word_seqs has only two axes [path][word] unique_word_seqs = k2.ragged.remove_axis(unique_word_seqs, 0) # word_fsas is an FsaVec with axes [path][state][arc] word_fsas = k2.linear_fsa(unique_word_seqs) word_fsas_with_epsilon_loops = k2.add_epsilon_self_loops(word_fsas) # lats has phone IDs as labels and word IDs as aux_labels. # inv_lats has word IDs as labels and phone IDs as aux_labels inv_lats = k2.invert(lats) inv_lats = k2.arc_sort(inv_lats) # no-op if inv_lats is already arc-sorted path_lats = k2.intersect_device(inv_lats, word_fsas_with_epsilon_loops, b_to_a_map=path_to_seq_map, sorted_match_a=True) # path_lats has word IDs as labels and phone IDs as aux_labels path_lats = k2.top_sort(k2.connect(path_lats.to('cpu')).to(lats.device)) tot_scores = path_lats.get_tot_scores(True, True) # RaggedFloat currently supports float32 only. # We may bind Ragged<double> as RaggedDouble if needed. ragged_tot_scores = k2.RaggedFloat(seq_to_path_shape, tot_scores.to(torch.float32)) argmax_indexes = k2.ragged.argmax_per_sublist(ragged_tot_scores) # Since we invoked `k2.ragged.unique_sequences`, which reorders # the index from `paths`, we use `new2old` # here to convert argmax_indexes to the indexes into `paths`. # # Use k2.index here since argmax_indexes' dtype is torch.int32 best_path_indexes = k2.index(new2old, argmax_indexes) paths_2axes = k2.ragged.remove_axis(paths, 0) # best_paths is a k2.RaggedInt with 2 axes [path][arc_pos] best_paths = k2.index(paths_2axes, best_path_indexes) # labels is a k2.RaggedInt with 2 axes [path][phone_id] # Note that it contains -1s. labels = k2.index(lats.labels.contiguous(), best_paths) labels = k2.ragged.remove_values_eq(labels, -1) # lats.aux_labels is a k2.RaggedInt tensor with 2 axes, so # aux_labels is also a k2.RaggedInt with 2 axes aux_labels = k2.index(lats.aux_labels, best_paths.values()) best_path_fsas = k2.linear_fsa(labels) best_path_fsas.aux_labels = aux_labels return best_path_fsas
def levenshtein_alignment( refs: Fsa, hyps: Fsa, hyp_to_ref_map: torch.Tensor, sorted_match_ref: bool = False, ) -> Fsa: '''Get the levenshtein alignment of two FsaVecs This function supports both CPU and GPU. But it is very slow on CPU. Args: refs: An FsaVec (must have 3 axes, i.e., `len(refs.shape) == 3`. It is the output Fsa of the :func:`levenshtein_graph`. hyps: An FsaVec (must have 3 axes) on the same device as `refs`. It is the output Fsa of the :func:`levenshtein_graph`. hyp_to_ref_map: A 1-D torch.Tensor with dtype torch.int32 on the same device as `refs`. Map from FSA-id in `hpys` to the corresponding FSA-id in `refs` that we want to get levenshtein alignment with. E.g. might be an identity map, or all-to-zero, or something the user chooses. Requires - `hyp_to_ref_map.shape[0] == hyps.shape[0]` - `0 <= hyp_to_ref_map[i] < refs.shape[0]` sorted_match_ref: If true, the arcs of refs must be sorted by label (checked by calling code via properties), and we'll use a matching approach that requires this. Returns: Returns an FsaVec containing the alignment information and satisfing `ans.Dim0() == hyps.Dim0()`. Two attributes named `ref_labels` and `hyp_labels` will be added to the returned FsaVec. `ref_labels` contains the aligned sequences of refs and `hyp_labels` contains the aligned sequences of hyps. You can get the levenshtein distance by calling `get_tot_scores` on the returned FsaVec. Examples: >>> hyps = k2.levenshtein_graph([[1, 2, 3], [1, 3, 3, 2]]) >>> refs = k2.levenshtein_graph([[1, 2, 4]]) >>> alignment = k2.levenshtein_alignment( refs, hyps, hyp_to_ref_map=torch.tensor([0, 0], dtype=torch.int32), sorted_match_ref=True) >>> alignment.labels tensor([ 1, 2, 0, -1, 1, 0, 0, 0, -1], dtype=torch.int32) >>> alignment.ref_labels tensor([ 1, 2, 4, -1, 1, 2, 4, 0, -1], dtype=torch.int32) >>> alignment.hyp_labels tensor([ 1, 2, 3, -1, 1, 3, 3, 2, -1], dtype=torch.int32) >>> -alignment.get_tot_scores( use_double_scores=False, log_semiring=False)) tensor([1., 3.]) ''' assert hasattr(refs, "aux_labels") assert hasattr(hyps, "aux_labels") hyps.rename_tensor_attribute_("aux_labels", "hyp_labels") lattice = k2.intersect_device( refs, hyps, b_to_a_map=hyp_to_ref_map, sorted_match_a=sorted_match_ref) lattice = k2.remove_epsilon_self_loops(lattice) alignment = k2.shortest_path(lattice, use_double_scores=True).invert_() alignment.rename_tensor_attribute_("labels", "ref_labels") alignment.rename_tensor_attribute_("aux_labels", "labels") alignment.scores -= getattr( alignment, "__ins_del_score_offset_internal_attr_") return alignment
def whole_lattice_rescoring(lats: Fsa, G_with_epsilon_loops: Fsa) -> Fsa: '''Rescore the 1st pass lattice with an LM. In general, the G in HLG used to obtain `lats` is a 3-gram LM. This function replaces the 3-gram LM in `lats` with a 4-gram LM. Args: lats: The decoding lattice from the 1st pass. We assume it is the result of intersecting HLG with the network output. G_with_epsilon_loops: An LM. It is usually a 4-gram LM with epsilon self-loops. It should be arc sorted. Returns: Return a new lattice rescored with a given G. ''' assert len(lats.shape) == 3, f'{lats.shape}' assert hasattr(lats, 'lm_scores') assert G_with_epsilon_loops.shape == (1, None, None), \ f'{G_with_epsilon_loops.shape}' device = lats.device lats.scores = lats.scores - lats.lm_scores # Now lats contains only acoustic scores # We will use lm_scores from the given G, so remove lats.lm_scores here del lats.lm_scores assert hasattr(lats, 'lm_scores') is False # inverted_lats has word IDs as labels. # Its aux_labels are token IDs, which is a ragged tensor k2.RaggedInt # if lats.aux_labels is a ragged tensor inverted_lats = k2.invert(lats) num_seqs = lats.shape[0] b_to_a_map = torch.zeros(num_seqs, device=device, dtype=torch.int32) while True: try: rescoring_lats = k2.intersect_device(G_with_epsilon_loops, inverted_lats, b_to_a_map, sorted_match_a=True) break except RuntimeError as e: logging.info(f'Caught exception:\n{e}\n') # Usually, this is an OOM exception. We reduce # the size of the lattice and redo k2.intersect_device() # NOTE(fangjun): The choice of the threshold 1e-5 is arbitrary here # to avoid OOM. We may need to fine tune it. logging.info(f'num_arcs before: {inverted_lats.num_arcs}') inverted_lats = k2.prune_on_arc_post(inverted_lats, 1e-5, True) logging.info(f'num_arcs after: {inverted_lats.num_arcs}') rescoring_lats = k2.top_sort(k2.connect(rescoring_lats)) # inv_rescoring_lats has token IDs as labels # and word IDs as aux_labels. inv_rescoring_lats = k2.invert(rescoring_lats) return inv_rescoring_lats