1st cm

tokenization_qwen.py

@@ -13,15 +13,35 @@ import itertools
 
 import requests
 import unicodedata
-from typing import Collection, Dict, List, Set, Tuple, Union, Any, Callable, Optional
+from typing import (
+    Collection,
+    Dict,
+    List,
+    Set,
+    Tuple,
+    Union,
+    Any,
+    Callable,
+    Optional,
+)
 
 import tiktoken
 import numpy as np
 
 from transformers import PreTrainedTokenizer, AddedToken
 from transformers.utils import try_to_load_from_cache
-from transformers.tokenization_utils_base import BatchEncoding, PaddingStrategy, TruncationStrategy, \
-    TextInput, TextInputPair, PreTokenizedInput, PreTokenizedInputPair, TensorType, EncodedInput, EncodedInputPair
+from transformers.tokenization_utils_base import (
+    BatchEncoding,
+    PaddingStrategy,
+    TruncationStrategy,
+    TextInput,
+    TextInputPair,
+    PreTokenizedInput,
+    PreTokenizedInputPair,
+    TensorType,
+    EncodedInput,
+    EncodedInputPair,
+)
 
 import matplotlib.colors as mcolors
 from matplotlib.font_manager import FontProperties
@@ -40,10 +60,10 @@ IMEND = "<|im_end|>"
 # as different as possible to minimize the impact
 EXTRAS = tuple((f"<|extra_{i}|>" for i in range(205)))
 SPECIAL_TOKENS = (
-                     ENDOFTEXT,
-                     IMSTART,
-                     IMEND,
-                 ) + EXTRAS
+    ENDOFTEXT,
+    IMSTART,
+    IMEND,
+) + EXTRAS
 
 LANGUAGES = {
     "en": "english",
@@ -62,14 +82,16 @@ def _load_tiktoken_bpe(tiktoken_bpe_file: str) -> Dict[bytes, int]:
         contents = f.read()
     return {
         base64.b64decode(token): int(rank)
-        for token, rank in (line.split() for line in contents.splitlines() if line)
+        for token, rank in (
+            line.split() for line in contents.splitlines() if line
+        )
     }
 
 
 def _list_find(
-        input_list: List[Any],
-        candidates: Tuple[Any],
-        start: int = 0,
+    input_list: List[Any],
+    candidates: Tuple[Any],
+    start: int = 0,
 ):
     for i in range(start, len(input_list)):
         if input_list[i] in candidates:
@@ -78,12 +100,12 @@ def _list_find(
 
 
 def _replace_closed_tag(
-        input_tokens: List[Any],
-        start_tags: Union[Any, Tuple[Any]],
-        end_tags: Union[Any, Tuple[Any]],
-        inclusive_replace_func: Callable,
-        exclusive_replace_func: Callable = lambda x: x,
-        audio_info: Dict = None
+    input_tokens: List[Any],
+    start_tags: Union[Any, Tuple[Any]],
+    end_tags: Union[Any, Tuple[Any]],
+    inclusive_replace_func: Callable,
+    exclusive_replace_func: Callable = lambda x: x,
+    audio_info: Dict = None,
 ):
     if isinstance(start_tags, (str, int)):
         start_tags = (start_tags,)
@@ -98,12 +120,16 @@ def _replace_closed_tag(
         start = _list_find(input_tokens, start_tags, end)
         if start == -1:
             break
-        output_tokens.extend(exclusive_replace_func(input_tokens[end: start]))
+        output_tokens.extend(exclusive_replace_func(input_tokens[end:start]))
         tag_idx = start_tags.index(input_tokens[start])
         end = _list_find(input_tokens, (end_tags[tag_idx],), start)
         if end == -1:
             raise ValueError("Unclosed audio token")
-        output_tokens.extend(inclusive_replace_func(input_tokens[start: end + 1], audio_info, audio_idx))
+        output_tokens.extend(
+            inclusive_replace_func(
+                input_tokens[start : end + 1], audio_info, audio_idx
+            )
+        )
         end += 1
         audio_idx += 1
     output_tokens.extend(exclusive_replace_func(input_tokens[end:]))
@@ -116,12 +142,12 @@ class QWenTokenizer(PreTrainedTokenizer):
     vocab_files_names = VOCAB_FILES_NAMES
 
     def __init__(
-            self,
-            vocab_file,
-            errors="replace",
-            audio_start_tag='<audio>',
-            audio_end_tag='</audio>',
-            **kwargs,
+        self,
+        vocab_file,
+        errors="replace",
+        audio_start_tag="<audio>",
+        audio_end_tag="</audio>",
+        **kwargs,
     ):
         super().__init__(**kwargs)
         self.audio_start_tag = audio_start_tag
@@ -129,7 +155,7 @@ class QWenTokenizer(PreTrainedTokenizer):
         self.audio_pad_tag = "[[[AUDIO:modality]]]"
 
         self.AUDIO_ST = (
-            '[[[AUDIO:modality]]]',
+            "[[[AUDIO:modality]]]",
             # Transcription Tag
             "<|startoftranscript|>",  # Transcription
             "<|startofanalysis|>",  # Analysis
@@ -146,7 +172,9 @@ class QWenTokenizer(PreTrainedTokenizer):
             "<|notimestamps|>",
             "<|sil|>",
             "<|timestamps|>",
-            *[f"<|{i * 0.01:.2f}|>" for i in range(3001)],  # timestamps 0.00-30.00
+            *[
+                f"<|{i * 0.01:.2f}|>" for i in range(3001)
+            ],  # timestamps 0.00-30.00
             # Output Instruction
             "<|caption_audiocaps|>",  # Audiocaps caption style
             "<|caption_clotho|>",  # Clotho caption style
@@ -164,12 +192,15 @@ class QWenTokenizer(PreTrainedTokenizer):
             "<|endofword|>",
             "<|delim|>",  # delimiter of timestamps pair in audio grounding
             "<|emotion_recognition|>",  # emotion recognition
+            "<|emotion_transcription|>",
             "<|music_description|>",  # music description
             "<|note_analysis|>",  # note analysis
             "<|pitch|>",  # note analysis: pitch
             *[f"<|midi_pitch_{i}|>" for i in range(128)],  # midi pitch 0-127
             "<|velocity|>",  # note analysis: velocity
-            *[f"<|midi_velocity_{i}|>" for i in range(128)],  # midi velocity 0-127
+            *[
+                f"<|midi_velocity_{i}|>" for i in range(128)
+            ],  # midi velocity 0-127
             "<|sonic|>",  # note analysis:  sonic
             "<|instrument|>",  # note analysis:  instrument
             "<|speaker_meta|>",  # meta information of speaker
@@ -186,25 +217,28 @@ class QWenTokenizer(PreTrainedTokenizer):
             "<|entities|>",  # speech language understanding: entities
             "<|speech_edit|>",  # speech edit
             audio_start_tag,
-            audio_end_tag
+            audio_end_tag,
         )
 
         self.errors = errors  # how to handle errors in decoding
 
-        self.mergeable_ranks = _load_tiktoken_bpe(vocab_file)  # type: dict[bytes, int]
+        self.mergeable_ranks = _load_tiktoken_bpe(
+            vocab_file
+        )  # type: dict[bytes, int]
         self.special_tokens = {
             token: index
             for index, token in enumerate(
                 SPECIAL_TOKENS + self.AUDIO_ST, start=len(self.mergeable_ranks)
-
             )
         }
         self.audio_start_id = self.special_tokens[self.audio_start_tag]
         self.audio_end_id = self.special_tokens[self.audio_end_tag]
         self.audio_pad_id = self.special_tokens[self.audio_pad_tag]
-        print(f"audio_start_id: {self.audio_start_id}, "
-              f"audio_end_id: {self.audio_end_id}, "
-              f"audio_pad_id: {self.audio_pad_id}.")
+        print(
+            f"audio_start_id: {self.audio_start_id}, "
+            f"audio_end_id: {self.audio_end_id}, "
+            f"audio_pad_id: {self.audio_pad_id}."
+        )
 
         enc = tiktoken.Encoding(
             "Qwen",
@@ -213,7 +247,7 @@ class QWenTokenizer(PreTrainedTokenizer):
             special_tokens=self.special_tokens,
         )
         assert (
-                len(self.mergeable_ranks) + len(self.special_tokens) == enc.n_vocab
+            len(self.mergeable_ranks) + len(self.special_tokens) == enc.n_vocab
         ), f"{len(self.mergeable_ranks) + len(self.special_tokens)} != {enc.n_vocab} in encoding"
 
         self.decoder = {
@@ -230,7 +264,7 @@ class QWenTokenizer(PreTrainedTokenizer):
     def __getstate__(self):
         # for pickle lovers
         state = self.__dict__.copy()
-        del state['tokenizer']
+        del state["tokenizer"]
         return state
 
     def __setstate__(self, state):
@@ -251,7 +285,7 @@ class QWenTokenizer(PreTrainedTokenizer):
         return self.mergeable_ranks
 
     def convert_tokens_to_ids(
-            self, tokens: Union[bytes, str, List[Union[bytes, str]]]
+        self, tokens: Union[bytes, str, List[Union[bytes, str]]]
     ) -> List[int]:
         ids = []
         if isinstance(tokens, (str, bytes)):
@@ -266,13 +300,21 @@ class QWenTokenizer(PreTrainedTokenizer):
                 ids.append(self.mergeable_ranks.get(token))
         return ids
 
-    def _add_tokens(self, new_tokens: Union[List[str], List[AddedToken]], special_tokens: bool = False) -> int:
+    def _add_tokens(
+        self,
+        new_tokens: Union[List[str], List[AddedToken]],
+        special_tokens: bool = False,
+    ) -> int:
         if not special_tokens and new_tokens:
-            raise ValueError('Adding regular tokens is not supported')
+            raise ValueError("Adding regular tokens is not supported")
         for token in new_tokens:
-            surface_form = token.content if isinstance(token, AddedToken) else token
-            if surface_form not in SPECIAL_TOKENS  + self.AUDIO_ST:
-                raise ValueError('Adding unknown special tokens is not supported')
+            surface_form = (
+                token.content if isinstance(token, AddedToken) else token
+            )
+            if surface_form not in SPECIAL_TOKENS + self.AUDIO_ST:
+                raise ValueError(
+                    "Adding unknown special tokens is not supported"
+                )
         return 0
 
     def save_vocabulary(self, save_directory: str, **kwargs) -> Tuple[str]:
@@ -290,12 +332,12 @@ class QWenTokenizer(PreTrainedTokenizer):
         return (file_path,)
 
     def tokenize(
-            self,
-            text: str,
-            allowed_special: Union[Set, str] = "all",
-            disallowed_special: Union[Collection, str] = (),
-            audio_info: Dict = None,
-            **kwargs,
+        self,
+        text: str,
+        allowed_special: Union[Set, str] = "all",
+        disallowed_special: Union[Collection, str] = (),
+        audio_info: Dict = None,
+        **kwargs,
     ) -> List[Union[bytes, str]]:
         """
         Converts a string in a sequence of tokens.
@@ -321,61 +363,89 @@ class QWenTokenizer(PreTrainedTokenizer):
 
         # this implementation takes a detour: text -> token id -> token surface forms
         for t in self.tokenizer.encode(
-                text, allowed_special=allowed_special, disallowed_special=disallowed_special
+            text,
+            allowed_special=allowed_special,
+            disallowed_special=disallowed_special,
         ):
             tokens.append(self.decoder[t])
 
         def _encode_audiourl(audio_tokens, audio_info, audio_idx):
-            assert audio_tokens[0] == self.audio_start_tag and audio_tokens[-1] == self.audio_end_tag
-            audio_token_span = audio_info['audio_span_tokens'][audio_idx]
-            out_audio_tokens = [self.audio_start_tag] + [self.audio_pad_tag] * (audio_token_span - 2) + [
-                self.audio_end_tag]
+            assert (
+                audio_tokens[0] == self.audio_start_tag
+                and audio_tokens[-1] == self.audio_end_tag
+            )
+            audio_token_span = audio_info["audio_span_tokens"][audio_idx]
+            out_audio_tokens = (
+                [self.audio_start_tag]
+                + [self.audio_pad_tag] * (audio_token_span - 2)
+                + [self.audio_end_tag]
+            )
             return out_audio_tokens
 
-        return _replace_closed_tag(tokens, self.audio_start_tag, self.audio_end_tag, _encode_audiourl,
-                                   audio_info=audio_info)
+        return _replace_closed_tag(
+            tokens,
+            self.audio_start_tag,
+            self.audio_end_tag,
+            _encode_audiourl,
+            audio_info=audio_info,
+        )
 
     def _batch_encode_plus(
-            self,
-            batch_text_or_text_pairs: Union[
-                List[TextInput],
-                List[TextInputPair],
-                List[PreTokenizedInput],
-                List[PreTokenizedInputPair],
-                List[EncodedInput],
-                List[EncodedInputPair],
-            ],
-            add_special_tokens: bool = True,
-            padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD,
-            truncation_strategy: TruncationStrategy = TruncationStrategy.DO_NOT_TRUNCATE,
-            max_length: Optional[int] = None,
-            stride: int = 0,
-            is_split_into_words: bool = False,
-            pad_to_multiple_of: Optional[int] = None,
-            return_tensors: Optional[Union[str, TensorType]] = None,
-            return_token_type_ids: Optional[bool] = None,
-            return_attention_mask: Optional[bool] = None,
-            return_overflowing_tokens: bool = False,
-            return_special_tokens_mask: bool = False,
-            return_offsets_mapping: bool = False,
-            return_length: bool = False,
-            verbose: bool = True,
-            **kwargs,
+        self,
+        batch_text_or_text_pairs: Union[
+            List[TextInput],
+            List[TextInputPair],
+            List[PreTokenizedInput],
+            List[PreTokenizedInputPair],
+            List[EncodedInput],
+            List[EncodedInputPair],
+        ],
+        add_special_tokens: bool = True,
+        padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD,
+        truncation_strategy: TruncationStrategy = TruncationStrategy.DO_NOT_TRUNCATE,
+        max_length: Optional[int] = None,
+        stride: int = 0,
+        is_split_into_words: bool = False,
+        pad_to_multiple_of: Optional[int] = None,
+        return_tensors: Optional[Union[str, TensorType]] = None,
+        return_token_type_ids: Optional[bool] = None,
+        return_attention_mask: Optional[bool] = None,
+        return_overflowing_tokens: bool = False,
+        return_special_tokens_mask: bool = False,
+        return_offsets_mapping: bool = False,
+        return_length: bool = False,
+        verbose: bool = True,
+        **kwargs,
     ) -> BatchEncoding:
 
         def get_input_ids(text):
             if isinstance(text, str):
                 tokens = self.tokenize(text, **kwargs)
                 return self.convert_tokens_to_ids(tokens)
-            elif isinstance(text, (list, tuple)) and len(text) > 0 and isinstance(text[0], str):
+            elif (
+                isinstance(text, (list, tuple))
+                and len(text) > 0
+                and isinstance(text[0], str)
+            ):
                 if is_split_into_words:
                     tokens = list(
-                        itertools.chain(*(self.tokenize(t, is_split_into_words=True, **kwargs) for t in text))
+                        itertools.chain(
+                            *(
+                                self.tokenize(
+                                    t, is_split_into_words=True, **kwargs
+                                )
+                                for t in text
+                            )
+                        )
                     )
                     return self.convert_tokens_to_ids(tokens)
                 else:
                     return self.convert_tokens_to_ids(text)
-            elif isinstance(text, (list, tuple)) and len(text) > 0 and isinstance(text[0], int):
+            elif (
+                isinstance(text, (list, tuple))
+                and len(text) > 0
+                and isinstance(text[0], int)
+            ):
                 return text
             else:
                 raise ValueError(
@@ -392,18 +462,22 @@ class QWenTokenizer(PreTrainedTokenizer):
         input_ids = []
         audio_info = kwargs.pop("audio_info", None)
         for pair_id in range(len(batch_text_or_text_pairs)):
-            kwargs['audio_info'] = audio_info[pair_id]
+            kwargs["audio_info"] = audio_info[pair_id]
             ids_or_pair_ids = batch_text_or_text_pairs[pair_id]
             # for ids_or_pair_ids in batch_text_or_text_pairs:
             if not isinstance(ids_or_pair_ids, (list, tuple)):
                 ids, pair_ids = ids_or_pair_ids, None
-            elif is_split_into_words and not isinstance(ids_or_pair_ids[0], (list, tuple)):
+            elif is_split_into_words and not isinstance(
+                ids_or_pair_ids[0], (list, tuple)
+            ):
                 ids, pair_ids = ids_or_pair_ids, None
             else:
                 ids, pair_ids = ids_or_pair_ids
 
             first_ids = get_input_ids(ids)
-            second_ids = get_input_ids(pair_ids) if pair_ids is not None else None
+            second_ids = (
+                get_input_ids(pair_ids) if pair_ids is not None else None
+            )
             input_ids.append((first_ids, second_ids))
 
         batch_outputs = self._batch_prepare_for_model(
@@ -473,23 +547,35 @@ class QWenTokenizer(PreTrainedTokenizer):
         raise NotImplementedError
 
     def _decode(
-            self,
-            token_ids: Union[int, List[int]],
-            skip_special_tokens: bool = False,
-            errors: str = None,
-            **kwargs,
+        self,
+        token_ids: Union[int, List[int]],
+        skip_special_tokens: bool = False,
+        errors: str = None,
+        **kwargs,
     ) -> str:
         if isinstance(token_ids, int):
             token_ids = [token_ids]
         audio_info = kwargs.pop("audio_info", None)
 
         def _decode_audiourl(audio_token_ids, audio_info, audio_idx):
-            assert audio_token_ids[0] == self.audio_start_id and audio_token_ids[-1] == self.audio_end_id
+            assert (
+                audio_token_ids[0] == self.audio_start_id
+                and audio_token_ids[-1] == self.audio_end_id
+            )
             audio_url = audio_info["audio_urls"][audio_idx]
-            return [self.audio_start_id] + self.tokenizer.encode(audio_url) + [self.audio_end_id]
+            return (
+                [self.audio_start_id]
+                + self.tokenizer.encode(audio_url)
+                + [self.audio_end_id]
+            )
 
-        token_ids = _replace_closed_tag(token_ids, self.audio_start_id, self.audio_end_id, _decode_audiourl,
-                                        audio_info=audio_info)
+        token_ids = _replace_closed_tag(
+            token_ids,
+            self.audio_start_id,
+            self.audio_end_id,
+            _decode_audiourl,
+            audio_info=audio_info,
+        )
 
         if skip_special_tokens:
             token_ids = [i for i in token_ids if i < self.eod_id]
@@ -498,18 +584,32 @@ class QWenTokenizer(PreTrainedTokenizer):
     def to_list_format(self, text: str):
         text = unicodedata.normalize("NFC", text)
         token_ids = self.tokenizer.encode(
-            text, allowed_special=set(self.AUDIO_ST + (ENDOFTEXT,)))
+            text, allowed_special=set(self.AUDIO_ST + (ENDOFTEXT,))
+        )
 
         def _encode_audio_info(tokens):
             if len(tokens) == 0:
                 return []
-            if tokens[0] == self.audio_start_id and tokens[-1] == self.audio_end_id:
-                key = 'audio'
+            if (
+                tokens[0] == self.audio_start_id
+                and tokens[-1] == self.audio_end_id
+            ):
+                key = "audio"
             else:
-                _tobytes = lambda x: x.encode('utf-8') if isinstance(x, str) else x
-                return [{'text': b''.join(map(_tobytes, map(self.decoder.get, tokens))).decode('utf-8')}]
-            _tobytes = lambda x: x.encode('utf-8') if isinstance(x, str) else x
-            val = b''.join(map(_tobytes, map(self.decoder.get, tokens[1:-1]))).decode('utf-8')
+                _tobytes = lambda x: (
+                    x.encode("utf-8") if isinstance(x, str) else x
+                )
+                return [
+                    {
+                        "text": b"".join(
+                            map(_tobytes, map(self.decoder.get, tokens))
+                        ).decode("utf-8")
+                    }
+                ]
+            _tobytes = lambda x: x.encode("utf-8") if isinstance(x, str) else x
+            val = b"".join(
+                map(_tobytes, map(self.decoder.get, tokens[1:-1]))
+            ).decode("utf-8")
             return [{key: val}]
 
         return _replace_closed_tag(
@@ -521,21 +621,25 @@ class QWenTokenizer(PreTrainedTokenizer):
         )
 
     def from_list_format(self, list_format: List[Dict]):
-        text = ''
+        text = ""
         num_audios = 0
         for ele in list_format:
-            if 'audio' in ele:
+            if "audio" in ele:
                 num_audios += 1
-                text += f'Audio {num_audios}:'
-                text += self.audio_start_tag + ele['audio'] + self.audio_end_tag
-                text += '\n'
-            elif 'text' in ele:
-                text += ele['text']
-            elif 'box' in ele:
-                if 'ref' in ele:
-                    text += self.ref_start_tag + ele['ref'] + self.ref_end_tag
-                for box in ele['box']:
-                    text += self.box_start_tag + '(%d,%d),(%d,%d)' % (box[0], box[1], box[2], box[3]) + self.box_end_tag
+                text += f"Audio {num_audios}:"
+                text += self.audio_start_tag + ele["audio"] + self.audio_end_tag
+                text += "\n"
+            elif "text" in ele:
+                text += ele["text"]
+            elif "box" in ele:
+                if "ref" in ele:
+                    text += self.ref_start_tag + ele["ref"] + self.ref_end_tag
+                for box in ele["box"]:
+                    text += (
+                        self.box_start_tag
+                        + "(%d,%d),(%d,%d)" % (box[0], box[1], box[2], box[3])
+                        + self.box_end_tag
+                    )
             else:
                 raise ValueError("Unsupport element: " + str(ele))
         return text
@@ -549,12 +653,16 @@ class QWenTokenizer(PreTrainedTokenizer):
         if len(audio_urls) > 0:
             audios, audio_lens, audio_span_tokens = [], [], []
             for audio_path in audio_urls:
-                if audio_path.startswith("http://") or audio_path.startswith("https://"):  # http
+                if audio_path.startswith("http://") or audio_path.startswith(
+                    "https://"
+                ):  # http
                     data = bytes(requests.get(audio_path, stream=True).content)
                     audio = load_bytesio_audio(data)
                 else:
                     audio = load_audio(audio_path)
-                L = (audio.shape[0] if audio.shape[0] <= 480000 else 480000)  # max_length < 30s
+                L = (
+                    audio.shape[0] if audio.shape[0] <= 480000 else 480000
+                )  # max_length < 30s
                 mel_len = L // 160
                 audio = pad_or_trim(audio.flatten())
                 mel = log_mel_spectrogram(audio)
@@ -563,17 +671,16 @@ class QWenTokenizer(PreTrainedTokenizer):
                 audio_len = [audio_len_after_cnn, audio_token_num]
                 audios.append(mel)
                 audio_lens.append(audio_len)
-                audio_span_tokens.append(audio_token_num + 2)  # add audio bos eos
+                audio_span_tokens.append(
+                    audio_token_num + 2
+                )  # add audio bos eos
             input_audio_lengths = torch.IntTensor(audio_lens)
             input_audios = torch.stack(audios, dim=0)
-            return {"input_audios": input_audios,
-                    "input_audio_lengths": input_audio_lengths,
-                    "audio_span_tokens": audio_span_tokens,
-                    "audio_urls": audio_urls}
+            return {
+                "input_audios": input_audios,
+                "input_audio_lengths": input_audio_lengths,
+                "audio_span_tokens": audio_span_tokens,
+                "audio_urls": audio_urls,
+            }
         else:
             return None
-
-
-
-
-