Function main

codeclash/analysis/viz/heatmap_returncode.py:17–151

(logs: Path)

def main(logs: Path):
    model_to_arena_to_return_codes = {}
    if not DATA_CACHE.exists():
        tournaments = [x.parent for x in logs.rglob("metadata.json")]
        for game_log_folder in tqdm(tournaments):
            arena = game_log_folder.name.split(".", 2)[1]
            with open(game_log_folder / "metadata.json") as f:
                metadata = json.load(f)
            try:
                p2m = {
                    x["name"]: x["config"]["model"]["model_name"].strip("@").split("/")[-1]
                    for x in metadata["config"]["players"]
                }
                for model in p2m.values():
                    if model not in model_to_arena_to_return_codes:
                        model_to_arena_to_return_codes[model] = {}
                    if arena not in model_to_arena_to_return_codes[model]:
                        model_to_arena_to_return_codes[model][arena] = {"actions": 0, "malformed": 0}
            except KeyError:
                continue

            for name in p2m.keys():
                traj_files = (game_log_folder / "players" / name).rglob("*.traj.json")
                for traj_file in traj_files:
                    with open(traj_file) as f:
                        traj = json.load(f)
                    msgs = traj["messages"][2:]  # Skip system prompt and first user message
                    observations = [msg for msg in msgs if msg["role"] == "user"]
                    actions = len(observations)
                    malformed = 0
                    for obs in observations:
                        if isinstance(obs["content"], str):
                            malformed += "<returncode>0</returncode>" not in obs["content"]
                        elif isinstance(obs["content"], list):
                            malformed += "<returncode>0</returncode>" not in obs["content"][0]["text"]
                        else:
                            print(f"Unknown content type: {type(obs['content'])}")
                    model_to_arena_to_return_codes[p2m[name]][arena]["actions"] += actions
                    model_to_arena_to_return_codes[p2m[name]][arena]["malformed"] += malformed

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

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

    model_to_return_code = defaultdict(tuple)
    for model, arena_to_return_codes in model_to_arena_to_return_codes.items():
        a, m = 0, 0
        for arena, data in arena_to_return_codes.items():
            malform_rate = data["malformed"] / data["actions"] * 100
            a += data["actions"]
            m += data["malformed"]
            print(
                f"- {model} in {arena}: {data['actions']} actions; {data['malformed']} malformed ({malform_rate:.2f}%)"
            )
        model_to_return_code[model] = (m, a, m / a * 100)
    print("=" * 20)