Upload model

config.json

@@ -12,7 +12,7 @@
   "discriminator_dilation": 1,
   "discriminator_dim": 256,
   "discriminator_dropout": 0.0,
-  "discriminator_input_dim": 512,
+  "discriminator_input_dim": 44,
   "discriminator_kernel": 3,
   "discriminator_linear_emb": false,
   "discriminator_max_pool": false,
@@ -25,14 +25,57 @@
   "generator_dropout": 0.0,
   "generator_input_dim": 512,
   "generator_kernel": 4,
-  "generator_output_dim": 40,
+  "generator_output_dim": 44,
   "generator_stride": 1,
   "model_type": "reborn_uasr",
+  "phones": [
+    "AH",
+    "N",
+    "S",
+    "IH",
+    "T",
+    "L",
+    "R",
+    "D",
+    "K",
+    "IY",
+    "Z",
+    "M",
+    "ER",
+    "EH",
+    "P",
+    "AE",
+    "B",
+    "AA",
+    "EY",
+    "F",
+    "OW",
+    "NG",
+    "G",
+    "V",
+    "AO",
+    "AY",
+    "SH",
+    "UW",
+    "W",
+    "HH",
+    "JH",
+    "Y",
+    "CH",
+    "TH",
+    "AW",
+    "UH",
+    "OY",
+    "DH",
+    "ZH",
+    "<SIL>"
+  ],
   "segmenter_dropout": 0.1,
   "segmenter_hidden_dim": 512,
   "segmenter_input_dim": 512,
   "segmenter_kernel_size": 7,
   "segmenter_type": "cnn",
+  "special_token_nums": 4,
   "torch_dtype": "float32",
   "transformers_version": "4.24.0"
 }

configuration_reborn.py

@@ -1,3 +1,4 @@
+import os
 from transformers import PretrainedConfig
 
 class RebornUASRConfig(PretrainedConfig):
@@ -37,6 +38,10 @@ class RebornUASRConfig(PretrainedConfig):
         generator_dropout: float = 0.0,
         generator_bn_apply: bool = False,
         generator_bn_init_weight: float = 30.0,
+
+        phones: list = [],
+        dict_fpath: str = "",
+        special_token_nums: int = 4, # [<s>, <pad>, </s>, <unk>]
         **kwargs
     ):
         super().__init__(**kwargs)
@@ -70,3 +75,27 @@ class RebornUASRConfig(PretrainedConfig):
         self.generator_bn_apply = generator_bn_apply
         self.generator_bn_init_weight = generator_bn_init_weight
 
+        self.special_token_nums = special_token_nums
+        if os.path.isfile(dict_fpath):
+            self.phones = self.read_phns_dict_from_fpath(dict_fpath)
+        else:
+            self.phones = phones
+        if len(self.phones) > 0:
+            self.generator_output_dim = len(self.phones) + self.special_token_nums
+            self.discriminator_input_dim = self.generator_output_dim
+    
+    def read_phns_dict_from_fpath(self, fpath: str):
+        phns = []
+        with open(fpath, "r") as f:
+            for l in f:
+                phn = l.strip().split('\t')[0].split(' ')[0]
+                phns.append(phn)
+        return phns
+
+def main():
+    config = RebornUASRConfig(dict_fpath="/home/andybi7676/Desktop/uasr-rl/data/ls_100h_new/text/prep/phones/dict.phn.txt")
+    print(config)
+    config.save_pretrained("reborn_uasr")
+
+if __name__ == "__main__":
+    main()

modeling_reborn.py

@@ -2,7 +2,7 @@ import torch
 import torch.nn as nn
 from transformers import PreTrainedModel
 from .configuration_reborn import RebornUASRConfig
-from typing import Optional, Tuple, Union
+from typing import Optional, Tuple, Union, List
 
 class RebornSegmenter(nn.Module):
     def __init__(self, config):
@@ -158,6 +158,176 @@ class RebornGenerator(nn.Module):
 
         return result
 
+def get_item(tensor):
+    # tpu-comment: making this a no-op for xla devices.
+    if torch.is_tensor(tensor) and tensor.device.type == "xla":
+        return tensor.detach()
+    if hasattr(tensor, "item"):
+        return tensor.item()
+    if hasattr(tensor, "__getitem__"):
+        return tensor[0]
+    return tensor
+
+def post_process(sentence: str, symbol: str):
+    if symbol == "sentencepiece":
+        sentence = sentence.replace(" ", "").replace("\u2581", " ").strip()
+    elif symbol == "wordpiece":
+        sentence = sentence.replace(" ", "").replace("_", " ").strip()
+    elif symbol == "letter":
+        sentence = sentence.replace(" ", "").replace("|", " ").strip()
+    elif symbol == "silence":
+        import re
+        sentence = sentence.replace("<SIL>", "")
+        sentence = re.sub(' +', ' ', sentence).strip()
+    elif symbol == "_EOW":
+        sentence = sentence.replace(" ", "").replace("_EOW", " ").strip()
+    elif symbol in {"subword_nmt", "@@ ", "@@"}:
+        if symbol == "subword_nmt":
+            symbol = "@@ "
+        sentence = (sentence + " ").replace(symbol, "").rstrip()
+    elif symbol == "none":
+        pass
+    elif symbol is not None:
+        raise NotImplementedError(f"Unknown post_process option: {symbol}")
+    return sentence
+
+class SimpleTokenizer(object):
+    def __init__(self,
+        phones: List[str],
+        bos="<s>",
+        pad="<pad>",
+        eos="</s>",
+        unk="<unk>",
+        extra_special_symbols=None,
+    ) -> None:
+        self.bos_word, self.unk_word, self.pad_word, self.eos_word = bos, unk, pad, eos
+        self.symbols = []
+        self.count = []
+        self.indices = {}
+        self.bos_index = self.add_symbol(bos)
+        self.pad_index = self.add_symbol(pad)
+        self.eos_index = self.add_symbol(eos)
+        self.unk_index = self.add_symbol(unk)
+        if extra_special_symbols:
+            for s in extra_special_symbols:
+                self.add_symbol(s)
+        self.nspecial = len(self.symbols)
+        for phone in phones:
+            self.add_symbol(phone)
+        self.postprocess_code = "silence"
+    
+    def add_symbol(self, word, n=1, overwrite=False):
+        """Adds a word to the dictionary"""
+        if word in self.indices and not overwrite:
+            idx = self.indices[word]
+            self.count[idx] = self.count[idx] + n
+            return idx
+        else:
+            idx = len(self.symbols)
+            self.indices[word] = idx
+            self.symbols.append(word)
+            self.count.append(n)
+            return idx
+
+    def __eq__(self, other):
+        return self.indices == other.indices
+
+    def __getitem__(self, idx):
+        if idx < len(self.symbols):
+            return self.symbols[idx]
+        return self.unk_word
+
+    def get_count(self, idx):
+        return self.count[idx]
+
+    def __len__(self):
+        """Returns the number of symbols in the dictionary"""
+        return len(self.symbols)
+
+    def __contains__(self, sym):
+        return sym in self.indices
+
+    def index(self, sym):
+        """Returns the index of the specified symbol"""
+        assert isinstance(sym, str)
+        if sym in self.indices:
+            return self.indices[sym]
+        return self.unk_index
+
+    def string(
+        self,
+        tensor,
+        bpe_symbol=None,
+        escape_unk=False,
+        extra_symbols_to_ignore=None,
+        unk_string=None,
+        include_eos=False,
+        separator=" ",
+    ):
+        """Helper for converting a tensor of token indices to a string.
+
+        Can optionally remove BPE symbols or escape <unk> words.
+        """
+        if torch.is_tensor(tensor) and tensor.dim() == 2:
+            return "\n".join(
+                self.string(
+                    t,
+                    bpe_symbol,
+                    escape_unk,
+                    extra_symbols_to_ignore,
+                    include_eos=include_eos,
+                )
+                for t in tensor
+            )
+
+        extra_symbols_to_ignore = set(extra_symbols_to_ignore or [])
+        if not include_eos:
+            extra_symbols_to_ignore.add(self.eos())
+
+        def token_string(i):
+            if i == self.unk():
+                if unk_string is not None:
+                    return unk_string
+                else:
+                    return self.unk_string(escape_unk)
+            else:
+                return self[i]
+
+        if hasattr(self, "bos_index"):
+            extra_symbols_to_ignore.add(self.bos())
+
+        sent = separator.join(
+            token_string(i)
+            for i in tensor
+            if get_item(i) not in extra_symbols_to_ignore
+        )
+
+        return post_process(sent, bpe_symbol)
+
+    def unk_string(self, escape=False):
+        """Return unknown string, optionally escaped as: <<unk>>"""
+        if escape:
+            return "<{}>".format(self.unk_word)
+        else:
+            return self.unk_word
+    
+    def bos(self):
+        """Helper to get index of beginning-of-sentence symbol"""
+        return self.bos_index
+
+    def pad(self):
+        """Helper to get index of pad symbol"""
+        return self.pad_index
+
+    def eos(self):
+        """Helper to get index of end-of-sentence symbol"""
+        return self.eos_index
+
+    def unk(self):
+        """Helper to get index of unk symbol"""
+        return self.unk_index
+
+
 class RebornUASRModel(PreTrainedModel):
     config_class = RebornUASRConfig
 
@@ -166,6 +336,9 @@ class RebornUASRModel(PreTrainedModel):
         self.pca = nn.Linear(1024, 512)
         self.segmenter = RebornSegmenter(config)
         self.generator = RebornGenerator(config)
+        self.tokenizer = None
+        if len(config.phones) > 0:
+            self.tokenizer = SimpleTokenizer(config.phones)
 
     def forward(
         self,
@@ -181,4 +354,28 @@ class RebornUASRModel(PreTrainedModel):
             'x_segmented': x_segmented,
             'x_generated': x_generated
         }
+    
+    def generate(self, x, padding_mask, merge_consecutive=True, remove_silence=True):
+        res = self.forward(x, padding_mask)
+        y_raw_logits = res['x_generated']['dense_x']
+        y_raw_padding = res['x_generated']['dense_padding_mask']
+        y_raw_logits[y_raw_padding][..., self.tokenizer.pad_index] = float('inf')
+        preds = y_raw_logits.argmax(-1)
+        hyps = []
+        postprocess_code = "silence" if remove_silence else "none"
+        for pred in preds:
+            if merge_consecutive:
+                # merge consecutive predictions
+                pred = torch.unique_consecutive(pred)
+            hyp = self.tokenizer.string(pred, bpe_symbol=postprocess_code)
+            hyps.append(hyp)
+        return hyps
+
+def main():
+    model_config = RebornUASRConfig.from_pretrained("/home/andybi7676/Desktop/uasr-rl/reborn_uasr/config.json")
+    print(model_config)
+    model = RebornUASRModel(model_config)
+    print(model.tokenizer.indices)
 
+if __name__ == "__main__":
+    main()

pytorch_model.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:319b78a86e5743fd8239760ea2628d8d64cd4bfe293423e2acaa204ee6954f4f
-size 12923917
+oid sha256:8f6ef4288440fc0e67b955fa0ffabdf48f8762577f304fd72ffd03131c5c840d
+size 12956685