()
| 127 | |
| 128 | |
| 129 | def main(): |
| 130 | parser = argparse.ArgumentParser() |
| 131 | parser = add_code_generation_args(parser) |
| 132 | args, _ = parser.parse_known_args() |
| 133 | |
| 134 | print("Loading tokenizer ...") |
| 135 | tokenizer = CodeGeeXTokenizer( |
| 136 | tokenizer_path=args.tokenizer_path, |
| 137 | mode="codegeex-13b") |
| 138 | |
| 139 | print("Loading state dict ...") |
| 140 | state_dict = paddle.load(args.load) |
| 141 | state_dict = state_dict["module"] |
| 142 | |
| 143 | print("Building CodeGeeX model ...") |
| 144 | model = model_provider(args) |
| 145 | model.set_state_dict(state_dict) |
| 146 | model.eval() |
| 147 | model.to(dtype="float16") |
| 148 | if args.quantize: |
| 149 | raise NotImplementedError("quantize") |
| 150 | |
| 151 | with open(args.prompt_file, "r") as f: |
| 152 | prompt = f.readlines() |
| 153 | prompt = "".join(prompt) |
| 154 | |
| 155 | times = {} |
| 156 | out_seq_lengths = [args.out_seq_length] |
| 157 | micro_batch_size = args.micro_batch_size |
| 158 | seq_length = args.max_position_embeddings |
| 159 | for out_seq_length in out_seq_lengths: |
| 160 | print(f"Generating with out_seq_len {out_seq_length}...") |
| 161 | |
| 162 | times[out_seq_length] = [] |
| 163 | for prompt in [prompt]: |
| 164 | t0 = time.perf_counter() |
| 165 | tokens = tokenizer.encode_code(prompt) |
| 166 | print(tokens) |
| 167 | print("Current prompt:") |
| 168 | print(prompt) |
| 169 | n_token_prompt = len(tokens) |
| 170 | print("N_token_prompt:", n_token_prompt) |
| 171 | token_stream = get_token_stream( |
| 172 | model, |
| 173 | tokenizer, |
| 174 | seq_length, |
| 175 | out_seq_length, |
| 176 | [copy.deepcopy(tokens) for _ in range(micro_batch_size)], |
| 177 | micro_batch_size=micro_batch_size, |
| 178 | topk=args.top_k, |
| 179 | topp=args.top_p, |
| 180 | temperature=args.temperature, |
| 181 | greedy=args.greedy, |
| 182 | ) |
| 183 | is_finished = [False for _ in range(micro_batch_size)] |
| 184 | for i, generated in enumerate(token_stream): |
| 185 | generated_tokens = generated[0] |
| 186 | for j in range(micro_batch_size): |
no test coverage detected