generate_samples_gpt.py 4.42 KB
Newer Older
Mohammad's avatar
Mohammad committed
1
# coding=utf-8
Mohammad's avatar
Mohammad committed
2
# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
Mohammad's avatar
Mohammad committed
3
4
5
6
7
8
9
10
11
12
13
14
15
#
# 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.

16
"""Sample Generate GPT"""
Mohammad's avatar
Mohammad committed
17

Mohammad's avatar
Mohammad committed
18
19
20
21
22
import os
import sys
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__),
                                             os.path.pardir)))

Neel Kant's avatar
Neel Kant committed
23
24
from megatron import get_args
from megatron import print_rank_0
Mohammad's avatar
Mohammad committed
25
from megatron import get_tokenizer
26
from megatron import mpu
Mohammad's avatar
Mohammad committed
27
28
from megatron.checkpointing import load_checkpoint
from megatron.initialize import initialize_megatron
29
from megatron.model import GPTModel
Mohammad's avatar
Mohammad committed
30
31
32
from megatron.training import get_model
from megatron.text_generation_utils import generate_and_write_samples_unconditional
from megatron.text_generation_utils import generate_samples_input_from_file
zihanl's avatar
zihanl committed
33
34
from megatron.text_generation_utils import generate_samples_prompt_input_from_file
from megatron.text_generation_utils import generate_samples_line_by_line_input_from_file
Mohammad's avatar
Mohammad committed
35
36
37
from megatron.text_generation_utils import generate_samples_interactive


38
def model_provider(pre_process=True, post_process=True):
Mohammad's avatar
Mohammad committed
39
40
    """Build the model."""

41
    print_rank_0('building GPT model ...')
42
43
    model = GPTModel(num_tokentypes=0, parallel_output=False,
                     pre_process=pre_process, post_process=post_process)
Mohammad's avatar
Mohammad committed
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82

    return model


def add_text_generate_args(parser):
    """Text generation arguments."""
    group = parser.add_argument_group(title='text generation')

    group.add_argument("--temperature", type=float, default=1.0,
                       help='Sampling temperature.')
    group.add_argument("--greedy", action='store_true', default=False,
                       help='Use greedy sampling.')
    group.add_argument("--top_p", type=float, default=0.0,
                       help='Top p sampling.')
    group.add_argument("--top_k", type=int, default=0,
                       help='Top k sampling.')
    group.add_argument("--out-seq-length", type=int, default=1024,
                       help='Size of the output generated text.')
    group.add_argument("--sample-input-file", type=str, default=None,
                       help='Get input from file instead of interactive mode, '
                       'each line is an input.')
    group.add_argument("--sample-output-file", type=str, default=None,
                       help='Output file got from --sample-input-file')
    group.add_argument("--num-samples", type=int, default=0,
                       help='Number of samples to generate unconditionally, '
                       'defaults to 0 and interactive conditional sampling')
    group.add_argument("--genfile", type=str,
                       help='Output file when generating unconditionally')
    group.add_argument("--recompute", action='store_true',
                       help='During generation recompute all attention '
                       'instead of using previously computed keys/values.')

    return parser


def main():
    """Main program."""

    initialize_megatron(extra_args_provider=add_text_generate_args,
83
84
85
                        args_defaults={'tokenizer_type': 'GPT2BPETokenizer',
                                       'no_load_rng': True,
                                       'no_load_optim': True})
Mohammad's avatar
Mohammad committed
86

Jared Casper's avatar
Jared Casper committed
87
88
89
90
91
    args = get_args()
    if args.num_layers_per_virtual_pipeline_stage is not None:
        print("Interleaved pipeline schedule is not yet supported for text generation.")
        exit()

Mohammad's avatar
Mohammad committed
92
93
    # Set up model and load checkpoint.
    model = get_model(model_provider)
Jared Casper's avatar
Jared Casper committed
94

Mohammad's avatar
Mohammad committed
95
96
97
    if args.load is not None:
        _ = load_checkpoint(model, None, None)

Jared Casper's avatar
Jared Casper committed
98
99
100
    assert len(model) == 1, "Above condition should have caught this"
    model = model[0]

Mohammad's avatar
Mohammad committed
101
102
    # Generate samples.
    if args.num_samples == 0:
Mostofa Patwary's avatar
Mostofa Patwary committed
103
        if args.sample_input_file != None:
zihanl's avatar
zihanl committed
104
            args.micro_batch_size = 1
zihanl's avatar
zihanl committed
105
            generate_samples_input_from_file(model)
Mohammad's avatar
Mohammad committed
106
107
108
109
110
111
112
113
114
        else:
            generate_samples_interactive(model)
    else:
        generate_and_write_samples_unconditional(model)


if __name__ == "__main__":

    main()