api.py 8.45 KB
Newer Older
mshoeybi's avatar
working  
mshoeybi committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
# 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.

"""Inference API."""


import torch

mshoeybi's avatar
mshoeybi committed
21
from megatron import mpu
mshoeybi's avatar
working  
mshoeybi committed
22
from .communication import broadcast_float_list
23
24
from .generation import (
        generate_tokens_probs_and_return_on_first_stage,
rprenger's avatar
rprenger committed
25
26
        score_and_return_on_first_stage,
        beam_search_and_return_on_first_stage)
mshoeybi's avatar
mshoeybi committed
27
28
29
30
31
32
33
34
from .tokenization import (
    tokenize_prompts,
    detokenize_generations)

def generate_and_post_process(model,
                              prompts=None,
                              tokens_to_generate=0,
                              return_output_log_probs=False,
mshoeybi's avatar
mshoeybi committed
35
36
                              top_k_sampling=0,
                              top_p_sampling=0.0,
37
38
                              top_p_decay=0.0,
                              top_p_bound=0.0,
mshoeybi's avatar
mshoeybi committed
39
                              temperature=1.0,
mshoeybi's avatar
mshoeybi committed
40
                              add_BOS=False,
41
42
                              use_eod_token_for_early_termination=True,
                              stop_on_double_eol=False,
43
44
                              stop_on_eol=False,
                              random_seed=-1):
mshoeybi's avatar
mshoeybi committed
45
    """Run inference and post-process outputs, i.e., detokenize,
mshoeybi's avatar
mshoeybi committed
46
    move to cpu and convert to list."""
mshoeybi's avatar
mshoeybi committed
47
48

    # Main inference.
49
    tokens, lengths, output_log_probs = generate(
mshoeybi's avatar
mshoeybi committed
50
51
52
53
        model,
        prompts=prompts,
        tokens_to_generate=tokens_to_generate,
        return_output_log_probs=return_output_log_probs,
mshoeybi's avatar
mshoeybi committed
54
55
        top_k_sampling=top_k_sampling,
        top_p_sampling=top_p_sampling,
56
57
        top_p_decay=top_p_decay,
        top_p_bound=top_p_bound,
mshoeybi's avatar
mshoeybi committed
58
        temperature=temperature,
mshoeybi's avatar
mshoeybi committed
59
        add_BOS=add_BOS,
60
61
        use_eod_token_for_early_termination=use_eod_token_for_early_termination,
        stop_on_double_eol=stop_on_double_eol,
62
63
        stop_on_eol=stop_on_eol,
        random_seed=random_seed)
mshoeybi's avatar
mshoeybi committed
64
65
66
67
68
69
70
71

    # Only post-process on first stage.
    if mpu.is_pipeline_first_stage():
        tokens, prompts_plus_generations, prompts_plus_generations_segments = \
            detokenize_generations(tokens, lengths, True)

        if return_output_log_probs:
            output_log_probs = output_log_probs.cpu().numpy().tolist()
72
73
            for i, (prob, seg) in enumerate(zip(output_log_probs, prompts_plus_generations_segments)):
                output_log_probs[i] = prob[:len(seg)-1]
mshoeybi's avatar
mshoeybi committed
74
75

        return prompts_plus_generations, prompts_plus_generations_segments, \
76
            output_log_probs, tokens
mshoeybi's avatar
mshoeybi committed
77
78

    return None
mshoeybi's avatar
working  
mshoeybi committed
79
80
81
82
83

def generate(model,
             prompts=None,
             tokens_to_generate=0,
             return_output_log_probs=False,
mshoeybi's avatar
mshoeybi committed
84
85
             top_k_sampling=0,
             top_p_sampling=0.0,
86
87
             top_p_decay=0.0,
             top_p_bound=0.0,
mshoeybi's avatar
mshoeybi committed
88
             temperature=1.0,
mshoeybi's avatar
mshoeybi committed
89
             add_BOS=False,
90
91
             use_eod_token_for_early_termination=True,
             stop_on_double_eol=False,
92
93
             stop_on_eol=False,
             random_seed=-1):
mshoeybi's avatar
mshoeybi committed
94
95
96
97
98
99
100
    """Given prompts and input parameters, run inference and return:
       tokens: prompts plus the generated tokens.
       lengths: length of the prompt + generations. Note that we can
           discard tokens in the tokens tensor that are after the
           corresponding length.
       output_log_probs: log probs of the tokens.
    """
mshoeybi's avatar
working  
mshoeybi committed
101
102

    # Make sure input params are avaialble to all ranks.
mshoeybi's avatar
mshoeybi committed
103
    values = [tokens_to_generate,
104
              return_output_log_probs,
105
              top_k_sampling, top_p_sampling, top_p_decay, top_p_bound,
106
107
              temperature, add_BOS, use_eod_token_for_early_termination,
              stop_on_double_eol,
108
109
              stop_on_eol,
              random_seed]
110
    values_float_tensor = broadcast_float_list(12, float_list=values)
mshoeybi's avatar
working  
mshoeybi committed
111
112
    tokens_to_generate = int(values_float_tensor[0].item())
    return_output_log_probs = bool(values_float_tensor[1].item())
mshoeybi's avatar
mshoeybi committed
113
114
    top_k_sampling = int(values_float_tensor[2].item())
    top_p_sampling = values_float_tensor[3].item()
115
116
    top_p_decay = values_float_tensor[4].item()
    top_p_bound = values_float_tensor[5].item()
117
118
119
120
121
122
    temperature = values_float_tensor[6].item()
    add_BOS = bool(values_float_tensor[7].item())
    use_eod_token_for_early_termination = bool(values_float_tensor[8].item())
    stop_on_double_eol = bool(values_float_tensor[9].item())
    stop_on_eol = bool(values_float_tensor[10].item())
    random_seed = int(values_float_tensor[11].item())
123
124
125

    if random_seed != -1:
        torch.random.manual_seed(random_seed)
mshoeybi's avatar
working  
mshoeybi committed
126
127
128
129
130

    # Tokenize prompts and get the batch.
    # Note that these tensors are broadcaseted to all ranks.
    if torch.distributed.get_rank() == 0:
        assert prompts is not None
rprenger's avatar
rprenger committed
131
    
mshoeybi's avatar
working  
mshoeybi committed
132
    context_tokens_tensor, context_length_tensor = tokenize_prompts(
mshoeybi's avatar
mshoeybi committed
133
        prompts=prompts, tokens_to_generate=tokens_to_generate, add_BOS=add_BOS)
mshoeybi's avatar
working  
mshoeybi committed
134

135
    if tokens_to_generate == 0:
136
137
        return score_and_return_on_first_stage(
            model, context_tokens_tensor, context_length_tensor)
138
    
mshoeybi's avatar
working  
mshoeybi committed
139
140
141
142
143
    # Main inference function.
    # Note that the outputs are available on the first stage.
    return generate_tokens_probs_and_return_on_first_stage(
        model, context_tokens_tensor, context_length_tensor,
        return_output_log_probs=return_output_log_probs,
mshoeybi's avatar
mshoeybi committed
144
145
        top_k=top_k_sampling,
        top_p=top_p_sampling,
146
147
        top_p_decay=top_p_decay,
        top_p_bound=top_p_bound,
mshoeybi's avatar
mshoeybi committed
148
        temperature=temperature,
149
150
151
        use_eod_token_for_early_termination=use_eod_token_for_early_termination,
        stop_on_double_eol=stop_on_double_eol,
        stop_on_eol=stop_on_eol)
rprenger's avatar
rprenger committed
152
153
154
155
156

def beam_search_and_post_process(model,
                                 prompts=None,
                                 tokens_to_generate=0,
                                 beam_size=0,
157
158
                                 add_BOS=False,
                                 stop_token=50256,
159
160
                                 num_return_gen=1,
                                 length_penalty=1):
rprenger's avatar
rprenger committed
161
162
163
164
165
166
167
168
    """Run beam search and post-process outputs, i.e., detokenize,
    move to cpu and convert to list."""

    # Main inference.
    tokens, scores = beam_search(model,
                                 prompts=prompts,
                                 tokens_to_generate=tokens_to_generate,
                                 beam_size=beam_size,
169
170
                                 add_BOS=add_BOS,
                                 stop_token=stop_token,
171
172
                                 num_return_gen=num_return_gen,
                                 length_penalty=length_penalty)
rprenger's avatar
rprenger committed
173
174
175
176
    # Only post-process on first stage.
    if mpu.is_pipeline_first_stage():
        lengths = tokens.size(1)*torch.ones(beam_size, dtype=torch.int64, device=torch.cuda.current_device()) 
        tokens, prompts_plus_generations, prompts_plus_generations_segments = detokenize_generations(tokens, lengths, True)
rprenger's avatar
rprenger committed
177
178
        scores = scores.cpu().numpy().tolist()
        return prompts_plus_generations, prompts_plus_generations_segments, scores
rprenger's avatar
rprenger committed
179
180
181

    return None

182
def beam_search(model, prompts=None, tokens_to_generate=0, beam_size=0, add_BOS=False, stop_token=50256, num_return_gen=1, length_penalty=1):
rprenger's avatar
rprenger committed
183
184
185
    # Make sure input params are avaialble to all ranks.
    values = [tokens_to_generate,
              beam_size,
186
187
188
189
              add_BOS,
              stop_token,
              num_return_gen,
              length_penalty]
190
    values_float_tensor = broadcast_float_list(6, float_list=values)
rprenger's avatar
rprenger committed
191
192
193
    tokens_to_generate = int(values_float_tensor[0].item())
    beam_size = int(values_float_tensor[1].item())
    add_BOS = bool(values_float_tensor[2].item())
194
195
196
    stop_token = int(values_float_tensor[3].item())
    num_return_gen = int(values_float_tensor[4].item())
    length_penalty = values_float_tensor[5].item()
rprenger's avatar
rprenger committed
197
198
199
200

    context_tokens_tensor, context_length_tensor = tokenize_prompts(
        prompts=prompts, tokens_to_generate=tokens_to_generate, add_BOS=add_BOS)
    
201
    return beam_search_and_return_on_first_stage(model, context_tokens_tensor, context_length_tensor, 
202
            beam_size, stop_token=stop_token, num_return_gen=num_return_gen, length_penalty=length_penalty)