172 lines
7.0 KiB
Python
172 lines
7.0 KiB
Python
# 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={'<image>': 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
|