def load_responses(checkpoint, dataset):
    """Load responses for a specific checkpoint and dataset"""
    checkpoint_path = os.path.join(DATA_DIR, checkpoint, dataset)
    responses = []
    
    # Find first JSONL file in the directory
    jsonl_files = [f for f in os.listdir(checkpoint_path) if f.endswith(".jsonl") and f"_t{TEMPERATURE}" in f]
    if not jsonl_files:
        return []
    
    with open(os.path.join(checkpoint_path, jsonl_files[0])) as f:
        for line in f:
            data = json.loads(line)
            responses.append({
                "idx": data["idx"],
                "input": data["question"],
                "pred": data["code"][0],  # Assuming single prediction per example
                "score": data["score"][0],
                "stop_reason": data["finish_reason"][0]  # Add stop reason
            })
    
    # Sort by index to ensure alignment between checkpoints
    responses.sort(key=lambda x: x["idx"])
    return responses
Load responses for a specific checkpoint and dataset
examples/simplelr_math_eval/visualizer.py