File size: 3,450 Bytes
d90b3a8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 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 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 |
#!/usr/bin/env python
# Copyright (c) 2024 EleutherAI
# This file is based on code by the authors denoted below and has been modified from its original version.
#
# Copyright (c) 2024, 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.
from megatron.utils import print_rank_0, setup_for_inference_or_eval
from megatron.text_generation_utils import (
generate_samples_input_from_file,
generate_samples_from_prompt,
generate_samples_unconditional,
generate_samples_interactive,
precompute_logits,
)
def main(input_args=None, overwrite_values=None):
"""
Generate text/sample model
"""
model, neox_args = setup_for_inference_or_eval(
use_cache=True, input_args=input_args, overwrite_values=overwrite_values
)
if neox_args.recompute:
model.module.inference_mode(
use_cache=False
) # don't use kv cache if recomputing
if neox_args.text_gen_type == "unconditional":
print_rank_0(
f"Generating samples unconditionally and saving results to {neox_args.sample_output_file}"
)
generate_samples_unconditional(
neox_args=neox_args,
model=model,
number_of_samples=neox_args.num_samples,
output_file=neox_args.sample_output_file,
maximum_tokens=neox_args.maximum_tokens,
recompute=neox_args.recompute,
temperature=neox_args.temperature,
top_k=neox_args.top_k,
top_p=neox_args.top_p,
)
elif neox_args.text_gen_type == "input-file":
print_rank_0(
f"Generating samples from input file {neox_args.sample_input_file}"
)
assert neox_args.sample_input_file is not None
generate_samples_input_from_file(
neox_args=neox_args,
model=model,
input_file=neox_args.sample_input_file,
output_file=neox_args.sample_output_file,
maximum_tokens=neox_args.maximum_tokens,
prompt_end=neox_args.prompt_end,
recompute=neox_args.recompute,
temperature=neox_args.temperature,
top_k=neox_args.top_k,
top_p=neox_args.top_p,
)
elif neox_args.text_gen_type == "interactive":
generate_samples_interactive(
neox_args=neox_args,
model=model,
recompute=neox_args.recompute,
temperature=neox_args.temperature,
maximum_tokens=neox_args.maximum_tokens,
prompt_end=neox_args.prompt_end,
top_k=neox_args.top_k,
top_p=neox_args.top_p,
)
elif neox_args.text_gen_type == "precompute":
precompute_logits(neox_args=neox_args, model=model)
else:
raise ValueError(
f"`text_gen_type` either not specified or not recognised: {neox_args.text_gen_type}"
)
if __name__ == "__main__":
main()
|