Function main

codeclash/analysis/viz/line_chart_steps_per_round.py:15–110

()

def main():
    model_to_round_steps = {}

    if not DATA_CACHE.exists():
        tournaments = [x.parent for x in LOCAL_LOG_DIR.rglob("metadata.json")]
        for game_log_folder in tqdm(tournaments, desc="Scanning tournaments"):
            try:
                with open(game_log_folder / "metadata.json") as f:
                    metadata = json.load(f)
            except Exception:
                continue

            try:
                p2m = {
                    x["name"]: x["config"]["model"]["model_name"].strip("@").split("/")[-1]
                    for x in metadata["config"]["players"]
                }
            except Exception:
                # malformed metadata
                continue

            # ensure models exist
            for model in set(p2m.values()):
                model_to_round_steps.setdefault(model, [[] for _ in range(ROUNDS)])

            for player_name, model in p2m.items():
                traj_files = (game_log_folder / "players" / player_name).rglob("*.traj.json")
                for traj_file in traj_files:
                    m = traj_file.name.rsplit("_r", 1)
                    if len(m) != 2:
                        continue
                    try:
                        round_part = m[1]
                        round_idx = int(round_part.split(".")[0])
                    except Exception:
                        continue
                    if round_idx < 1 or round_idx > ROUNDS:
                        continue

                    try:
                        with open(traj_file) as f:
                            traj = json.load(f)
                    except Exception:
                        continue

                    num_steps = sum(1 for _ in traj.get("messages", []) if _["role"] == "assistant")
                    model_to_round_steps[model][round_idx - 1].append(num_steps)

        with open(DATA_CACHE, "w") as f:
            json.dump(model_to_round_steps, f, indent=2)

    with open(DATA_CACHE) as f:
        model_to_round_steps = json.load(f)

    # Compute averages per round and plot
    model_to_avg = {}
    for model, rounds_lists in model_to_round_steps.items():
        # pad/truncate