| 167 | env.write_file(f"{batch['batch_id']}", batch) |
| 168 | |
| 169 | def main(): |
| 170 | parser = argparse.ArgumentParser() |
| 171 | parser.add_argument("--output_dir", type=str, required=True) |
| 172 | parser.add_argument("--model_name", type=str, required=True) |
| 173 | parser.add_argument("--base_url", type=str, required=True) |
| 174 | parser.add_argument("--api_key", type=str, required=True) |
| 175 | parser.add_argument("--max_turns", type=int, default=100) |
| 176 | parser.add_argument("--persona", type=str, default='general') |
| 177 | parser.add_argument("--hp", type=int, default=5) |
| 178 | parser.add_argument("--remove", action="store_true", default=False, help="Enable removal option") |
| 179 | parser.add_argument("--input_file", type=str) |
| 180 | args = parser.parse_args() |
| 181 | |
| 182 | os.makedirs(args.output_dir, exist_ok=True) |
| 183 | |
| 184 | data = load_data(args.input_file) |
| 185 | cache = load_data(args.output_dir) |
| 186 | cache = [line['batch_id'] for line in cache] |
| 187 | data = [line for line in data if line['batch_id'] not in cache] |
| 188 | |
| 189 | print(f'reload cache with {len(cache)} items. Running the remaining {len(data)} items.') |
| 190 | |
| 191 | _run(args, data) |
| 192 | exit(0) |
| 193 | num_processes = min(5, len(data)) |
| 194 | |
| 195 | if len(data) < 5: |
| 196 | _run(args, data) |
| 197 | exit() |
| 198 | |
| 199 | length = len(data) // num_processes + 1 |
| 200 | pool = multiprocessing.Pool(processes=num_processes) |
| 201 | |
| 202 | collects = [] |
| 203 | for ids in range(num_processes): |
| 204 | collect = data[ids * length:(ids + 1) * length] |
| 205 | collects.append(pool.apply_async(_run, (args, collect))) |
| 206 | pool.close() |
| 207 | pool.join() |
| 208 | |
| 209 | print('All done.') |
| 210 | |
| 211 | if __name__ == "__main__": |
| 212 | main() |