generate_samples_gpt.py 5.45 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

    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.')
zihanl's avatar
zihanl committed
75
76
77
78
    group.add_argument('--spec-toks', type=str, default=None,
                       help='additional special tokens')
    group.add_argument('--line-by-line', action="store_true",
                       help='generate samples line by line')
zihanl's avatar
zihanl committed
79
                       
zihanl's avatar
zihanl committed
80
81
82
83
84
85
86
87
    group.add_argument('--prompt', action="store_true",
                       help='generate samples based on prompting')
    group.add_argument('--prompt-file', type=str, default="",
                       help='prompting file')
    group.add_argument('--prompt-type', type=str, default="",
                       help='prompt type (context or keyphrase)')
    group.add_argument('--num-prompt-examples', type=int, default=10,
                       help='number of prompt examples')
zihanl's avatar
zihanl committed
88
89
90
91
    group.add_argument("--noknowledge", action='store_true', default=False,
                       help='Do not use knowledge in prompting')
    group.add_argument('--dynamic-prompt', action='store_true', default=False,
                       help='using different prompts for different test samples')
Mohammad's avatar
Mohammad committed
92
93
94
95
96
97
98
99

    return parser


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

    initialize_megatron(extra_args_provider=add_text_generate_args,
100
101
102
                        args_defaults={'tokenizer_type': 'GPT2BPETokenizer',
                                       'no_load_rng': True,
                                       'no_load_optim': True})
Mohammad's avatar
Mohammad committed
103

Jared Casper's avatar
Jared Casper committed
104
105
106
107
108
    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
109
110
    # Set up model and load checkpoint.
    model = get_model(model_provider)
Jared Casper's avatar
Jared Casper committed
111

Mohammad's avatar
Mohammad committed
112
113
114
    if args.load is not None:
        _ = load_checkpoint(model, None, None)

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

Mohammad's avatar
Mohammad committed
118
119
    # Generate samples.
    if args.num_samples == 0:
Mostofa Patwary's avatar
Mostofa Patwary committed
120
        if args.sample_input_file != None:
zihanl's avatar
zihanl committed
121
            args.micro_batch_size = 1
zihanl's avatar
zihanl committed
122
            generate_samples_input_from_file(model)
Mohammad's avatar
Mohammad committed
123
124
125
126
127
128
129
130
131
        else:
            generate_samples_interactive(model)
    else:
        generate_and_write_samples_unconditional(model)


if __name__ == "__main__":

    main()