eval_cube_api.py

import os
import random
import numpy as np
import argparse
import json
from tqdm import tqdm
from openai import OpenAI
from anthropic import Anthropic
import io
import base64
from concurrent.futures import ThreadPoolExecutor, as_completed
from datasets import load_dataset
from cube_helper import extract_happycube_solution, extract_2d_array, SolutionValidator

parser = argparse.ArgumentParser()
parser.add_argument('--model_name', type=str, choices=["gpt-4o", "o3", "o4-mini", "deepseek-reasoner", "gemini-2.5-pro-preview-05-06", 
                                                       'claude-3-7-sonnet-20250219', 'doubao-1-5-thinking-vision-pro-250428'], default='gpt-4o')
parser.add_argument('--subset', type=str, default='cube_easy', choices=["cube", "cube_easy", "cube_perception"])
parser.add_argument('--num_workers', type=int, default=1)
parser.add_argument('--seed', type=int, default=20)
args = parser.parse_args()
random.seed(args.seed)
np.random.seed(args.seed)

def encode_pil_image(pil_image):
    buffered = io.BytesIO()
    pil_image.save(buffered, format="PNG")
    img_bytes = buffered.getvalue()
    return base64.b64encode(img_bytes).decode("utf-8")

def get_message(sample):
    prompt = sample['question']
    image_info = []
    if args.subset in ['cube', 'cube_perception']:
        pil_image = sample['image']
        base64_image = encode_pil_image(pil_image)
        if 'claude' not in args.model_name:
            image_info.append({
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:image/png;base64,{base64_image}",
                            },
                        })
        else:
            image_info.append({
                            "type": "image", 
                            "source": {
                                "type": "base64",
                                "media_type": 'image/png',
                                "data": base64_image
                            }
                        })
    messages = [
            {
                "role": "user",
                "content": [
                    *image_info,
                    { "type": "text", "text": prompt}
                ]
            }
        ]
    return messages

def remove_image(messages):
    messages[0]['content'] = [content if 'image' not in content['type'] else '<|image|>' for content in messages[0]['content']]
    return messages

def call_api(model_name, index, sample, client, log_file):  
    if os.path.exists(log_file):
        with open(log_file, 'r', encoding='utf8') as f:
                ids = [json.loads(l)['id'] for l in f.readlines()]
        if index in ids:
            print(f'{index} already done!')
            return

    messages = get_message(sample) 
    try:
        if 'deepseek' in model_name or 'doubao' in model_name:
            output = client.chat.completions.create(
                model=model_name,
                messages=messages,
                stream=False,
                temperature=0.,
            )
        elif 'o4' in model_name:
            output = client.chat.completions.create(    
                model=model_name,    
                messages=messages,    
                stream=False,    
                reasoning_effort='medium',    
                max_completion_tokens=25000
            )
        elif 'o3' in model_name:
            output = client.chat.completions.create(    
                model=model_name,    
                messages=messages,    
                stream=False,    
                reasoning_effort='medium',    
                max_completion_tokens=40000
            )
        elif 'gemini' in model_name:
            output = client.chat.completions.create(
                model=model_name,
                messages=messages,
                stream=False,
                reasoning_effort='medium',
                max_completion_tokens=25000
            )
        elif '4o' in model_name:
            output = client.chat.completions.create(
                model=model_name,
                messages=messages,
                stream=False,
                temperature=0.,
                max_tokens=16000,
            )
        elif 'claude' in model_name:
            output = client.messages.create(
                        model=model_name, 
                        max_tokens=16000,
                        thinking={
                            "type": "enabled",
                            "budget_tokens": 12000
                        },
                        messages=messages
                        )
        else:
            raise "Unsupported model"
        
        if 'claude' not in model_name:
            response = output.choices[0].message.content.strip()
            try:
                reasoning_content = output.choices[0].message.reasoning_content.strip()
            except:
                reasoning_content = ""
        else:
            reasoning_content = output.content[0].thinking
            response = output.content[1].text if len(output.content) > 1 else '' 

        if args.subtask in ['cube', 'cube_easy']:
            face_arrays = sample['face_arrays']
            validator = SolutionValidator()
            solution = extract_happycube_solution(response)
            if solution is None:
                solution = extract_happycube_solution(reasoning_content)

            correct = validator.validate(face_arrays, solution) if solution is not None else False

            result = {
                'id': index,
                'face_arrays': face_arrays,
                'messages': remove_image(messages),
                'response': response,
                'correct': correct,
                'reasoning_content': reasoning_content,
                'usage': str(output.usage)
            }
            print(f'{index} finished')
            with open(log_file, "a", encoding="utf-8") as log:
                log.write(json.dumps(result, ensure_ascii=False) + "\n")
        
        else: # cube_perception
            solution = extract_2d_array(response)
            gt_answer = np.array(sample['answer'])
            if solution is not None:
                per_cell_acc = np.equal(solution, gt_answer).mean()
                whole_piece_acc = bool(np.all(np.equal(solution, gt_answer)))
            else:
                per_cell_acc, whole_piece_acc = 0., 0.

            result = {
                'id': index,
                'messages': remove_image(messages),
                'answer': sample['answer'],
                'response': response,
                'per_cell_acc': per_cell_acc,
                'whole_piece_acc': whole_piece_acc,
                'reasoning_content': reasoning_content,
                'usage': str(output.usage)
            }
            print(f'{index} finished')
            with open(log_file, "a", encoding="utf-8") as log:
                log.write(json.dumps(result, ensure_ascii=False) + "\n")

    except Exception as e:
        print(f"Error processing sample {index}: {e}")


def get_result(log_file, subset):
    if subset in ['cube', 'cube_easy']:
        with open(log_file, 'r') as f:
            accuracies = [json.loads(l)['correct'] for l in f.readlines()]
        return {'accuracy': np.mean(accuracies)}
    else: 
        # cube_perception
        with open(log_file, 'r') as f:
            per_cell_accuracies, whole_piece_accuracies = [], []
            for l in f.readlines():
                per_cell_accuracies.append(json.loads(l)['per_cell_acc'])
                whole_piece_accuracies.append(json.loads(l)['whole_piece_acc'])
        return {'per_cell_accuracy': np.mean(per_cell_accuracies), 'whole_piece_accuracy': np.mean(whole_piece_accuracies)}
    
if __name__ == '__main__':
    save_path = f'output/{args.subset}/'
    log_file = f'{save_path}/{args.model_name}.jsonl'
    os.makedirs(save_path, exist_ok=True)

    dataset = load_dataset('mrble/MARBLE', args.subset)['train']

    if 'deepseek' in args.model_name:
        client = OpenAI(api_key='<YOUR_API_KEY>', base_url="https://api.deepseek.com")
    elif 'gemini' in args.model_name:
        client = OpenAI(api_key='<YOUR_API_KEY>', base_url='https://generativelanguage.googleapis.com/v1beta/openai/')
    elif 'doubao' in args.model_name:
        client = OpenAI(api_key='<YOUR_API_KEY>', base_url="https://ark.cn-beijing.volces.com/api/v3")
    elif 'claude' in args.model_name:
        client = Anthropic(api_key='<YOUR_API_KEY>')
    else:
        client = OpenAI(api_key='<YOUR_API_KEY>')
    
    if args.num_workers == 1:
        for i, sample in tqdm(enumerate(dataset)):
            call_api(args.model_name, i, sample, client, log_file)
    else:
        with ThreadPoolExecutor(max_workers=args.num_workers) as executor:  # Adjust max_workers based on your rate limit
            futures = [
                executor.submit(
                    call_api, args.model_name, i, sample, client, log_file
                ) for i, sample in enumerate(dataset)
            ]

            for _ in tqdm(as_completed(futures), total=len(futures)):
                pass
    
    result = get_result(log_file, args.subtask)

    with open('output/final_result.jsonl', 'a') as f:
        f.write(
            json.dumps(
                {
                    'subset': args.subset,
                    'model': args.model_name,
                    **result
                }
            ) + '\n'
        )