Function generate_samples_interactive

codegeex/megatron/code_generation_utils.py:377–484

(model, print_frequency=24)

def generate_samples_interactive(model, print_frequency=24):
    args = get_args()
    tokenizer = get_tokenizer()

    context_count = 0
    model.eval()
    with torch.no_grad():
        while True:
            terminate_runs = 0
            raw_text_len = 0

            if (
                    mpu.is_pipeline_first_stage()
                    and mpu.get_tensor_model_parallel_rank() == 0
            ):
                os.system("clear")
                raw_text = input("\nContext prompt (stop to exit) >>> ")
                while not raw_text:
                    print("Prompt should not be empty!")
                    raw_text = input("\nContext prompt (stop to exit) >>> ")
                raw_text_len = len(raw_text)

                if "stop" in raw_text:
                    terminate_runs = 1
                else:
                    context_tokens = tokenizer.tokenize(raw_text)
                    context_length = len(context_tokens)

                    if context_length >= (args.seq_length // 2):
                        print(
                            "\nContext length",
                            context_length,
                            "\nPlease give smaller context (half of the "
                            "sequence length)!",
                            flush=True,
                        )
                        continue
            else:
                context_tokens = tokenizer.tokenize("EMPTY TEXT")
                context_length = 0

            input_info = [terminate_runs, raw_text_len, context_length]
            input_info_tensor = torch.cuda.LongTensor(input_info)
            torch.distributed.all_reduce(
                input_info_tensor, group=mpu.get_model_parallel_group()
            )
            terminate_runs = input_info_tensor[0].item()
            raw_text_len = input_info_tensor[1].item()
            context_length = input_info_tensor[2].item()

            if terminate_runs == 1:
                return

            # For pipeline parallel we send context tokens to other stages
            # so they get the lengths correct
            if (
                    mpu.get_tensor_model_parallel_rank() == 0
                    and args.pipeline_model_parallel_size > 1