# coding=utf-8 # Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Tokenization utilities.""" import re import torch from icecream import ic def detokenize_generations(tokens_gpu_tensor, lengths_gpu_tensor, return_segments, tokenizer): """Detokenize the generated tokens.""" prompts_plus_generations = [] if return_segments: prompts_plus_generations_segments = [] tokens = tokens_gpu_tensor.cpu().numpy().tolist() lengths = lengths_gpu_tensor.cpu().numpy().tolist() for sequence_tokens, length in zip(tokens, lengths): sequence_tokens = sequence_tokens[:length] prompts_plus_generations.append( tokenizer.detokenize(sequence_tokens)) if return_segments: from tokenizers.decoders import Metaspace if hasattr(tokenizer, 'tokenizer'): if isinstance(tokenizer.tokenizer.decoder, Metaspace): words = tokenizer.tokenizer.decode(sequence_tokens) else: words = [] for token in sequence_tokens: word = tokenizer.tokenizer.decoder[token] word = bytearray( [tokenizer.tokenizer.byte_decoder[c] for c in word]).decode( 'utf-8', errors='replace') words.append(word) prompts_plus_generations_segments.append(words) else: words = tokenizer.detokenize(sequence_tokens) # else: # words = [] # for token in sequence_tokens: # word = tokenizer.tokenizer.decoder[token] # word = bytearray( # [tokenizer.tokenizer.byte_decoder[c] for c in word]).decode( # 'utf-8', errors='replace') # words.append(word) prompts_plus_generations_segments.append(words) if return_segments: return tokens, prompts_plus_generations, \ prompts_plus_generations_segments return tokens, prompts_plus_generations def tokenize_prompts(prompts=None, tokens_to_generate=None, add_BOS=None, rank=0, tokenizer=None, ignore_dist=False): """Tokenize prompts and make them avaiable on all ranks.""" # On all ranks set to None so we can pass them to functions sizes_list = None prompts_tokens_cuda_long_tensor = None prompts_length_cuda_long_tensor = None # On the specified rank, build the above. attention_mask = None if ignore_dist or torch.distributed.get_rank() == rank: assert prompts is not None assert tokens_to_generate is not None # Tensor of tokens padded and their unpadded length. prompts_tokens_cuda_long_tensor, prompts_length_cuda_long_tensor, attention_mask = \ _tokenize_prompts_and_batch( prompts, tokens_to_generate, add_BOS, tokenizer) # We need the sizes of these tensors for the boradcast sizes_list = [prompts_tokens_cuda_long_tensor.size(0), # Batch size prompts_tokens_cuda_long_tensor.size(1)] # Sequence lenght return prompts_tokens_cuda_long_tensor, prompts_length_cuda_long_tensor, attention_mask def _tokenize_prompts_and_batch(prompts, tokens_to_generate, add_BOS, tokenizer): """Given a set of prompts and number of tokens to generate: - tokenize prompts - set the sequence length to be the max of length of prompts plus the number of tokens we would like to generate - pad all the sequences to this length so we can convert them into a 2D tensor. """ # Tokenize all the prompts. # if add_BOS: # prompts_tokens = [[tokenizer.bos] + tokenizer.tokenize(prompt) # for prompt in prompts] # else: # prompts_tokens = [tokenizer.tokenize(prompt) for prompt in prompts] prompts_tokens = [_tokenize_prompt( prompt, tokenizer, add_BOS) for prompt in prompts] # Now we have a list of list of tokens which each list has a different # size. We want to extend this list to: # - incorporate the tokens that need to be generated # - make all the sequences equal length. # Get the prompts length. prompts_length = [len(prompt_tokens) for prompt_tokens in prompts_tokens] # Get the max prompts length. max_prompt_len = max(prompts_length) # Number of tokens in the each sample of the batch. samples_length = max_prompt_len + tokens_to_generate # Now update the list of list to be of the same size: samples_length. for prompt_tokens, prompt_length in zip(prompts_tokens, prompts_length): padding_size = samples_length - prompt_length prompt_tokens.extend([tokenizer.eos_token_id] * padding_size) # Now we are in a structured format, we can convert to tensors. prompts_tokens_tensor = torch.LongTensor(prompts_tokens) prompts_length_tensor = torch.LongTensor(prompts_length) attention_mask = torch.zeros(prompts_tokens_tensor.shape[:2]) for i, l in enumerate(prompts_length_tensor): attention_mask[i, :l] = 1 return prompts_tokens_tensor, prompts_length_tensor, attention_mask def _tokenize_prompt(prompt, tokenizer, add_BOS=False, media_info={'': 65}): media_tokens = {k: -int(i+1) for i, k in enumerate(media_info.keys())} media_lengths = media_info.copy() if add_BOS: prompt_chunk = [tokenizer.bos_token_id] else: prompt_chunk = [] # Pure Text if all([media_token not in prompt for media_token in media_tokens.keys()]): enc_chunk = prompt_chunk + \ tokenizer(prompt, add_special_tokens=False)['input_ids'] # Multi-Modal Text else: enc_chunk = prompt_chunk pattern = '|'.join(map(re.escape, list(media_tokens.keys()))) chunk_strs = re.split(f'({pattern})', prompt) chunk_strs = [x for x in chunk_strs if len(x) > 0] for idx, chunk_str in enumerate(chunk_strs): if chunk_str in media_tokens: enc_chunk += [media_tokens[chunk_str]] * \ media_lengths[chunk_str] else: tmp_chunk = tokenizer(chunk_str, add_special_tokens=False)[ 'input_ids'] # if idx < len(chunk_strs) - 1: # Last chunk should not have eos # tmp_chunk += [tokenizer.eod_id] enc_chunk += tmp_chunk return enc_chunk