new file: codecmanipulator.py

codecmanipulator.py

@@ -0,0 +1,203 @@
+import json
+import numpy as np
+import einops
+
+
+class CodecManipulator(object):
+    r"""
+    **mm tokenizer v0.1**
+    see codeclm/hf/mm_tokenizer_v0.1_hf/id2vocab.json
+
+    text tokens: 
+        llama tokenizer 0~31999
+    
+    special tokens: "32000": "<EOD>", "32001": "<SOA>", "32002": "<EOA>", "32003": "<SOI>", "32004": "<EOI>", "32005": "<SOV>", "32006": "<EOV>", "32007": "<s_local>", "32008": "<e_local>", "32009": "<s_global>", "32010": "<e_global>", "32011": "<semantic>", "32012": "<acoustic>", "32013": "<low_level>", "32014": "<dac_16k>", "32015": "<dac_44k>", "32016": "<xcodec>", "32017": "<placeholder>", "32018": "<semantic_mert>", "32019": "<semantic_hubert>", "32020": "<visual>", "32021": "<semanticodec>"
+
+    mm tokens:
+        dac_16k: 4 codebook, 1024 vocab, 32022 - 36117
+        dac_44k: 9 codebook, 1024 vocab, 36118 - 45333
+        xcodec: 12 codebook, 1024 vocab, 45334 - 57621
+        semantic mert: 1024, 57622 - 58645
+        semantic hubert: 512, 58646 - 59157
+        visual: 64000, not included in v0.1
+        semanticodec 100tps 16384: semantic=16384, 59158 - 75541, acoustic=8192, 75542 - 83733
+    """
+    def __init__(self, codec_type, quantizer_begin=None, n_quantizer=None, teacher_forcing=False, data_feature="codec"):
+        self.codec_type = codec_type
+        self.mm_v0_2_cfg = {
+            "dac16k": {"codebook_size": 1024, "num_codebooks": 4, "global_offset": 32022, "sep": ["<dac_16k>"], "fps": 50},
+            "dac44k": {"codebook_size": 1024, "num_codebooks": 9, "global_offset": 36118, "sep": ["<dac_44k>"]},
+            "xcodec": {"codebook_size": 1024, "num_codebooks": 12, "global_offset": 45334, "sep": ["<xcodec>"], "fps": 50},
+            "mert": {"codebook_size": 1024, "global_offset": 57622, "sep": ["<semantic_mert>"]},
+            "hubert": {"codebook_size": 512, "global_offset": 58646, "sep": ["<semantic_hubert>"]},
+            "semantic/s": {"codebook_size": 16384, "num_codebooks": 1, "global_offset": 59158, "sep": ["<semanticodec>", "<semantic>"]},
+            "semantic/a": {"codebook_size": 8192, "num_codebooks": 1, "global_offset": 75542, "sep": ["<semanticodec>", "<acoustic>"]},
+            "semanticodec": {"codebook_size": [16384, 8192], "num_codebooks": 2, "global_offset": 59158, "sep": ["<semanticodec>"], "fps": 50},
+            "special_tokens": {
+                '<EOD>': 32000, '<SOA>': 32001, '<EOA>': 32002, '<SOI>': 32003, '<EOI>': 32004, '<SOV>': 32005, '<EOV>': 32006, '<s_local>': 32007, '<e_local>': 32008, '<s_global>': 32009, '<e_global>': 32010, '<semantic>': 32011, '<acoustic>': 32012, '<stage_1>': 32013, '<dac_16k>': 32014, '<dac_44k>': 32015, '<xcodec>': 32016, '<stage_2>': 32017, '<semantic_mert>': 32018, '<semantic_hubert>': 32019, '<visual>': 32020, '<semanticodec>': 32021
+            },
+            "metadata": {
+                "len": 83734,
+                "text_range": [0, 31999],
+                "special_range": [32000, 32021],
+                "mm_range": [32022, 83733]
+            },
+            "codec_range": {
+                "dac16k": [32022, 36117],
+                "dac44k": [36118, 45333],
+                "xcodec": [45334, 57621],
+                # "hifi16k": [53526, 57621],
+                "mert": [57622, 58645],
+                "hubert": [58646, 59157],
+                "semantic/s": [59158, 75541],
+                "semantic/a": [75542, 83733],
+                "semanticodec": [59158, 83733]
+            }
+        }
+        self.sep = self.mm_v0_2_cfg[self.codec_type]["sep"]
+        self.sep_ids = [self.mm_v0_2_cfg["special_tokens"][s] for s in self.sep]
+        self.codebook_size = self.mm_v0_2_cfg[self.codec_type]["codebook_size"]
+        self.num_codebooks = self.mm_v0_2_cfg[self.codec_type]["num_codebooks"]
+        self.global_offset = self.mm_v0_2_cfg[self.codec_type]["global_offset"]
+        self.fps = self.mm_v0_2_cfg[self.codec_type]["fps"] if "fps" in self.mm_v0_2_cfg[self.codec_type] else None
+
+        self.quantizer_begin = quantizer_begin if quantizer_begin is not None else 0
+        self.n_quantizer = n_quantizer if n_quantizer is not None else self.num_codebooks  
+        self.teacher_forcing = teacher_forcing 
+        self.data_feature = data_feature
+
+
+    def offset_tok_ids(self, x, global_offset=0, codebook_size=2048, num_codebooks=4):
+        """
+        x: (K, T)
+        """
+        if isinstance(codebook_size, int):
+            assert x.max() < codebook_size, f"max(x)={x.max()}, codebook_size={codebook_size}"
+        elif isinstance(codebook_size, list):
+            for i, cs in enumerate(codebook_size):
+                assert x[i].max() < cs, f"max(x)={x[i].max()}, codebook_size={cs}, layer_id={i}"
+        else:
+            raise ValueError(f"codebook_size={codebook_size}")
+        assert x.min() >= 0, f"min(x)={x.min()}"
+        assert x.shape[0] == num_codebooks or x.shape[0] == self.n_quantizer, \
+            f"x.shape[0]={x.shape[0]}, num_codebooks={num_codebooks}, n_quantizer={self.n_quantizer}"
+
+        _x = x.copy()
+        _x = _x.astype(np.uint32)
+        cum_offset = 0
+        quantizer_begin = self.quantizer_begin
+        quantizer_end = quantizer_begin+self.n_quantizer
+        for k in range(self.quantizer_begin, quantizer_end): # k: quantizer_begin to quantizer_end - 1
+            if isinstance(codebook_size, int):
+                _x[k] += global_offset + k * codebook_size
+            elif isinstance(codebook_size, list):
+                _x[k] += global_offset + cum_offset
+                cum_offset += codebook_size[k]
+            else:
+                raise ValueError(f"codebook_size={codebook_size}")
+        return _x[quantizer_begin:quantizer_end]
+
+    def unoffset_tok_ids(self, x, global_offset=0, codebook_size=2048, num_codebooks=4):
+        """
+        x: (K, T)
+        """
+        if isinstance(codebook_size, int):
+            assert x.max() < global_offset + codebook_size * num_codebooks, f"max(x)={x.max()}, codebook_size={codebook_size}"
+        elif isinstance(codebook_size, list):
+            assert x.max() < global_offset + sum(codebook_size), f"max(x)={x.max()}, codebook_size={codebook_size}"
+        assert x.min() >= global_offset, f"min(x)={x.min()}, global_offset={global_offset}"
+        assert x.shape[0] == num_codebooks or x.shape[0] == self.n_quantizer, \
+            f"x.shape[0]={x.shape[0]}, num_codebooks={num_codebooks}, n_quantizer={self.n_quantizer}"
+        
+        _x = x.copy()
+        _x = _x.astype(np.uint32)
+        cum_offset = 0
+        quantizer_begin = self.quantizer_begin
+        quantizer_end = quantizer_begin+self.n_quantizer
+        for k in range(quantizer_begin, quantizer_end):
+            if isinstance(codebook_size, int):
+                _x[k-quantizer_begin] -= global_offset + k * codebook_size
+            elif isinstance(codebook_size, list):
+                _x[k-quantizer_begin] -= global_offset + cum_offset
+                cum_offset += codebook_size[k]
+            else:
+                raise ValueError(f"codebook_size={codebook_size}")
+        return _x
+
+    def flatten(self, x):
+        if len(x.shape) > 2:
+            x = x.squeeze()
+        assert x.shape[0] == self.num_codebooks or x.shape[0] == self.n_quantizer, \
+            f"x.shape[0]={x.shape[0]}, num_codebooks={self.num_codebooks}, n_quantizer={self.n_quantizer}"
+        return einops.rearrange(x, 'K T -> (T K)')
+
+    def unflatten(self, x, n_quantizer=None):
+        x = x.squeeze()
+        assert len(x.shape) == 1
+        assert x.shape[0] % self.num_codebooks == 0 or x.shape[0] % self.n_quantizer == 0, \
+            f"x.shape[0]={x.shape[0]}, num_codebooks={self.num_codebooks}, n_quantizer={self.n_quantizer}"
+        if n_quantizer!=self.num_codebooks:
+            return einops.rearrange(x, '(T K) -> K T', K=n_quantizer)
+        return einops.rearrange(x, '(T K) -> K T', K=self.num_codebooks)
+    
+    # def check_codec_type_from_path(self, path):
+    #     if self.codec_type == "hifi16k":
+    #         assert "academicodec_hifi_16k_320d_large_uni" in path
+    
+    def get_codec_type_from_range(self, ids):
+        ids_range = [ids.min(), ids.max()]
+        codec_range = self.mm_v0_2_cfg["codec_range"]
+        for codec_type, r in codec_range.items():
+            if ids_range[0] >= r[0] and ids_range[1] <= r[1]:
+                return codec_type
+        raise ValueError(f"ids_range={ids_range}, codec_range={codec_range}")
+
+    def npy2ids(self, npy):
+        if isinstance(npy, str):
+            data = np.load(npy)
+        elif isinstance(npy, np.ndarray):
+            data = npy
+        else:
+            raise ValueError(f"not supported type: {type(npy)}")
+        # data = data.squeeze()
+
+        assert len(data.shape)==2,  f'data shape: {data.shape} is not (n_codebook, seq_len)'
+        data = self.offset_tok_ids(
+            data, 
+            global_offset=self.global_offset, 
+            codebook_size=self.codebook_size, 
+            num_codebooks=self.num_codebooks, 
+        )
+        data = self.flatten(data)
+        codec_range = self.get_codec_type_from_range(data)
+        assert codec_range == self.codec_type, f"get_codec_type_from_range(data)={codec_range}, self.codec_type={self.codec_type}"
+        data = data.tolist()
+        return data
+    
+    def ids2npy(self, token_ids):
+        # make sure token_ids starts with codebook 0
+        if isinstance(self.codebook_size, int):
+            codebook_0_range = (self.global_offset + self.quantizer_begin*self.codebook_size, self.global_offset + (self.quantizer_begin+1)*self.codebook_size)
+        elif isinstance(self.codebook_size, list):
+            codebook_0_range = (self.global_offset, self.global_offset + self.codebook_size[0])
+        assert token_ids[0] >= codebook_0_range[0] \
+            and token_ids[0] < codebook_0_range[1], f"token_ids[0]={token_ids[self.quantizer_begin]}, codebook_0_range={codebook_0_range}"
+        data = np.array(token_ids)
+        data = self.unflatten(data, n_quantizer=self.n_quantizer)
+        data = self.unoffset_tok_ids(
+            data, 
+            global_offset=self.global_offset, 
+            codebook_size=self.codebook_size, 
+            num_codebooks=self.num_codebooks, 
+        )
+        return data
+
+    def npy_to_json_str(self, npy_path):
+        data = self.npy2ids(npy_path)
+        return json.dumps({"text": data, "src": npy_path, "codec": self.codec_type})
+    
+    def sep(self):
+        return ''.join(self.sep)
+    
+    def sep_ids(self):
+        return self.sep_ids

mmtokenizer.py

@@ -0,0 +1,367 @@
+from abc import ABC
+from abc import abstractmethod
+
+
+class AbstractTokenizer(ABC):
+    """Abstract class for tokenizer."""
+
+    def __init__(self, name):
+        self.name = name
+        super().__init__()
+
+    @property
+    @abstractmethod
+    def vocab_size(self):
+        pass
+
+    @property
+    @abstractmethod
+    def vocab(self):
+        """Dictionary from vocab text token to id token."""
+        pass
+
+    @property
+    @abstractmethod
+    def inv_vocab(self):
+        """Dictionary from vocab id token to text token."""
+        pass
+
+    @abstractmethod
+    def tokenize(self, text):
+        pass
+
+    def detokenize(self, token_ids):
+        raise NotImplementedError('detokenizer is not implemented for {} '
+                                  'tokenizer'.format(self.name))
+
+    @property
+    def cls(self):
+        raise NotImplementedError('CLS is not provided for {} '
+                                  'tokenizer'.format(self.name))
+
+    @property
+    def sep(self):
+        raise NotImplementedError('SEP is not provided for {} '
+                                  'tokenizer'.format(self.name))
+
+    @property
+    def pad(self):
+        raise NotImplementedError('PAD is not provided for {} '
+                                  'tokenizer'.format(self.name))
+
+    @property
+    def eod(self):
+        raise NotImplementedError('EOD is not provided for {} '
+                                  'tokenizer'.format(self.name))
+
+    @property
+    def mask(self):
+        raise NotImplementedError('MASK is not provided for {} '
+                                  'tokenizer'.format(self.name))
+
+
+class _SentencePieceTokenizer(AbstractTokenizer):
+    """SentencePieceTokenizer-Megatron wrapper"""
+
+    def __init__(self, model_file, vocab_extra_ids=0):
+        name = 'SentencePieceTokenizer'
+        super().__init__(name)
+
+        import sentencepiece
+        self.tokenizer = sentencepiece.SentencePieceProcessor(model_file=model_file)
+        self._initalize(vocab_extra_ids)
+
+    def _populate_vocab(self):
+        self._vocab = {}
+        self._inv_vocab = {}
+
+        for i in range(len(self.tokenizer)):
+            t = self.tokenizer.id_to_piece(i)
+            self._inv_vocab[i] = t
+            self._vocab[t] = i
+
+    def _initalize(self, vocab_extra_ids):
+        self._populate_vocab()
+        self._special_tokens = {}
+        self._inv_special_tokens = {}
+
+        self._t5_tokens = []
+
+        def _add_special_token(t):
+            if t not in self._vocab:
+                next_id = len(self._vocab)
+                self._vocab[t] = next_id
+                self._inv_vocab[next_id] = t
+            self._special_tokens[t] = self._vocab[t]
+            self._inv_special_tokens[self._vocab[t]] = t
+
+        _add_special_token('<CLS>')
+        self._cls_id = self._vocab['<CLS>']
+        _add_special_token('<SEP>')
+        self._sep_id = self._vocab['<SEP>']
+        _add_special_token('<EOD>')
+        self._eod_id = self._vocab['<EOD>']
+        _add_special_token('<MASK>')
+        self._mask_id = self._vocab['<MASK>']
+
+        pad_id = self.tokenizer.pad_id()
+        try:
+            pad_token = self.tokenizer.id_to_piece(pad_id)
+        except IndexError:
+            pad_token = '<PAD>'
+        _add_special_token(pad_token)
+        self._pad_id = self._vocab[pad_token]
+
+        bos_id = self.tokenizer.bos_id()
+        try:
+            bos_token = self.tokenizer.id_to_piece(bos_id)
+        except IndexError:
+            bos_token = '<BOS>'
+        _add_special_token(bos_token)
+        self._bos_id = self._vocab[bos_token]
+
+        eos_id = self.tokenizer.eos_id()
+        try:
+            eos_token = self.tokenizer.id_to_piece(eos_id)
+        except IndexError:
+            eos_token = '<EOS>'
+        _add_special_token(eos_token)
+        self._eos_id = self._vocab[eos_token]
+
+        for i in range(vocab_extra_ids):
+            t = "<extra_id_{}>".format(i)
+            _add_special_token(t)
+            self._t5_tokens += [t]
+
+    @property
+    def vocab_size(self):
+        return len(self._vocab)
+
+    @property
+    def vocab(self):
+        return self._vocab
+
+    @property
+    def inv_vocab(self):
+        return self._inv_vocab
+
+    @property
+    def decoder(self):
+        return self._inv_vocab
+
+    @property
+    def encoder(self):
+        return self._vocab
+
+    # From:
+    # https://github.com/NVIDIA/NeMo/blob/c8fa217e811d60d11d014827c7f3845ff6c99ae7/nemo/collections/common/tokenizers/sentencepiece_tokenizer.py#L89
+    def tokenize(self, text):
+        ids = []
+        idx = 0
+
+        while 1:
+            indices = {}
+            for token in self._special_tokens:
+                try:
+                    indices[token] = text[idx:].index(token)
+                except ValueError:
+                    continue
+            if len(indices) == 0:
+                break
+
+            next_token = min(indices, key=indices.get)
+            next_idx = idx + indices[next_token]
+
+            ids.extend(self.tokenizer.encode_as_ids(text[idx:next_idx]))
+            ids.append(self._special_tokens[next_token])
+            idx = next_idx + len(next_token)
+
+        ids.extend(self.tokenizer.encode_as_ids(text[idx:]))
+        return ids
+
+    # From:
+    # https://github.com/NVIDIA/NeMo/blob/c8fa217e811d60d11d014827c7f3845ff6c99ae7/nemo/collections/common/tokenizers/sentencepiece_tokenizer.py#L125
+    def detokenize(self, ids):
+        text = ""
+        last_i = 0
+
+        for i, id in enumerate(ids):
+            if id in self._inv_special_tokens:
+                text += self.tokenizer.decode_ids(ids[last_i:i]) + " "
+                text += self._inv_special_tokens[id] + " "
+                last_i = i + 1
+
+        text += self.tokenizer.decode_ids(ids[last_i:])
+        return text
+
+    @property
+    def cls(self):
+        return self._cls_id
+
+    @property
+    def sep(self):
+        return self._sep_id
+
+    @property
+    def pad(self):
+        return self._pad_id
+
+    @property
+    def bos_token_id(self):
+        return self._bos_id
+
+    @property
+    def bos(self):
+        return self._bos_id
+
+    @property
+    def eod(self):
+        return self._eod_id
+
+    @property
+    def eos_token_id(self):
+        return self._eos_id
+
+    @property
+    def eos(self):
+        return self._eos_id
+
+    @property
+    def mask(self):
+        return self._mask_id
+
+    @property
+    def additional_special_tokens_ids(self):
+        return [self.vocab[k] for k in self._t5_tokens]
+
+class _MMSentencePieceTokenizer(_SentencePieceTokenizer):
+    """SentencePieceTokenizer-Megatron wrapper"""
+
+    def __init__(self, model_file, vocab_extra_ids=0):
+        super().__init__(model_file, vocab_extra_ids)
+
+
+    def _initalize(self, vocab_extra_ids):
+        self._populate_vocab()
+        self._special_tokens = {}
+        self._inv_special_tokens = {}
+
+        self._t5_tokens = []
+
+        def _add_special_token(t):
+            if t not in self._vocab:
+                next_id = len(self._vocab)
+                self._vocab[t] = next_id
+                self._inv_vocab[next_id] = t
+            self._special_tokens[t] = self._vocab[t]
+            self._inv_special_tokens[self._vocab[t]] = t
+
+        _add_special_token('<CLS>')
+        self._cls_id = self._vocab['<CLS>']
+        _add_special_token('<SEP>')
+        self._sep_id = self._vocab['<SEP>']
+        _add_special_token('<EOD>')
+        self._eod_id = self._vocab['<EOD>']
+        _add_special_token('<MASK>')
+        self._mask_id = self._vocab['<MASK>']
+
+        _add_special_token('<SOA>')
+        self._soa_id = self._vocab['<SOA>']
+        _add_special_token('<EOA>')
+        self._eoa_id = self._vocab['<EOA>']
+        _add_special_token('<SOV>')
+        self._sov_id = self._vocab['<SOV>']
+        _add_special_token('<EOV>')
+        self._eov_id = self._vocab['<EOV>']
+        _add_special_token('<SOI>')
+        self._soi_id = self._vocab['<SOI>']
+        _add_special_token('<EOI>')
+        self._eoi_id = self._vocab['<EOI>']
+        _add_special_token('<s_local>')
+        self._s_local_id = self._vocab['<s_local>']
+        _add_special_token('<e_local>')
+        self._e_local_id = self._vocab['<e_local>']
+        _add_special_token('<s_global>')
+        self._s_global_id = self._vocab['<s_global>']
+        _add_special_token('<e_global>')
+        self._e_global_id = self._vocab['<e_global>']
+        _add_special_token('<stage_1>')
+        self._stage_1_id = self._vocab['<stage_1>']
+        _add_special_token('<stage_2>')
+        self._stage_2_id = self._vocab['<stage_2>']
+        pad_id = self.tokenizer.pad_id()
+        try:
+            pad_token = self.tokenizer.id_to_piece(pad_id)
+        except IndexError:
+            pad_token = '<PAD>'
+        _add_special_token(pad_token)
+        self._pad_id = self._vocab[pad_token]
+
+        bos_id = self.tokenizer.bos_id()
+        try:
+            bos_token = self.tokenizer.id_to_piece(bos_id)
+        except IndexError:
+            bos_token = '<BOS>'
+        _add_special_token(bos_token)
+        self._bos_id = self._vocab[bos_token]
+
+        eos_id = self.tokenizer.eos_id()
+        try:
+            eos_token = self.tokenizer.id_to_piece(eos_id)
+        except IndexError:
+            eos_token = '<EOS>'
+        _add_special_token(eos_token)
+        self._eos_id = self._vocab[eos_token]
+
+        for i in range(vocab_extra_ids):
+            t = "<extra_id_{}>".format(i)
+            _add_special_token(t)
+            self._t5_tokens += [t]
+    
+    @property
+    def soa(self):
+        return self._soa_id
+    
+    @property
+    def eoa(self):
+        return self._eoa_id
+    
+    @property
+    def sov(self):
+        return self._sov_id
+    
+    @property
+    def eov(self):
+        return self._eov_id
+    
+    @property
+    def soi(self):
+        return self._soi_id
+    
+    @property
+    def eoi(self):
+        return self._eoi_id
+    
+    @property
+    def s_local(self):
+        return self._s_local_id
+    
+    @property
+    def e_local(self):
+        return self._e_local_id
+    
+    @property
+    def s_global(self):
+        return self._s_global_id
+
+    @property
+    def e_global(self):
+        return self._e_global_id
+    
+    @property
+    def stage_1(self):
+        return self._stage_1_id
+
+    @property
+    def stage_2(self):
+        return self._stage_2_id