(model)
| 33 | |
| 34 | |
| 35 | def run_generation_distributed(model): |
| 36 | args = get_args() |
| 37 | if hasattr(args, "language_tgt_type"): |
| 38 | language_type = args.language_tgt_type |
| 39 | else: |
| 40 | language_type = args.language_type |
| 41 | print(f"Connecting to tcp://{args.channel_ip}:{args.channel_port}") |
| 42 | context = zmq.Context() |
| 43 | socket = context.socket(zmq.REQ) |
| 44 | socket.connect(f"tcp://{args.channel_ip}:{args.channel_port}") |
| 45 | output_file_path = args.output_prefix + f"_finished_rank{args.gen_rank}.jsonl" |
| 46 | unfinished_output_file_path = args.output_prefix + f"_unfinished_rank{args.gen_rank}.jsonl" |
| 47 | problems = {} |
| 48 | print("Building tokenizer...") |
| 49 | tokenizer = get_tokenizer() |
| 50 | |
| 51 | with open(output_file_path, "w") as f: |
| 52 | with open(unfinished_output_file_path, "w") as unfinished_f: |
| 53 | while True: |
| 54 | socket.send_json({"rank": args.gen_rank, "action": "pull"}) |
| 55 | resp = socket.recv_json() |
| 56 | try: |
| 57 | if "codecontest" in args.dataset.lower(): |
| 58 | if resp["contest_name"] is None: |
| 59 | break |
| 60 | elif resp["task_id"] is None: |
| 61 | break |
| 62 | |
| 63 | if "codecontest" in args.dataset.lower(): |
| 64 | current_spec = problems[resp["contest_name"]] |
| 65 | prompt = current_spec.prompt |
| 66 | else: |
| 67 | current_spec = resp["task_id"] |
| 68 | prompt = current_spec["prompt"] |
| 69 | |
| 70 | temperature = None if "temperature" not in resp else resp["temperature"] |
| 71 | topp = None if "topp" not in resp else resp["topp"] |
| 72 | |
| 73 | f.flush() |
| 74 | unfinished_f.flush() |
| 75 | tokens = tokenizer.tokenize(prompt) |
| 76 | n_token_prompt = len(tokens) |
| 77 | if n_token_prompt >= args.seq_length: |
| 78 | continue |
| 79 | if "micro_batch_size" in resp: |
| 80 | micro_batch_size = resp["micro_batch_size"] |
| 81 | else: |
| 82 | micro_batch_size = args.micro_batch_size |
| 83 | if args.beam_search: |
| 84 | beams = get_token_stream( |
| 85 | model, |
| 86 | [ |
| 87 | copy.deepcopy(tokens) |
| 88 | for _ in range(micro_batch_size) |
| 89 | ], |
| 90 | return_scores=args.return_scores, |
| 91 | prompt_length=n_token_prompt, |
| 92 | micro_batch_size=micro_batch_size, |
no test coverage detected