RASR FSA builder by orthography (#86)

Add an FSA builder class which creates an FSA given an orthography. Moreover, the classes have been refactored in order to achieve a better structure and extensibility. Finally, more documentation has been added.

Author: Icemole

i6_models/parts/rasr_fsa.py

@@ -1,14 +1,33 @@
 from __future__ import annotations
 
-__all__ = ["RasrFsaBuilder", "WeightedFsa", "RasrFsaBuilderV2", "WeightedFsaV2"]
+__all__ = ["WeightedFsa", "WeightedFsaV2", "RasrFsaBuilder", "RasrFsaBuilderV2", "RasrFsaBuilderByOrthography"]
 
+from abc import ABC, abstractmethod
 from functools import reduce
-from typing import Iterable, NamedTuple, Tuple, Union
+from typing import TYPE_CHECKING, Any, Iterable, List, NamedTuple, Tuple, Union
 
 import numpy as np
 import torch
 
 
+if TYPE_CHECKING:
+    import librasr
+
+
+FsaTuple = Tuple[int, int, np.ndarray, np.ndarray]
+"""
+FSA as a tuple containing
+* number of states S
+* number of edges E
+* integer edge array of shape [E, 3] where each row is an edge
+    consisting of from-state, to-state and the emission idx
+* float weight array of shape [E,]
+
+This format is how RASR outputs FSAs when retrieving an FSA by orthography/sequence tag:
+https://github.com/rwth-i6/rasr/blob/2bf347fb70f1298950a4adbda39197242f78a619/src/Python/AllophoneStateFsaBuilder.cc#L60
+"""
+
+
 class WeightedFsa(NamedTuple):
     """
     Convenience class that represents an FSA. It supports scaling the weights of the
@@ -87,44 +106,143 @@ def to(self, device: Union[str, torch.device]) -> WeightedFsaV2:
         )
 
 
-class RasrFsaBuilder:
+AppendedFsa = Tuple[List[int], List[int], torch.Tensor, torch.Tensor]
+"""Data structure used for FSA appending (see function below)."""
+
+
+def _append_fsa(original_fsa: AppendedFsa, fsa_to_append: FsaTuple) -> AppendedFsa:
     """
-    Builder class that wraps around the librasr.AllophoneStateFsaBuilder,
-    bringing the FSAs into the correct format for the `i6_native_ops.fbw.fbw_loss`.
+    Appends an FSA :paramref:`fsa_to_append` at the end of another FSA :paramref:`original_fsa`.
+
+    :param original_fsa: Original FSA.
+    :param fsa_to_append: FSA to concatenate to :paramref:`original_fsa`.
+    :return: FSA with the number of states/edges, the edges, and the weights of :paramref:`fsa_to_append`
+        appended at the end of :paramref:`original_fsa`.
+    """
+    edges = torch.from_numpy(np.int32(fsa_to_append[2])).reshape((3, fsa_to_append[1]))
+    return (
+        original_fsa[0] + [fsa_to_append[0]],  # num states
+        original_fsa[1] + [fsa_to_append[1]],  # num edges
+        torch.hstack([original_fsa[2], edges]),  # edges
+        torch.cat([original_fsa[3], torch.from_numpy(fsa_to_append[3])]),  # weights
+    )
+
+
+class _AbstractRasrFsaBuilder(ABC):
+    """
+    Builder class that wraps around the `librasr.AllophoneStateFsaBuilder` class.
+    Creates a single FSA, and joins a batch of FSAs with the correct format for the corresponding `i6_native_ops` loss.
+
     Use of this class requires a working installation of the python package `librasr`.
-    Hence, the package is locally imported in case other classes are accessed from
-    this module.
-    This class provides an explicit implementation of the `__getstate__` and `__setstate__`
-    functions, necessary for pickling as the C++-class `librasr.AllophoneStateFsaBuilder`
-    is not picklable.
+    Hence, the package is locally imported in case other classes are accessed from this module.
 
-    :param config_path: path to the RASR fsa exporter config
-    :param tdp_scale: multiply the weights by this scale
+    This class provides an explicit implementation of the `__getstate__` and `__setstate__` functions.
+    This is necessary for pickling as the C++ class `librasr.AllophoneStateFsaBuilder` is not picklable.
     """
 
     def __init__(self, config_path: str, tdp_scale: float = 1.0):
+        """
+        :param config_path: Path to the RASR FSA exporter config. The FSA builder will be created from here.
+        :param tdp_scale: Transition scale to be applied to the weights of the FSA.
+        """
+        self.config_path = config_path
+        self.builder = self.get_builder(config_path=self.config_path)
+        self.tdp_scale = tdp_scale
+
+    def get_builder(self, config_path: str) -> librasr.AllophoneStateFsaBuilder:
         import librasr
 
-        self.config_path = config_path
         config = librasr.Configuration()
-        config.set_from_file(self.config_path)
-        self.builder = librasr.AllophoneStateFsaBuilder(config)
-        self.tdp_scale = tdp_scale
+        config.set_from_file(config_path)
+        return librasr.AllophoneStateFsaBuilder(config)
 
     def __getstate__(self):
         state = self.__dict__.copy()
         del state["builder"]
         return state
 
     def __setstate__(self, state):
-        import librasr
-
         self.__dict__.update(state)
-        config = librasr.Configuration()
-        config.set_from_file(self.config_path)
-        self.builder = librasr.AllophoneStateFsaBuilder(config)
+        self.builder = self.get_builder(config_path=self.config_path)
+
+    def apply_tdp_scale_to_fsa_tuple(self, fsa: FsaTuple, tdp_scale: float) -> FsaTuple:
+        """
+        Scales the weights of an FSA represented as a tuple by the factor (TDP scale) provided.
+
+        :param fsa: FSA as a tuple containing
+            * number of states S
+            * number of edges E
+            * integer edge array of shape [E, 3] where each row is an edge
+                consisting of from-state, to-state and the emission idx
+            * float weight array of shape [E,]
+        :param tdp_scale: TDP scale by which the weights must be multiplied.
+        :return: FSA with scaled weights corresponding to :paramref:`tdp_scale`.
+        """
+        if tdp_scale == 1.0:
+            # No scaling.
+            return fsa
+        else:
+            return (fsa[0], fsa[1], fsa[2], fsa[3] * tdp_scale)
+
+    @abstractmethod
+    def build_single(self, single_identifier: Any) -> FsaTuple:
+        """
+        Builds a single FSA by calling the respective builder function.
+        The specific implementation depends on the type of FSA builder that is being created.
+
+        Note: it's recommended that the TDP scale is applied here.
+        For that, :funcref:`apply_tdp_scale_to_fsa_tuple` can be called.
+
+        :param single_identifier: Identifier of the sequence for which an FSA must be built.
+        :return: FSA as a tuple corresponding to the sequence identified by :paramref:`single_identifier`.
+            The returned value contains the following fields in order:
+            * number of states S
+            * number of edges E
+            * integer edge array of shape [E, 3] where each row is an edge
+                consisting of from-state, to-state and the emission idx
+            * float weight array of shape [E,]
+        """
+        ...
 
-    def build_single(self, seq_tag: str) -> Tuple[int, int, np.ndarray, np.ndarray]:
+    @abstractmethod
+    def build_batched_fsa(self, fsas: Iterable[FsaTuple]) -> Union[WeightedFsa, WeightedFsaV2]:
+        """
+        Creates the final FSA to be used by the corresponding `fbw` op from `i6_native_ops`.
+
+        :param fsas: Sequence of FSAs to be batched together.
+        :return: Single FSA which bundles together all FSAs provided as parameter.
+            The final object is compatible with the corresponding `fbw` op from `i6_native_ops`.
+        """
+        ...
+
+    def build_batch(self, multiple_identifiers: Iterable[Any]) -> Union[WeightedFsa, WeightedFsaV2]:
+        """
+        Build and concatenate the FSAs for a batch of data.
+
+        :funcref:`build_single` is called once for each item in :paramref:`multiple_identifiers`
+        in order to obtain the individual FSAs.
+
+        :param multiple_identifiers: Multiple elements for which the builder should create the FSAs.
+            In order to build each individual FSA,
+            :funcref:`build_single` should be called once for each item in :paramref:`multiple_identifiers`.
+        :return: Single FSA which joins all other FSAs retrieved.
+        """
+
+        fsas: Iterable[FsaTuple] = map(self.build_single, multiple_identifiers)
+
+        return self.build_batched_fsa(fsas)
+
+
+class RasrFsaBuilder(_AbstractRasrFsaBuilder):
+    """
+    Builder class that wraps around the librasr.AllophoneStateFsaBuilder,
+    bringing the FSAs into the correct format for the `i6_native_ops.fbw.fbw_loss`.
+
+    :param config_path: path to the RASR fsa exporter config
+    :param tdp_scale: multiply the weights by this scale
+    """
+
+    def build_single(self, seq_tag: str) -> FsaTuple:
         """
         Build the FSA for the given sequence tag in the corpus.
 
@@ -139,7 +257,7 @@ def build_single(self, seq_tag: str) -> Tuple[int, int, np.ndarray, np.ndarray]:
         raw_fsa = self.builder.build_by_segment_name(seq_tag)
         return raw_fsa
 
-    def build_batch(self, seq_tags: Iterable[str]) -> WeightedFsa:
+    def build_batched_fsa(self, fsas: Iterable[FsaTuple]) -> WeightedFsa:
         """
         Build and concatenate the FSAs for a batch of sequence tags
         and reformat as an input to `i6_native_ops.fbw.fbw_loss`.
@@ -149,23 +267,17 @@ def build_batch(self, seq_tags: Iterable[str]) -> WeightedFsa:
         the batch.
         Additionally we apply an optional scale to the weights.
 
-        :param seq_tags: an iterable object of sequence tags
+        :param fsas: Sequence of FSAs as a tuple containing:
+            * number of states S
+            * number of edges E
+            * integer edge array of shape [E, 3] where each row is an edge
+                consisting of from-state, to-state and the emission idx
+            * float weight array of shape [E,]
         :return: a concatenated FSA
         """
 
-        def append_fsa(a, b):
-            edges = torch.from_numpy(np.int32(b[2])).reshape((3, b[1]))
-            return (
-                a[0] + [b[0]],  # num states
-                a[1] + [b[1]],  # num edges
-                torch.hstack([a[2], edges]),  # edges
-                torch.cat([a[3], torch.from_numpy(b[3])]),  # weights
-            )
-
-        # concatenate all FSAs in the batch into a single one where state ids are not yet unique
-        fsas = map(self.build_single, seq_tags)
         empty_fsa = ([], [], torch.empty((3, 0), dtype=torch.int32), torch.empty((0,)))
-        num_states, num_edges, all_edges, all_weights = reduce(append_fsa, fsas, empty_fsa)
+        num_states, num_edges, all_edges, all_weights = reduce(_append_fsa, fsas, empty_fsa)
         num_edges = torch.tensor(num_edges, dtype=torch.int32)
         num_states = torch.tensor(num_states, dtype=torch.int32)
 
@@ -194,13 +306,33 @@ def append_fsa(a, b):
         return out_fsa
 
 
-class RasrFsaBuilderV2(RasrFsaBuilder):
+class _RasrFsaBuilderFbw2(_AbstractRasrFsaBuilder):
     """
-    An update of the RasrFsaBuilder that is compatible with the fbw2 op from i6_native_ops.
+    Abstract base class for building an FSA.
+    Internally uses allophones to model the FSA by means of `librasr.AllophoneStateFsaBuilder`.
+
+    The implementation is compatible with the `fbw2` op from `i6_native_ops`.
+
+    The user must overwrite the :funcref:`build_single` method.
+
+    The TDP scale must be explicitly called when running :funcref:`build_single`.
+    For that, :funcref:`apply_tdp_scale_to_fsa_tuple` can be used.
+
+    Using any subclass requires a working installation of the python package `librasr`.
     """
 
-    def build_batch(self, seq_tags: Iterable[str]) -> WeightedFsaV2:
-        fsas = list(map(self.build_single, seq_tags))
+    def build_batched_fsa(self, fsas: Iterable[FsaTuple]) -> WeightedFsaV2:
+        """
+        Joins a set of FSAs represented as tuples into a single :classref:`WeightedFsaV2` object.
+
+        :param fsas: FSAs to be concatenated, represented as tuples with the following fields:
+            * number of states S
+            * number of edges E
+            * integer edge array of shape [E, 3] where each row is an edge
+                consisting of from-state, to-state and the emission idx
+            * float weight array of shape [E,]
+        :return: Single FSA object corresponding to the joined FSAs passed as parameter.
+        """
 
         num_states = [f[0] for f in fsas]
         num_edges = [f[1] for f in fsas]
@@ -222,7 +354,55 @@ def build_batch(self, seq_tags: Iterable[str]) -> WeightedFsaV2:
             torch.IntTensor(np.array([start_states, end_states])),
         )
 
-        if self.tdp_scale != 1.0:
-            out_fsa *= self.tdp_scale
-
         return out_fsa
+
+
+class RasrFsaBuilderV2(_RasrFsaBuilderFbw2):
+    """
+    Builds an FSA given a sequence tag.
+    The orthography will be pulled from the corpus provided in the configuration file.
+    Internally uses allophones to model the FSA by means of `librasr.AllophoneStateFsaBuilder`.
+
+    The implementation is compatible with the `fbw2` op from `i6_native_ops`.
+    """
+
+    def build_single(self, single_identifier: str) -> FsaTuple:
+        """
+        Build the FSA for the given sequence tag in the corpus.
+
+        :param single_identifier: sequence tag
+        :return: FSA as a tuple containing
+            * number of states S
+            * number of edges E
+            * integer edge array of shape [E, 3] where each row is an edge
+                consisting of from-state, to-state and the emission idx
+            * float weight array of shape [E,]
+        """
+        raw_fsa = self.builder.build_by_segment_name(single_identifier)
+
+        return self.apply_tdp_scale_to_fsa_tuple(raw_fsa, self.tdp_scale)
+
+
+class RasrFsaBuilderByOrthography(_RasrFsaBuilderFbw2):
+    """
+    Builds an FSA given an orthography.
+    Internally uses allophones to model the FSA by means of `librasr.AllophoneStateFsaBuilder`.
+
+    The implementation is compatible with the `fbw2` op from `i6_native_ops`.
+    """
+
+    def build_single(self, single_identifier: str) -> FsaTuple:
+        """
+        Build the FSA for the given orthography in the corpus.
+
+        :param single_identifier: Segment orthography.
+        :return: FSA as a tuple containing
+            * number of states S
+            * number of edges E
+            * integer edge array of shape [E, 3] where each row is an edge
+                consisting of from-state, to-state and the emission idx
+            * float weight array of shape [E,]
+        """
+        raw_fsa = self.builder.build_by_orthography(single_identifier)
+
+        return self.apply_tdp_scale_to_fsa_tuple(raw_fsa, self.tdp_scale)