Function main

WebAgent/ParallelMuse/compressed_reasoning_aggregation.py:252–291

()

async def main():
    data_path = "[YOUR-ROLLOUT-FILE-PATH-HERE]"  # TODO

    report_sem = asyncio.Semaphore(64)
    merge_sem = asyncio.Semaphore(32)
    sem = {
        'report': report_sem,
        'merge': merge_sem
    }
    mode = 'converge_info'

    dataset = read_jsonl(data_path)

    tasks = []
    clustered_dataset = cluster_by_question(dataset)
    for cluster in clustered_dataset:
        filtered_cluster = []
        for traj in cluster:
            if 'prediction' in traj.keys() and traj['prediction'] != '[No Prediction]':
                filtered_cluster.append(traj)
    
        tasks.append(call_converge(sem, filtered_cluster, data_path))

    results = []

    with open(f"{data_path.replace('.jsonl', f'_{mode}.jsonl')}", "a") as f:
        for future in tqdm(asyncio.as_completed(tasks), total=len(tasks), desc=f"Converging ..."):
            try:
                result = await future
                results.append(result)
                f.write(json.dumps(result, ensure_ascii=False) + "\n")
                f.flush()
                os.fsync(f.fileno())
            except Exception as e:
                exception_type = type(e).__name__
                exception_message = str(e)
                traceback_info = ''.join(traceback.format_tb(e.__traceback__))
                error_message = f'{exception_type}: {exception_message}\n' \
                                f'Traceback:\n{traceback_info}'
                print(f"[ERROR]: {error_message}")


if __name__ == '__main__':