add input/output symbol table. (k2-fsa#206)

* add input/output symbol table.

* add validation for SymbolTable.

Author: csukuangfj

k2/csrc/fsa_utils.cu

@@ -159,9 +159,8 @@ static Fsa AcceptorFromStream(std::string first_line, std::istringstream &is,
     K2_CHECK_EQ(original_final_states.size(), original_final_weights.size());
     int32_t super_final_state = max_state + 1;
     for (std::size_t i = 0; i != original_final_states.size(); ++i) {
-      arcs.emplace_back(original_final_states[i],
-                        super_final_state,
-                        -1,     // kFinalSymbol
+      arcs.emplace_back(original_final_states[i], super_final_state,
+                        -1,  // kFinalSymbol
                         scale * original_final_weights[i]);
     }
   }
@@ -178,8 +177,7 @@ static Fsa AcceptorFromStream(std::string first_line, std::istringstream &is,
 }
 
 static Fsa TransducerFromStream(std::string first_line, std::istringstream &is,
-                                bool openfst,
-                                Array1<int32_t> *aux_labels) {
+                                bool openfst, Array1<int32_t> *aux_labels) {
   K2_CHECK(aux_labels != nullptr);
 
   std::vector<int32_t> state_aux_labels;
@@ -240,16 +238,15 @@ static Fsa TransducerFromStream(std::string first_line, std::istringstream &is,
     K2_CHECK_EQ(openfst, true);
     K2_CHECK_EQ(original_final_states.size(), original_final_weights.size());
     int32_t super_final_state = max_state + 1;
-    for (std::size_t = 0; i != original_final_states.size(); ++i) {
-      arcs.emplace_back(original_final_states[i],
-                        super_final_state,
-                        -1,             // kFinalSymbol
+    for (std::size_t i = 0; i != original_final_states.size(); ++i) {
+      arcs.emplace_back(original_final_states[i], super_final_state,
+                        -1,  // kFinalSymbol
                         scale * original_final_weights[i]);
       // TODO(guoguo) We are not sure yet what to put as the auxiliary label for
       //              arcs entering the super final state. The only real choices
       //              are kEpsilon or kFinalSymbol. We are using kEpsilon for
       //              now.
-      state_aux_labels.push_back(0);    // kEpsilon
+      state_aux_labels.push_back(0);  // kEpsilon
     }
   }
 

k2/python/k2/__init__.py

@@ -1,10 +1,12 @@
 from .array import Array
 from .fsa import Fsa
+from .symbol_table import SymbolTable
 from _k2 import Arc
 
 # please keep the list sorted
 __all__ = [
     'Arc',
     'Array',
     'Fsa',
+    'SymbolTable',
 ]

k2/python/k2/fsa.py

@@ -7,6 +7,7 @@
 
 import torch
 
+from .symbol_table import SymbolTable
 from _k2 import _Fsa
 from _k2 import _as_float
 from _k2 import _fsa_from_str
@@ -17,7 +18,7 @@
 
 class Fsa(object):
 
-    def __init__(self, s: str, negate_scores: bool = False):
+    def __init__(self, s: str, openfst: bool = False):
         '''Create an Fsa from a string.
 
         The given string `s` consists of lines with the following format:
@@ -49,14 +50,14 @@ def __init__(self, s: str, negate_scores: bool = False):
         Args:
           s:
             The input string. Refer to the above comment for its format.
-          negate_scores:
+          openfst:
             Optional. If true, the string form has the weights as costs,
             not scores, so we negate as we read.
         '''
         fsa: _Fsa
         aux_labels: Optional[torch.Tensor]
 
-        fsa, aux_labels = _fsa_from_str(s, negate_scores)
+        fsa, aux_labels = _fsa_from_str(s, openfst)
 
         self._fsa = fsa
         self._aux_labels = aux_labels
@@ -106,20 +107,43 @@ def from_tensor(cls,
         ans._aux_labels = aux_labels
         return ans
 
-    def to_str(self, negate_scores: bool = False) -> str:
+    def set_isymbol(self, isym: SymbolTable) -> None:
+        '''Set the input symbol table.
+
+        Args:
+          isym:
+            The input symbol table.
+        Returns:
+          None.
+        '''
+        self.isym = isym
+
+    def set_osymbol(self, osym: SymbolTable) -> None:
+        '''Set the output symbol table.
+
+        Args:
+          osym:
+            The output symbol table.
+
+        Returns:
+          None.
+        '''
+        self.osym = osym
+
+    def to_str(self, openfst: bool = False) -> str:
         '''Convert an Fsa to a string.
 
         Note:
           The returned string can be used to construct an Fsa.
 
         Args:
-          negate_scores:
+          openfst:
             Optional. If true, we negate the score during the conversion,
 
         Returns:
           A string representation of the Fsa.
         '''
-        return _fsa_to_str(self._fsa, negate_scores, self._aux_labels)
+        return _fsa_to_str(self._fsa, openfst, self._aux_labels)
 
     @property
     def arcs(self) -> torch.Tensor:
@@ -221,6 +245,7 @@ def to_dot(self) -> Digraph:
             src_state, dst_state, label = arc.tolist()
             src_state = str(src_state)
             dst_state = str(dst_state)
+            label = int(label)
             if label == -1:
                 final_state = dst_state
             if src_state not in seen:
@@ -235,10 +260,16 @@ def to_dot(self) -> Digraph:
                     dot.node(dst_state, label=dst_state, **default_node_attr)
                 seen.add(dst_state)
             if self._aux_labels is not None:
-                aux_label = f':{self._aux_labels[i]}'
+                aux_label = int(self._aux_labels[i])
+                if hasattr(self, 'osym'):
+                    aux_label = self.osym.get(aux_label)
+                aux_label = f':{aux_label}'
             else:
                 aux_label = ''
 
+            if hasattr(self, 'isym') and label != -1:
+                label = self.isym.get(label)
+
             dot.edge(src_state,
                      dst_state,
                      label=f'{label}{aux_label}/{weight:.2f}')

k2/python/k2/symbol_table.py

@@ -0,0 +1,110 @@
+# Copyright (c)  2020  Mobvoi Inc.        (authors: Fangjun Kuang)
+#
+# See ../../../LICENSE for clarification regarding multiple authors
+
+from typing import Dict
+from typing import Union
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class SymbolTable(object):
+    _id2sym: Dict[int, str]
+    '''Map an integer to a symbol.
+    '''
+
+    _sym2id: Dict[str, int]
+    '''Map a symbol to an integer.
+    '''
+
+    def __post_init__(self):
+        for idx, sym in self._id2sym.items():
+            assert self._sym2id[sym] == idx
+            assert idx >= 0
+
+        for sym, idx in self._sym2id.items():
+            assert idx >= 0
+            assert self._id2sym[idx] == sym
+
+        eps_sym = '<eps>'
+        if 0 not in self._id2sym:
+            self._id2sym[0] = eps_sym
+            self._sym2id[eps_sym] = 0
+        else:
+            assert self._id2sym[0] == eps_sym
+            assert self._sym2id[eps_sym] == 0
+
+    @staticmethod
+    def from_str(s: str) -> 'SymbolTable':
+        '''Build a symbol table from a string.
+
+        The string consists of lines. Every line has two fields separated
+        by space(s), tab(s) or both. The first field is the symbol and the
+        second the integer id of the symbol.
+
+        Args:
+          s:
+            The input string with the format described above.
+        Returns:
+          An instance of :class:`SymbolTable`.
+        '''
+        id2sym: Dict[int, str] = dict()
+        sym2id: Dict[str, int] = dict()
+
+        for line in s.split('\n'):
+            fields = line.split()
+            if len(fields) == 0:
+                continue  # skip empty lines
+            assert len(fields) == 2, \
+                    f'Expect a line with 2 fields. Given: {len(fields)}'
+            sym, idx = fields[0], int(fields[1])
+            assert sym not in sym2id, f'Duplicated symbol {sym}'
+            assert idx not in id2sym, f'Duplicated id {idx}'
+            id2sym[idx] = sym
+            sym2id[sym] = idx
+
+        return SymbolTable(_id2sym=id2sym, _sym2id=sym2id)
+
+    @staticmethod
+    def from_file(filename: str) -> 'SymbolTable':
+        '''Build a symbol table from file.
+
+        Every line in the symbol table file has two fields separated by
+        space(s), tab(s) or both. The following is an example file:
+
+        .. code-block::
+
+            <eps> 0
+            a 1
+            b 2
+            c 3
+
+        Args:
+          filename:
+            Name of the symbol table file. Its format is documented above.
+
+        Returns:
+          An instance of :class:`SymbolTable`.
+
+        '''
+        with open(filename, 'r') as f:
+            return SymbolTable.from_str(f.read().strip())
+
+    def get(self, k: Union[int, str]) -> Union[str, int]:
+        '''Get a symbol for an id or get an id for a symbol
+
+        Args:
+          k:
+            If it is an id, it tries to find the symbol corresponding
+            to the id; if it is a symbol, it tries to find the id
+            corresponding to the symbol.
+
+        Returns:
+          An id or a symbol depending on the given ``k``.
+        '''
+        if isinstance(k, int):
+            return self._id2sym[k]
+        elif isinstance(k, str):
+            return self._sym2id(k)
+        else:
+            raise ValueError(f'Unsupported type {type(k)}.')

k2/python/tests/fsa_test.py

@@ -25,8 +25,8 @@ 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
+            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
@@ -39,8 +39,8 @@ def test_acceptor_from_str(self):
         expected_str = '''
             0 1 2 -1.2
             0 2 10 -2.2
-            1 3 3 -3.2
-            1 6 -1 -4.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
@@ -53,16 +53,16 @@ def test_acceptor_from_str(self):
         expected_str = '''
             0 1 2 1.2
             0 2 10 2.2
-            1 3 3 3.2
-            1 6 -1 4.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(negate_scores=True))
+            fsa.to_str(openfst=True))
 
         arcs = fsa.arcs
         assert isinstance(arcs, torch.Tensor)
@@ -94,8 +94,8 @@ 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
+            1 3  3  33  -3.2
             2 6 -1  26  -5.2
             2 4  2  22  -6.2
             3 6 -1  36  -7.2
@@ -107,13 +107,13 @@ def test_transducer_from_str(self):
         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))
+            torch.tensor([22, 100, 16, 33, 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
+            1 3 3 33 -3.2
             2 6 -1 26 -5.2
             2 4 2 22 -6.2
             3 6 -1 36 -7.2
@@ -126,16 +126,47 @@ def test_transducer_from_str(self):
         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
+            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_str(negate_scores=True))
+            fsa.to_str(openfst=True))
+
+    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')
+        # the fsa is saved to /tmp/fsa.pdf
 
 
 if __name__ == '__main__':