()
| 117 | |
| 118 | |
| 119 | def main(): |
| 120 | parser = argparse.ArgumentParser() |
| 121 | parser = add_code_generation_args(parser) |
| 122 | args, _ = parser.parse_known_args() |
| 123 | |
| 124 | print("Loading tokenizer ...") |
| 125 | tokenizer = CodeGeeXTokenizer( |
| 126 | tokenizer_path=args.tokenizer_path, |
| 127 | mode="codegeex-13b") |
| 128 | |
| 129 | print("Loading state dict ...") |
| 130 | state_dict = torch.load(args.load, map_location="cpu") |
| 131 | state_dict = state_dict["module"] |
| 132 | |
| 133 | print("Building CodeGeeX model ...") |
| 134 | model = model_provider(args) |
| 135 | model.load_state_dict(state_dict) |
| 136 | model.eval() |
| 137 | model.half() |
| 138 | if args.quantize: |
| 139 | model = quantize_oneflow(model, weight_bit_width=8) |
| 140 | model.cuda() |
| 141 | torch.cuda.synchronize() |
| 142 | with open(args.prompt_file, "r") as f: |
| 143 | prompt = f.readlines() |
| 144 | prompt = "".join(prompt) |
| 145 | |
| 146 | times = {} |
| 147 | out_seq_lengths = [args.out_seq_length] |
| 148 | micro_batch_size = args.micro_batch_size |
| 149 | seq_length = args.max_position_embeddings |
| 150 | for out_seq_length in out_seq_lengths: |
| 151 | print(f"Generating with out_seq_len {out_seq_length}...") |
| 152 | |
| 153 | times[out_seq_length] = [] |
| 154 | for prompt in [prompt]: |
| 155 | t0 = time.perf_counter() |
| 156 | tokens = tokenizer.encode_code(prompt) |
| 157 | print(tokens) |
| 158 | print("Current prompt:") |
| 159 | print(prompt) |
| 160 | n_token_prompt = len(tokens) |
| 161 | print("N_token_prompt:", n_token_prompt) |
| 162 | token_stream = get_token_stream( |
| 163 | model, |
| 164 | tokenizer, |
| 165 | seq_length, |
| 166 | out_seq_length, |
| 167 | [copy.deepcopy(tokens) for _ in range(micro_batch_size)], |
| 168 | micro_batch_size=micro_batch_size, |
| 169 | topk=args.top_k, |
| 170 | topp=args.top_p, |
| 171 | temperature=args.temperature, |
| 172 | greedy=args.greedy, |
| 173 | ) |
| 174 | is_finished = [False for _ in range(micro_batch_size)] |
| 175 | for i, generated in enumerate(token_stream): |
| 176 | generated_tokens = generated[0] |
no test coverage detected