()
| 116 | |
| 117 | |
| 118 | def main(): |
| 119 | parser = argparse.ArgumentParser() |
| 120 | parser = add_code_generation_args(parser) |
| 121 | args, _ = parser.parse_known_args() |
| 122 | |
| 123 | print("Loading tokenizer ...") |
| 124 | tokenizer = CodeGeeXTokenizer( |
| 125 | tokenizer_path=args.tokenizer_path, |
| 126 | mode="codegeex-13b") |
| 127 | |
| 128 | print("Loading state dict ...") |
| 129 | state_dict = torch.load(args.load, map_location="cpu") |
| 130 | state_dict = state_dict["module"] |
| 131 | |
| 132 | print("Building CodeGeeX model ...") |
| 133 | model = model_provider(args) |
| 134 | model.load_state_dict(state_dict) |
| 135 | model.eval() |
| 136 | model.half() |
| 137 | if args.quantize: |
| 138 | model = quantize(model, weight_bit_width=8, backend="torch") |
| 139 | model.cuda() |
| 140 | torch.cuda.synchronize() |
| 141 | |
| 142 | with open(args.prompt_file, "r") as f: |
| 143 | prompt = f.readlines() |
| 144 | prompt = "".join(prompt) |
| 145 | |
| 146 | out_seq_lengths = [args.out_seq_length] |
| 147 | for out_seq_length in out_seq_lengths: |
| 148 | print(f"Generating with out_seq_len {out_seq_length}...") |
| 149 | while True: |
| 150 | print("\nPlease Input Query (Ctrl-D to save multiple lines, 'stop' to exit) >>> ") |
| 151 | prompts = [] |
| 152 | while True: |
| 153 | try: |
| 154 | line = input() |
| 155 | except EOFError: |
| 156 | break |
| 157 | prompts.append(line) |
| 158 | prompt = "\n".join(prompts) |
| 159 | prompt = prompt.strip() |
| 160 | if not prompt: |
| 161 | print('Query should not be empty!') |
| 162 | continue |
| 163 | if prompt == "stop": |
| 164 | return |
| 165 | try: |
| 166 | t0 = time.perf_counter() |
| 167 | generated_code = codegeex.generate( |
| 168 | model, |
| 169 | tokenizer, |
| 170 | prompt, |
| 171 | out_seq_length=out_seq_length, |
| 172 | seq_length=args.max_position_embeddings, |
| 173 | top_k=args.top_k, |
| 174 | top_p=args.top_p, |
| 175 | temperature=args.temperature, |
no test coverage detected