diff --git a/modelscope/models/nlp/qwen/backbone.py b/modelscope/models/nlp/qwen/backbone.py index 3e838402..ea03bf49 100644 --- a/modelscope/models/nlp/qwen/backbone.py +++ b/modelscope/models/nlp/qwen/backbone.py @@ -5,7 +5,7 @@ import importlib import math -from typing import Optional, Tuple, Union +from typing import TYPE_CHECKING, Callable, List, Optional, Tuple, Union import torch import torch.nn.functional as F @@ -13,7 +13,10 @@ import torch.utils.checkpoint from torch import nn from torch.cuda.amp import autocast from torch.nn import CrossEntropyLoss -from transformers import PreTrainedTokenizer +from transformers import (GenerationConfig, PreTrainedTokenizer, + StoppingCriteriaList) +from transformers.generation.logits_process import LogitsProcessorList +from transformers.generation.utils import GenerateOutput from transformers.modeling_outputs import (BaseModelOutputWithPast, CausalLMOutputWithPast) from transformers.modeling_utils import PreTrainedModel @@ -30,26 +33,37 @@ from modelscope.utils.constant import Tasks from modelscope.utils.logger import get_logger from ... import MODELS from .configuration import QWenConfig +from .qwen_generation_utils import (HistoryType, StopWordsLogitsProcessor, + decode_tokens, get_stop_words_ids, + make_context) + +if TYPE_CHECKING: + from transformers.generation.streamers import BaseStreamer try: from einops import rearrange except ImportError: rearrange = None + try: from flash_attn.layers.rotary import apply_rotary_emb_func from einops import rearrange + use_flash_rotary = True - print('use flash_attn rotary') except ImportError: use_flash_rotary = False - print('import flash_attn rotary fail') + print( + 'Warning: import flash_attn rotary fail, please install FlashAttention rotary to get better performance ' + 'https://github.com/Dao-AILab/flash-attention/tree/main/csrc/rotary') try: from flash_attn.ops.rms_norm import rms_norm - print('use flash_attn rms_norm') except ImportError: rms_norm = None - print('import flash_attn rms_norm fail') + print( + 'Warning: import flash_attn rms_norm fail, please install FlashAttention layer_norm to get better performance ' + 'https://github.com/Dao-AILab/flash-attention/tree/main/csrc/layer_norm' + ) logger = get_logger() @@ -62,21 +76,24 @@ try: from flash_attn.flash_attn_interface import flash_attn_unpadded_func except ImportError: flash_attn_unpadded_func = None + print('Warning: import flash_attn fail, please install FlashAttention ' + 'https://github.com/Dao-AILab/flash-attention') class FlashSelfAttention(torch.nn.Module): - def __init__(self, - causal=False, - softmax_scale=None, - attention_dropout=0.0, - device=None, - dtype=None): + def __init__( + self, + causal=False, + softmax_scale=None, + attention_dropout=0.0, + ): super().__init__() assert flash_attn_unpadded_func is not None, ( 'Please install FlashAttention first, ' 'e.g., with pip install flash-attn') - assert rearrange is not None, 'Please install einops first, e.g., with pip install einops' + assert (rearrange is not None + ), 'Please install einops first, e.g., with pip install einops' self.causal = causal self.softmax_scale = softmax_scale self.dropout_p = attention_dropout @@ -89,10 +106,12 @@ class FlashSelfAttention(torch.nn.Module): seqlen_k = k.shape[1] q, k, v = [rearrange(x, 'b s ... -> (b s) ...') for x in [q, k, v]] cu_seqlens_q = torch.arange( - 0, (batch_size + 1) * seqlen_q, + 0, + (batch_size + 1) * seqlen_q, step=seqlen_q, dtype=torch.int32, - device=q.device) + device=q.device, + ) if self.training: assert seqlen_k == seqlen_q @@ -102,10 +121,12 @@ class FlashSelfAttention(torch.nn.Module): else: is_causal = seqlen_q == seqlen_k cu_seqlens_k = torch.arange( - 0, (batch_size + 1) * seqlen_k, + 0, + (batch_size + 1) * seqlen_k, step=seqlen_k, dtype=torch.int32, - device=q.device) + device=q.device, + ) self.dropout_p = 0 output = flash_attn_unpadded_func( q, @@ -117,7 +138,8 @@ class FlashSelfAttention(torch.nn.Module): seqlen_k, self.dropout_p, softmax_scale=self.softmax_scale, - causal=is_causal) + causal=is_causal, + ) output = rearrange(output, '(b s) ... -> b s ...', b=batch_size) return output @@ -156,15 +178,16 @@ class QWenAttention(nn.Module): self.projection_size = config.kv_channels * config.num_attention_heads assert self.projection_size % config.num_attention_heads == 0 - self.hidden_size_per_attention_head = \ - self.projection_size // config.num_attention_heads + self.hidden_size_per_attention_head = ( + self.projection_size // config.num_attention_heads) self.c_attn = nn.Linear(config.hidden_size, 3 * self.projection_size) self.c_proj = nn.Linear( config.hidden_size, self.projection_size, bias=not config.no_bias) - if self.use_flash_attn: + self.is_fp32 = not (config.bf16 or config.fp16) + if self.use_flash_attn and flash_attn_unpadded_func is not None and not self.is_fp32: self.core_attention_flash = FlashSelfAttention( causal=True, attention_dropout=config.attn_pdrop) @@ -179,8 +202,17 @@ class QWenAttention(nn.Module): dim = ( self.rotary_ndims if self.rotary_ndims is not None else self.hidden_size_per_attention_head) - self.rotary_emb = RotaryEmbedding( - dim, base=config.rotary_emb_base, ntk_alpha=config.ntk_alpha) + self.rotary_emb = RotaryEmbedding(dim, base=config.rotary_emb_base) + + self.use_dynamic_ntk = config.use_dynamic_ntk + self.use_logn_attn = config.use_logn_attn + + logn_list = [ + math.log(i, self.seq_length) if i > self.seq_length else 1 + for i in range(1, 32768) + ] + self.logn_tensor = torch.Tensor(logn_list)[None, :, None, None] + self._ntk_cached = 1.0 self.attn_dropout = nn.Dropout(config.attn_pdrop) @@ -192,7 +224,8 @@ class QWenAttention(nn.Module): [], value.size(-1)**0.5, dtype=attn_weights.dtype, - device=attn_weights.device) + device=attn_weights.device, + ) query_length, key_length = query.size(-2), key.size(-2) causal_mask = self.bias[:, :, key_length @@ -231,7 +264,8 @@ class QWenAttention(nn.Module): q_seq_len, k_seq_len, dtype=torch.float32, - device=query.device) + device=query.device, + ) scale_factor = 1.0 if self.scale_attn_weights: @@ -283,15 +317,17 @@ class QWenAttention(nn.Module): new_shape = tensor.size()[:-2] + (num_heads * attn_head_size, ) return tensor.view(new_shape) - def forward(self, - hidden_states: Optional[Tuple[torch.FloatTensor]], - layer_past: Optional[Tuple[torch.Tensor]] = None, - attention_mask: Optional[torch.FloatTensor] = None, - head_mask: Optional[torch.FloatTensor] = None, - encoder_hidden_states: Optional[torch.Tensor] = None, - encoder_attention_mask: Optional[torch.FloatTensor] = None, - output_attentions: Optional[bool] = False, - use_cache: Optional[bool] = False): + def forward( + self, + hidden_states: Optional[Tuple[torch.FloatTensor]], + layer_past: Optional[Tuple[torch.Tensor]] = None, + attention_mask: Optional[torch.FloatTensor] = None, + head_mask: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.Tensor] = None, + encoder_attention_mask: Optional[torch.FloatTensor] = None, + output_attentions: Optional[bool] = False, + use_cache: Optional[bool] = False, + ): mixed_x_layer = self.c_attn(hidden_states) query, key, value = mixed_x_layer.split(self.split_size, dim=2) @@ -303,13 +339,22 @@ class QWenAttention(nn.Module): kv_seq_len = hidden_states.size()[1] if layer_past: kv_seq_len += layer_past[0].shape[1] - rotary_pos_emb = self.rotary_emb(kv_seq_len).to(hidden_states.device) + if (self.use_dynamic_ntk and kv_seq_len == hidden_states.size()[1] + and not self.training): + context_value = math.log(kv_seq_len / self.seq_length, 2) + 1 + ntk_alpha = 2**math.ceil(context_value) - 1 + ntk_alpha = max(ntk_alpha, 1) + self._ntk_cached = ntk_alpha + else: + ntk_alpha = self._ntk_cached + rotary_pos_emb = self.rotary_emb( + kv_seq_len, ntk_alpha=ntk_alpha).to(hidden_states.device) if rotary_pos_emb is not None: if isinstance(rotary_pos_emb, tuple): rotary_pos_emb = rotary_pos_emb else: - rotary_pos_emb = ((rotary_pos_emb, ) * 2) + rotary_pos_emb = (rotary_pos_emb, ) * 2 if rotary_pos_emb is not None: q_pos_emb, k_pos_emb = rotary_pos_emb @@ -329,7 +374,16 @@ class QWenAttention(nn.Module): else: present = None - if self.use_flash_attn: + if self.use_logn_attn and not self.training: + if self.logn_tensor.device != query.device: + self.logn_tensor = self.logn_tensor.to( + query.device).type_as(query) + seq_start = key.size(1) - query.size(1) + seq_end = key.size(1) + logn_tensor = self.logn_tensor[:, seq_start:seq_end, :, :] + query = query * logn_tensor.expand_as(query) + + if self.use_flash_attn and flash_attn_unpadded_func is not None and not self.is_fp32 and query.is_cuda: q, k, v = query, key, value context_layer = self.core_attention_flash(q, k, v) @@ -347,7 +401,7 @@ class QWenAttention(nn.Module): attn_output = self.c_proj(context_layer) outputs = (attn_output, present) if output_attentions: - if self.use_flash_attn: + if self.use_flash_attn and flash_attn_unpadded_func is not None and not self.is_fp32: raise ValueError( 'Cannot output attentions while using flash-attn') else: @@ -358,9 +412,8 @@ class QWenAttention(nn.Module): class QWenMLP(nn.Module): - def __init__(self, intermediate_size, config): + def __init__(self, config): super().__init__() - self.w1 = nn.Linear( config.hidden_size, config.ffn_hidden_size // 2, @@ -369,17 +422,14 @@ class QWenMLP(nn.Module): config.hidden_size, config.ffn_hidden_size // 2, bias=not config.no_bias) - ff_dim_in = config.ffn_hidden_size // 2 self.c_proj = nn.Linear( ff_dim_in, config.hidden_size, bias=not config.no_bias) def forward(self, hidden_states): - a1 = self.w1(hidden_states) a2 = self.w2(hidden_states) intermediate_parallel = a1 * F.silu(a2) - output = self.c_proj(intermediate_parallel) return output @@ -390,17 +440,13 @@ class QWenBlock(nn.Module): super().__init__() self.num_expert = num_expert self.layer_number = layer_idx - self.apply_residual_connection_post_layernorm = config.apply_residual_connection_post_layernorm + self.apply_residual_connection_post_layernorm = ( + config.apply_residual_connection_post_layernorm) hidden_size = config.hidden_size - self.apply_residual_connection_post_layernorm = config.apply_residual_connection_post_layernorm + self.apply_residual_connection_post_layernorm = ( + config.apply_residual_connection_post_layernorm) self.bf16 = config.bf16 - if config.n_inner is not None: - inner_dim = config.n_inner - else: - ff_mult = 4 * 2 / 3 - inner_dim = ff_mult * hidden_size - self.ln_1 = RMSNorm( hidden_size, eps=config.layer_norm_epsilon, @@ -411,7 +457,7 @@ class QWenBlock(nn.Module): eps=config.layer_norm_epsilon, ) - self.mlp = QWenMLP(inner_dim, config) + self.mlp = QWenMLP(config) def forward( self, @@ -432,7 +478,8 @@ class QWenBlock(nn.Module): attention_mask=attention_mask, head_mask=head_mask, use_cache=use_cache, - output_attentions=output_attentions) + output_attentions=output_attentions, + ) attn_output = attn_outputs[0] outputs = attn_outputs[1:] @@ -492,7 +539,8 @@ class QWenPreTrainedModel(TorchModel, PreTrainedModel): p.data.normal_( mean=0.0, std=(self.config.initializer_range - / math.sqrt(2 * self.config.n_layer))) + / math.sqrt(2 * self.config.n_layer)), + ) def _set_gradient_checkpointing(self, module, value=False): if isinstance(module, QWenModel): @@ -572,12 +620,16 @@ class QWenModel(QWenPreTrainedModel): output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, ): - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_attentions = ( + output_attentions if output_attentions is not None else + self.config.output_attentions) output_hidden_states = ( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states) use_cache = use_cache if use_cache is not None else self.config.use_cache - return_dict = return_dict if return_dict is not None else self.config.use_return_dict + return_dict = ( + return_dict + if return_dict is not None else self.config.use_return_dict) if input_ids is not None and inputs_embeds is not None: raise ValueError( @@ -612,7 +664,8 @@ class QWenModel(QWenPreTrainedModel): past_length, input_shape[-1] + past_length, dtype=torch.long, - device=device) + device=device, + ) position_ids = position_ids.unsqueeze(0).view(-1, input_shape[-1]) if attention_mask is not None: @@ -620,13 +673,11 @@ class QWenModel(QWenPreTrainedModel): raise ValueError('batch_size has to be defined and > 0') attention_mask = attention_mask.view(batch_size, -1) attention_mask = attention_mask[:, None, None, :] - attention_mask = attention_mask.to(dtype=self.dtype) attention_mask = (1.0 - attention_mask) * torch.finfo( self.dtype).min encoder_attention_mask = None - head_mask = self.get_head_mask(head_mask, self.config.n_layer) if inputs_embeds is None: @@ -693,7 +744,6 @@ class QWenModel(QWenPreTrainedModel): all_self_attentions = all_self_attentions + (outputs[1], ) hidden_states = self.ln_f(hidden_states) - hidden_states = hidden_states.view(output_shape) if not return_dict: @@ -705,39 +755,57 @@ class QWenModel(QWenPreTrainedModel): last_hidden_state=hidden_states, past_key_values=presents, hidden_states=all_hidden_states, - attentions=all_self_attentions) + attentions=all_self_attentions, + ) class RotaryEmbedding(torch.nn.Module): - def __init__(self, dim, base=10000, ntk_alpha=1.0): + def __init__(self, dim, base=10000): super().__init__() - base = base * ntk_alpha**(dim / (dim - 2)) - inv_freq = 1.0 / (base**(torch.arange(0, dim, 2).float() / dim)) - self.register_buffer('inv_freq', inv_freq) + self.dim = dim + self.base = base + self.inv_freq = 1.0 / (base**(torch.arange(0, dim, 2).float() / dim)) if importlib.util.find_spec('einops') is None: raise RuntimeError('einops is required for Rotary Embedding') self._rotary_pos_emb_cache = None self._seq_len_cached = 0 + self._ntk_alpha_cached = 1.0 - def update_rotary_pos_emb_cache(self, max_seq_len, offset=0): + def update_rotary_pos_emb_cache(self, + max_seq_len, + offset=0, + ntk_alpha=1.0): seqlen = max_seq_len + offset - if seqlen > self._seq_len_cached: + if seqlen > self._seq_len_cached or ntk_alpha != self._ntk_alpha_cached: + base = self.base * ntk_alpha**(self.dim / (self.dim - 2)) + ''' + self.inv_freq = 1.0 / ( + base**(torch.arange( + 0, self.dim, 2, device=self.inv_freq.device).float() + / self.dim)) + ''' + self.inv_freq = torch.arange( + 0, self.dim, 2, device=self.inv_freq.device).float() / self.dim + self.inv_freq = 1.0 / (base**self.inv_freq) self._seq_len_cached = seqlen + self._ntk_alpha_cached = ntk_alpha seq = torch.arange(seqlen, device=self.inv_freq.device) freqs = torch.outer(seq.type_as(self.inv_freq), self.inv_freq) emb = torch.cat((freqs, freqs), dim=-1) from einops import rearrange + self._rotary_pos_emb_cache = rearrange(emb, 'n d -> 1 n 1 d') - def forward(self, max_seq_len, offset=0): - self.update_rotary_pos_emb_cache(max_seq_len, offset) + def forward(self, max_seq_len, offset=0, ntk_alpha=1.0): + self.update_rotary_pos_emb_cache(max_seq_len, offset, ntk_alpha) return self._rotary_pos_emb_cache[:, offset:offset + max_seq_len] def _rotate_half(x): from einops import rearrange + x = rearrange(x, '... (j d) -> ... j d', j=2) x1, x2 = x.unbind(dim=-2) return torch.cat((-x2, x1), dim=-1) @@ -771,7 +839,7 @@ class RMSNorm(torch.nn.Module): return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) def forward(self, x): - if rms_norm is not None: + if rms_norm is not None and x.is_cuda: return rms_norm(x, self.weight, self.eps) else: output = self._norm(x.float()).type_as(x) diff --git a/modelscope/models/nlp/qwen/configuration.py b/modelscope/models/nlp/qwen/configuration.py index bce39f12..fa1bbcbc 100644 --- a/modelscope/models/nlp/qwen/configuration.py +++ b/modelscope/models/nlp/qwen/configuration.py @@ -3,9 +3,7 @@ # This source code is licensed under the license found in the # LICENSE file in the root directory of this source tree. -from typing import List, Optional - -from transformers import GenerationConfig, PretrainedConfig +from transformers import PretrainedConfig from modelscope.utils.logger import get_logger @@ -43,7 +41,8 @@ class QWenConfig(PretrainedConfig): kv_channels=128, rotary_pct=1.0, rotary_emb_base=10000, - ntk_alpha=1.0, + use_dynamic_ntk=False, + use_logn_attn=False, use_flash_attn=True, ffn_hidden_size=22016, no_bias=True, @@ -72,7 +71,8 @@ class QWenConfig(PretrainedConfig): self.kv_channels = kv_channels self.rotary_pct = rotary_pct self.rotary_emb_base = rotary_emb_base - self.ntk_alpha = ntk_alpha + self.use_dynamic_ntk = use_dynamic_ntk + self.use_logn_attn = use_logn_attn self.use_flash_attn = use_flash_attn self.ffn_hidden_size = ffn_hidden_size self.no_bias = no_bias diff --git a/modelscope/models/nlp/qwen/qwen_generation_utils.py b/modelscope/models/nlp/qwen/qwen_generation_utils.py index 9adbc5df..2bc9df5d 100644 --- a/modelscope/models/nlp/qwen/qwen_generation_utils.py +++ b/modelscope/models/nlp/qwen/qwen_generation_utils.py @@ -84,11 +84,11 @@ def get_ltor_masks_and_position_ids( attention_mask[b, 0, (i + 1):, :(i + 1)] = 0 # Reset positions. if reset_position_ids: - position_ids[b, (i + 1):] -= (i + 1 - prev_index) + position_ids[b, (i + 1):] -= i + 1 - prev_index prev_index = i + 1 # Convert attention mask to binary: - attention_mask = (attention_mask < 0.5) + attention_mask = attention_mask < 0.5 return attention_mask, loss_mask, position_ids @@ -103,7 +103,8 @@ def get_batch(context_tokens: torch.LongTensor, eod_id: int): eod_id, reset_position_ids=False, reset_attention_mask=False, - eod_mask_loss=False) + eod_mask_loss=False, + ) return tokens, attention_mask, position_ids @@ -120,11 +121,13 @@ def get_stop_words_ids(chat_format, tokenizer): def make_context( tokenizer: PreTrainedTokenizer, query: str, - history: List[Tuple[str, str]] = [], + history: List[Tuple[str, str]] = None, system: str = '', max_window_size: int = 6144, chat_format: str = 'chatml', ): + if history is None: + history = [] if chat_format == 'chatml': im_start, im_end = '<|im_start|>', '<|im_end|>' @@ -150,10 +153,13 @@ def make_context( response_tokens = im_start_tokens + response_tokens_part + im_end_tokens next_context_tokens = nl_tokens + query_tokens + nl_tokens + response_tokens - prev_chat = f'\n{im_start}{query_text}{im_end}\n{im_start}{response_text}{im_end}' + prev_chat = ( + f'\n{im_start}{query_text}{im_end}\n{im_start}{response_text}{im_end}' + ) - current_context_size = len(system_tokens) + len( - next_context_tokens) + len(context_tokens) + current_context_size = ( + len(system_tokens) + len(next_context_tokens) + + len(context_tokens)) if current_context_size < max_window_size: context_tokens = next_context_tokens + context_tokens raw_text = prev_chat + raw_text @@ -322,9 +328,9 @@ class StopWordsLogitsProcessor(LogitsProcessor): stop_words_ids)) self.eos_token_id = eos_token_id for stop_token_seq in self.stop_words_ids: - assert len( - stop_token_seq - ) > 0, 'Stop words token sequences {} cannot have an empty list'.format( + assert ( + len(stop_token_seq) > 0 + ), 'Stop words token sequences {} cannot have an empty list'.format( stop_words_ids) def __call__(self, input_ids: torch.LongTensor, @@ -332,7 +338,7 @@ class StopWordsLogitsProcessor(LogitsProcessor): stopped_samples = self._calc_stopped_samples(input_ids) for i, should_stop in enumerate(stopped_samples): if should_stop: - scores[i, self.eos_token_id] = float(2**30) + scores[i, self.eos_token_id] = float(2**15) return scores def _tokens_match(self, prev_tokens: torch.LongTensor, @@ -365,10 +371,10 @@ class StopWordsLogitsProcessor(LogitsProcessor): def top_k_logits(logits, top_k=0, top_p=0.0, filter_value=-float('Inf')): - """ This function has been mostly taken from huggingface conversational - ai code at - https://medium.com/huggingface/how-to-build-a-state-of-the-art- - conversational-ai-with-transfer-learning-2d818ac26313 """ + """This function has been mostly taken from huggingface conversational + ai code at + https://medium.com/huggingface/how-to-build-a-state-of-the-art- + conversational-ai-with-transfer-learning-2d818ac26313""" if top_k > 0: # Remove all tokens with a probability less than the @@ -388,8 +394,8 @@ def top_k_logits(logits, top_k=0, top_p=0.0, filter_value=-float('Inf')): sorted_indices_to_remove = cumulative_probs > top_p # Shift the indices to the right to keep also the first token # above the threshold - sorted_indices_to_remove[..., 1:] \ - = sorted_indices_to_remove[..., :-1].clone() + sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[ + ..., :-1].clone() sorted_indices_to_remove[..., 0] = 0 for i in range(sorted_indices.size(0)): indices_to_remove = sorted_indices[i][sorted_indices_to_remove[i]] diff --git a/modelscope/models/nlp/qwen/text_generation.py b/modelscope/models/nlp/qwen/text_generation.py index 46abdf1e..a91819b0 100644 --- a/modelscope/models/nlp/qwen/text_generation.py +++ b/modelscope/models/nlp/qwen/text_generation.py @@ -45,6 +45,14 @@ class QWenForTextGeneration(QWenPreTrainedModel): super().__init__(config) self.transformer = QWenModel(config) self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False) + assert not (config.bf16 and config.fp16 + ), 'In config, bf16 and fp16 cannot both be true' + if config.bf16: + self.transformer.bfloat16() + self.lm_head.bfloat16() + if config.fp16: + self.transformer.half() + self.lm_head.half() self.post_init() def get_output_embeddings(self): @@ -241,6 +249,6 @@ class QWenForTextGeneration(QWenPreTrainedModel): stopping_criteria, prefix_allowed_tokens_fn, synced_gpus, - streamer, + streamer=streamer, **kwargs, ) diff --git a/modelscope/models/nlp/qwen/tokenization.py b/modelscope/models/nlp/qwen/tokenization.py index c0f66eac..0446854c 100644 --- a/modelscope/models/nlp/qwen/tokenization.py +++ b/modelscope/models/nlp/qwen/tokenization.py @@ -21,48 +21,41 @@ from modelscope.utils.logger import get_logger logger = get_logger() -TIKTOKEN_NAME = 'qwen.tiktoken' +VOCAB_FILES_NAMES = {'vocab_file': 'qwen.tiktoken'} class QWenTokenizer(PreTrainedTokenizer): """QWen tokenizer.""" """NOTE: This tokenizer will not handle special tokens to avoid injection attacks""" - @classmethod - def from_pretrained(cls, - pretrained_model_name_or_path, - cache_dir=None, - *inputs, - **kwargs): - merges_file = os.path.join(pretrained_model_name_or_path, - TIKTOKEN_NAME) - tokenizer = cls(merges_file, *inputs, **kwargs) - return tokenizer + vocab_files_names = VOCAB_FILES_NAMES - def __init__(self, - merges_file, - errors='replace', - max_len=None, - unk_token='<|endoftext|>', - bos_token='<|endoftext|>', - eos_token='<|endoftext|>', - pad_token=None, - add_prefix_space=False, - add_bos_token=False, - add_more_sp_tokens=True, - **kwargs): - bos_token = AddedToken( - bos_token, lstrip=False, rstrip=False) if isinstance( - bos_token, str) else bos_token - eos_token = AddedToken( - eos_token, lstrip=False, rstrip=False) if isinstance( - eos_token, str) else eos_token - unk_token = AddedToken( - unk_token, lstrip=False, rstrip=False) if isinstance( - unk_token, str) else unk_token - pad_token = AddedToken( - pad_token, lstrip=False, rstrip=False) if isinstance( - pad_token, str) else pad_token + def __init__( + self, + vocab_file, + errors='replace', + max_len=None, + unk_token='<|endoftext|>', + bos_token='<|endoftext|>', + eos_token='<|endoftext|>', + pad_token=None, + add_prefix_space=False, + add_bos_token=False, + add_more_sp_tokens=True, + **kwargs, + ): + bos_token = ( + AddedToken(bos_token, lstrip=False, rstrip=False) if isinstance( + bos_token, str) else bos_token) + eos_token = ( + AddedToken(eos_token, lstrip=False, rstrip=False) if isinstance( + eos_token, str) else eos_token) + unk_token = ( + AddedToken(unk_token, lstrip=False, rstrip=False) if isinstance( + unk_token, str) else unk_token) + pad_token = ( + AddedToken(pad_token, lstrip=False, rstrip=False) if isinstance( + pad_token, str) else pad_token) super().__init__( errors=errors, unk_token=unk_token, @@ -77,14 +70,21 @@ class QWenTokenizer(PreTrainedTokenizer): self.errors = errors # how to handle errors in decoding - name = 'QWen' + name = 'Qwen' ENDOFTEXT = '<|endoftext|>' IMSTART = '<|im_start|>' IMEND = '<|im_end|>' if add_more_sp_tokens: - special_tokens = (ENDOFTEXT, IMSTART, IMEND, '', '', '', - '', '') + tuple( - [f'' for i in range(200)]) + special_tokens = ( + ENDOFTEXT, + IMSTART, + IMEND, + '', + '', + '', + '', + '', + ) + tuple([f'' for i in range(200)]) else: special_tokens = (ENDOFTEXT, IMSTART, IMEND) @@ -100,7 +100,7 @@ class QWenTokenizer(PreTrainedTokenizer): for line in contents.splitlines() if line) } - mergeable_ranks = load_tiktoken_bpe(merges_file) + mergeable_ranks = load_tiktoken_bpe(vocab_file) special_tokens = { token: index for index, token in enumerate( @@ -113,9 +113,9 @@ class QWenTokenizer(PreTrainedTokenizer): mergeable_ranks=mergeable_ranks, special_tokens=special_tokens, ) - assert len(mergeable_ranks) + len( - special_tokens - ) == enc.n_vocab, f'{len(mergeable_ranks) + len(special_tokens)} != {enc.n_vocab} in encoding' + assert ( + len(mergeable_ranks) + len(special_tokens) == enc.n_vocab + ), f'{len(mergeable_ranks) + len(special_tokens)} != {enc.n_vocab} in encoding' self.mergeable_ranks = mergeable_ranks self.encoder = self.mergeable_ranks @@ -147,7 +147,7 @@ class QWenTokenizer(PreTrainedTokenizer): if len(ids) > self.max_len: logger.warning( 'Token indices sequence length is longer than the specified maximum ' - ' sequence length for this OpenAI GPT model ({} > {}). Running this' + ' sequence length for this model ({} > {}). Running this' ' sequence through the model will result in indexing errors'. format(len(ids), self.max_len)) return ids @@ -173,10 +173,6 @@ class QWenTokenizer(PreTrainedTokenizer): Args: text (`str`): The sequence to be encoded. - pair (`str`, *optional*): - A second sequence to be encoded with the first. - add_special_tokens (`bool`, *optional*, defaults to `False`): - Whether or not to add the special tokens associated with the corresponding model. kwargs (additional keyword arguments, *optional*): Will be passed to the underlying model specific encode method. See details in [`~PreTrainedTokenizerBase.__call__`] @@ -205,7 +201,34 @@ class QWenTokenizer(PreTrainedTokenizer): return self.tokenizer.n_vocab def _convert_id_to_token(self, index: int) -> str: - raise NotImplementedError + if index >= self.tokenizer.n_vocab: + return self.unk_token + return self.tokenizer.decode([index]) + + def _convert_token_to_id(self, token: str) -> int: + """Converts a token to an id using the vocab.""" + return self.encoder.get( + token.encode('UTF-8'), + self.tokenizer.encode(self.unk_token, allowed_special='all')[0], + ) + + @property + def all_special_tokens(self) -> List[str]: + """ + `List[str]`: All the special tokens (`''`, `''`, etc.) mapped to class attributes. + + Convert tokens of `tokenizers.AddedToken` type to string. + """ + all_toks = [str(s) for s in self.special_tokens.keys()] + return all_toks + + @property + def all_special_ids(self) -> List[int]: + """ + `List[int]`: List the ids of the special tokens(`''`, `''`, etc.) mapped to class attributes. + """ + all_ids = [v for v in self.special_tokens.values()] + return all_ids def _tokenize(self, text, **kwargs): """ @@ -220,9 +243,10 @@ class QWenTokenizer(PreTrainedTokenizer): self, token_ids: Union[int, List[int]], skip_special_tokens: bool = False, - clean_up_tokenization_spaces: bool = None, **kwargs, ) -> str: if isinstance(token_ids, int): token_ids = [token_ids] + if skip_special_tokens: + token_ids = [i for i in token_ids if i not in self.all_special_ids] return self.tokenizer.decode(token_ids)