()
| 226 | |
| 227 | |
| 228 | def main(): |
| 229 | parser = argparse.ArgumentParser() |
| 230 | parser = add_code_generation_args(parser) |
| 231 | args, _ = parser.parse_known_args() |
| 232 | |
| 233 | tokenizer, model = get_model(args) |
| 234 | |
| 235 | examples = [] |
| 236 | if args.example_path is None: |
| 237 | example_path = os.path.join(os.path.split(os.path.realpath(__file__))[0], "example_inputs.jsonl") |
| 238 | else: |
| 239 | example_path = args.example_path |
| 240 | |
| 241 | # Load examples for gradio DEMO |
| 242 | with open(example_path, "r", encoding="utf-8") as f: |
| 243 | for line in f: |
| 244 | examples.append(list(json.loads(line).values())) |
| 245 | |
| 246 | |
| 247 | def predict( |
| 248 | prompt, |
| 249 | lang, |
| 250 | seed, |
| 251 | out_seq_length, |
| 252 | temperature, |
| 253 | top_k, |
| 254 | top_p, |
| 255 | ): |
| 256 | set_random_seed(seed) |
| 257 | if lang != "None": |
| 258 | prompt = LANGUAGE_TAG[lang] + "\n" + prompt |
| 259 | |
| 260 | if enable_fastllm and args.fastllm: |
| 261 | model.direct_query = True |
| 262 | outputs = model.chat(tokenizer, |
| 263 | prompt, |
| 264 | max_length=out_seq_length, |
| 265 | top_p=top_p, |
| 266 | top_k=top_k, |
| 267 | temperature=temperature) |
| 268 | response = prompt + outputs[0] |
| 269 | elif enable_chatglm_cpp and args.chatglm_cpp: |
| 270 | inputs = tokenizer([prompt], return_tensors="pt") |
| 271 | pipeline = model |
| 272 | outputs = pipeline.generate(prompt, |
| 273 | max_length=inputs['input_ids'].shape[-1] + out_seq_length, |
| 274 | do_sample=temperature > 0, |
| 275 | top_p=top_p, |
| 276 | top_k=top_k, |
| 277 | temperature=temperature) |
| 278 | response = prompt + outputs |
| 279 | else: |
| 280 | inputs = tokenizer([prompt], return_tensors="pt") |
| 281 | inputs = inputs.to(model.device) |
| 282 | outputs = model.generate(**inputs, |
| 283 | max_length=inputs['input_ids'].shape[-1] + out_seq_length, |
| 284 | do_sample=True, |
| 285 | top_p=top_p, |
no test coverage detected