Add files using upload-large-folder tool

453a6ca verified over 1 year ago

6.64 kB

# Copyright (c) 2023 Alibaba PAI and Nvidia Megatron-LM Team.
#
# 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.

import torch

from megatron import get_args
from megatron.text_generation.communication import broadcast_int_list
from megatron.text_generation.communication import broadcast_tensor

from megatron_patch.tokenizer import get_tokenizer


def detokenize_generations(tokens_gpu_tensor, lengths_gpu_tensor,
                           return_segments):
    """
    Detokenize the generated tokens.
    Args:
        tokens_gpu_tensor (torch.Tensor): The generated tokens as a GPU tensor.
        lengths_gpu_tensor (torch.Tensor): The lengths of the generated tokens as a GPU tensor.
        return_segments (bool): Whether to return the tokenized segments or not.
    Returns:
        tuple: A tuple containing the generated tokens, the detokenized generations,
                and optionally the tokenized segments.
    """

    tokenizer = get_tokenizer()
    args = get_args()
    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.decode(sequence_tokens))
        if return_segments:
            words = []
            for token in sequence_tokens:
                if args.tokenizer_type in [
                        'SentencePieceTokenizer', 'GPTSentencePieceTokenizer'
                ]:
                    word = tokenizer.decoder[token]
                else:
                    word = tokenizer.decode(token)
                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):
    """
    Tokenize prompts and make them avaiable on all ranks.
    Args:
        prompts (list): List of prompts to be tokenized.
        tokens_to_generate (int): Number of tokens to generate.
        add_BOS (bool): Whether to add the BOS token or not.
        rank (int): Rank of the process. Only the process with this rank will tokenize the prompts.
    Returns:
        tuple: A tuple containing the tokenized prompts and their lengths.
    """

    # 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.
    if 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 = \
            _tokenize_prompts_and_batch(prompts, tokens_to_generate, add_BOS)
        # 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

    # First, broadcast the sizes.
    sizes_tensor = broadcast_int_list(2, int_list=sizes_list, rank=rank)

    # Now that we have the sizes, we can boradcast the tokens
    # and length tensors.
    sizes = sizes_tensor.tolist()
    prompts_tokens_cuda_long_tensor = broadcast_tensor(
        sizes, torch.int64, tensor=prompts_tokens_cuda_long_tensor, rank=rank)
    prompts_length_cuda_long_tensor = broadcast_tensor(
        sizes[0],
        torch.int64,
        tensor=prompts_length_cuda_long_tensor,
        rank=rank)

    return prompts_tokens_cuda_long_tensor, prompts_length_cuda_long_tensor


def _tokenize_prompts_and_batch(prompts, tokens_to_generate, add_BOS):
    """
    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.
    Args:
        prompts (list): List of prompts to be tokenized.
        tokens_to_generate (int): Number of tokens to generate.
        add_BOS (bool): Whether to add the BOS token or not.
    Returns:
        tuple: A tuple containing the tokenized prompts and their lengths.
    """

    # Tokenize all the prompts.
    tokenizer = get_tokenizer()
    if add_BOS:
        prompts_tokens = [[tokenizer.eod] + tokenizer.tokenize(prompt)
                          for prompt in prompts]
    else:
        prompts_tokens = [tokenizer.encode(prompt) 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.
    if not hasattr(tokenizer, 'eod'):
        tokenizer.eod = tokenizer.encode(tokenizer.eos_token)[0]
    for prompt_tokens, prompt_length in zip(prompts_tokens, prompts_length):
        padding_size = samples_length - prompt_length
        prompt_tokens.extend([tokenizer.eod] * padding_size)

    # Now we are in a structured format, we can convert to tensors.

    prompts_tokens_tensor = torch.cuda.LongTensor(prompts_tokens)
    prompts_length_tensor = torch.cuda.LongTensor(prompts_length)

    return prompts_tokens_tensor, prompts_length_tensor

	# Copyright (c) 2023 Alibaba PAI and Nvidia Megatron-LM Team.
	#
	# 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.

	import torch

	from megatron import get_args
	from megatron.text_generation.communication import broadcast_int_list
	from megatron.text_generation.communication import broadcast_tensor

	from megatron_patch.tokenizer import get_tokenizer


	def detokenize_generations(tokens_gpu_tensor, lengths_gpu_tensor,
	return_segments):
	"""
	Detokenize the generated tokens.

	Args:
	tokens_gpu_tensor (torch.Tensor): The generated tokens as a GPU tensor.
	lengths_gpu_tensor (torch.Tensor): The lengths of the generated tokens as a GPU tensor.
	return_segments (bool): Whether to return the tokenized segments or not.

	Returns:
	tuple: A tuple containing the generated tokens, the detokenized generations,
	and optionally the tokenized segments.

	"""

	tokenizer = get_tokenizer()
	args = get_args()
	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.decode(sequence_tokens))
	if return_segments:
	words = []
	for token in sequence_tokens:
	if args.tokenizer_type in [
	'SentencePieceTokenizer', 'GPTSentencePieceTokenizer'
	]:
	word = tokenizer.decoder[token]
	else:
	word = tokenizer.decode(token)
	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):
	"""
	Tokenize prompts and make them avaiable on all ranks.

	Args:
	prompts (list): List of prompts to be tokenized.
	tokens_to_generate (int): Number of tokens to generate.
	add_BOS (bool): Whether to add the BOS token or not.
	rank (int): Rank of the process. Only the process with this rank will tokenize the prompts.

	Returns:
	tuple: A tuple containing the tokenized prompts and their lengths.

	"""

	# 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.
	if 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 = \
	_tokenize_prompts_and_batch(prompts, tokens_to_generate, add_BOS)
	# 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

	# First, broadcast the sizes.
	sizes_tensor = broadcast_int_list(2, int_list=sizes_list, rank=rank)

	# Now that we have the sizes, we can boradcast the tokens
	# and length tensors.
	sizes = sizes_tensor.tolist()
	prompts_tokens_cuda_long_tensor = broadcast_tensor(
	sizes, torch.int64, tensor=prompts_tokens_cuda_long_tensor, rank=rank)
	prompts_length_cuda_long_tensor = broadcast_tensor(
	sizes[0],
	torch.int64,
	tensor=prompts_length_cuda_long_tensor,
	rank=rank)

	return prompts_tokens_cuda_long_tensor, prompts_length_cuda_long_tensor


	def _tokenize_prompts_and_batch(prompts, tokens_to_generate, add_BOS):
	"""
	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.

	Args:
	prompts (list): List of prompts to be tokenized.
	tokens_to_generate (int): Number of tokens to generate.
	add_BOS (bool): Whether to add the BOS token or not.

	Returns:
	tuple: A tuple containing the tokenized prompts and their lengths.

	"""

	# Tokenize all the prompts.
	tokenizer = get_tokenizer()
	if add_BOS:
	prompts_tokens = [[tokenizer.eod] + tokenizer.tokenize(prompt)
	for prompt in prompts]
	else:
	prompts_tokens = [tokenizer.encode(prompt) 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.
	if not hasattr(tokenizer, 'eod'):
	tokenizer.eod = tokenizer.encode(tokenizer.eos_token)[0]
	for prompt_tokens, prompt_length in zip(prompts_tokens, prompts_length):
	padding_size = samples_length - prompt_length
	prompt_tokens.extend([tokenizer.eod] * padding_size)

	# Now we are in a structured format, we can convert to tensors.

	prompts_tokens_tensor = torch.cuda.LongTensor(prompts_tokens)
	prompts_length_tensor = torch.cuda.LongTensor(prompts_length)

	return prompts_tokens_tensor, prompts_length_tensor