method.py

import numpy as np
import os
import utils
import arc_dataset
import draw
import ast
import openai
from openai import OpenAI

def VLSR_Extract_Rule(client, model="o4-mini", data=None, id=0, save_root="images/ARC-AGI/"):
    prompt = r"I will now provide you with several input and output images about 2D grids. You need to summarize the grid-changing rule from it. Output the rule you learned within \\boxed{}."
    conversation_history = []
    content = []
    example_num = len(data["example"])
    input_test = data["input_test"]
    gt = data["gt"]
    os.makedirs(save_root, exist_ok=True)


    content.append({"type": "text", "text": prompt})
    for i in range(example_num):
        input_example, output_example = data["example"][i]
        os.makedirs(os.path.join(save_root, str(id)), exist_ok=True)
        save_path_input = os.path.join(save_root, str(id), f"input_{i}.jpg")
        save_path_output = os.path.join(save_root, str(id), f"output_{i}.jpg")
        draw.generate_grid_image(grid=input_example, output_filename=save_path_input)
        draw.generate_grid_image(grid=output_example, output_filename=save_path_output)

        content.append({"type": "text", "text": "Example input " + str(i+1)})
        content.append({"type": "image_url", "image_url":{"url": f"data:image/jpeg;base64,{utils.encode_image(save_path_input)}" }},)
        content.append({"type": "text", "text": "Example output " + str(i+1)})
        content.append({"type": "image_url", "image_url":{"url": f"data:image/jpeg;base64,{utils.encode_image(save_path_output)}" }},)

    conversation_history.append({"role": "user",
        "content": content
        }
    )
    response = client.chat.completions.create(
        model=model,
        messages = conversation_history,
        temperature = 0.7,
        n = 1
    )
    rule = response.choices[0].message.content
    rule_extracted = utils.extract_answer(rule)
    return rule_extracted
    
def VLSR_Apply_Rule(client, model="o4-mini", data=None, id=0, rule=None):
    prompt = r"I will provide you with several input and output matrices. You need to find the matrix-changing rule from it and  apply it to the new input. Put the output matrix within \\boxed{}."
    conversation_history = []
    content = []
    example_num = len(data["example"])
    input_test = data["input_test"]
    gt = data["gt"]


    content.append({"type": "text", "text": prompt})
    if rule is not None:
        prompt_rule = r"Here is a possible rule for your reference. Note that the rule is described in color and each color represents a value in the matrix: [0:black; 1:blue; 2:red; 3:green; 4:yellow; 5:grey; 6:pink; 7:orange; 8:light blue; 9:brown]. You need to first check the correctness of the rule based on the examples. If the rule is correct, apply it to the new input. Otherwise, summarize a new rule and apply it to the new input."
        content.append({"type": "text", "text": prompt_rule})
        content.append({"type": "text", "text": "Rule: " + str(rule)})

    for i in range(example_num):
        input_example, output_example = data["example"][i]
        content.append({"type": "text", "text": "Example input " + str(i+1) + "\n"})
        content.append({"type": "text", "text": str(input_example)})
        content.append({"type": "text", "text": "Example output " + str(i+1) + "\n"})
        content.append({"type": "text", "text": str(output_example)})

    content.append({"type": "text", "text": "New Input: " + "\n"})
    content.append({"type": "text", "text": str(input_test)})

    conversation_history.append({"role": "user",
        "content": content
        }
    )
    response = client.chat.completions.create(
        model=model,
        messages = conversation_history,
        temperature = 0.7,
        n = 1
    )
    answer = response.choices[0].message.content
    answer_extracted = utils.extract_answer(answer)
    if answer_extracted is None or answer_extracted == "":
        return answer
    return answer_extracted

def VLSR(client, model="o4-mini", data=None, id=0, save_root="images/ARC-AGI/"):
    rule = VLSR_Extract_Rule(client, model, data, id, save_root)
    answer = VLSR_Apply_Rule(client, model, data, id, rule)
    return answer, rule

def MSSC_Check(client, model="o4-mini", data=None, id=0, save_root="images/ARC-AGI/", answer = None):
    prompt = r"I will now provide you with several input and output example images, which follows a specific changing rule. Then, I will give you another input and output pair, determine whether the new pair also follows the same changing rule. Add your final judgment at the end of your replay: \\boxed{True} or \\boxed{False}."
    conversation_history = []
    content = []
    example_num = len(data["example"])
    input_test = data["input_test"]
    gt = data["gt"]
    os.makedirs(save_root, exist_ok=True)

    content.append({"type": "text", "text": prompt})
    for i in range(example_num):
        os.makedirs(os.path.join(save_root, str(id)), exist_ok=True)
        input_example, output_example = data["example"][i]
        save_path_input = os.path.join(save_root, str(id), f"input_{i}.jpg")
        draw.generate_grid_image(grid=input_example, output_filename=save_path_input)
        save_path_output = os.path.join(save_root, str(id), f"output_{i}.jpg")
        draw.generate_grid_image(grid=output_example, output_filename=save_path_output)

        content.append({"type": "text", "text": "Example input " + str(i+1)})
        content.append({"type": "image_url", "image_url":{"url": f"data:image/jpeg;base64,{utils.encode_image(save_path_input)}" }},)
        content.append({"type": "text", "text": "Example output " + str(i+1)})
        content.append({"type": "image_url", "image_url":{"url": f"data:image/jpeg;base64,{utils.encode_image(save_path_input)}" }},)

    save_path_input = os.path.join(save_root, str(id), f"input_test.jpg")
    draw.generate_grid_image(grid=input_test, output_filename=save_path_input)

    try:
        answer_list = ast.literal_eval(answer)
    except:
        answer_list = utils.latex_matrix_to_list(answer)

    if answer_list is None:
        return None

    try:  
        save_path_output = os.path.join(save_root, str(id), f"output_pred.jpg")
        draw.generate_grid_image(grid=answer_list, output_filename=save_path_output)
    except:
        return None

    content.append({"type": "text", "text": "New input: "})
    content.append({"type": "image_url", "image_url":{"url": f"data:image/jpeg;base64,{utils.encode_image(save_path_input)}" }},)
    content.append({"type": "text", "text": "New output: "})
    content.append({"type": "image_url", "image_url":{"url": f"data:image/jpeg;base64,{utils.encode_image(save_path_output)}" }},)

    conversation_history.append({"role": "user",
        "content": content
        }
    )
    response = client.chat.completions.create(
        model=model,
        messages = conversation_history,
        temperature = 0.7,
        n = 1
    )
    judge = response.choices[0].message.content
    judge_extracted = utils.extract_answer(judge)
    return judge_extracted
        

def MSSC_Correction(client, model="o4-mini", data=None, id=0, answer = None, judge=None):

    prompt = r"I will now provide you with several input and output examples following a specific matrix changing rule. You need to summarize the matrix-changing rule and apply it to the new input. However, the original result you get is likely wrong. Check the result carefully and output a new output matrix within \\boxed{}."
    conversation_history = []
    content = []
    example_num = len(data["example"])
    input_test = data["input_test"]
    gt = data["gt"]

    for i in range(example_num):
        input_example, output_example = data["example"][i]
        content.append({"type": "text", "text": "Example input " + str(i+1) + "\n"})
        content.append({"type": "text", "text": str(input_example)})
        content.append({"type": "text", "text": "Example output " + str(i+1) + "\n"})
        content.append({"type": "text", "text": str(output_example)})

    content.append({"type": "text", "text": "New Input: " + "\n"})
    content.append({"type": "text", "text": str(input_test)})
    content.append({"type": "text", "text": "Original Output: " + "\n"})
    content.append({"type": "text", "text": str(answer)})

    conversation_history.append({"role": "user",
        "content": content
        }
    )
    response = client.chat.completions.create(
        model=model,
        messages = conversation_history,
        temperature = 0.7,
        n = 1
    )
    answer = response.choices[0].message.content
    answer_extracted = utils.extract_answer(answer)
    if answer_extracted is None or answer_extracted == "":
        return answer
    return answer_extracted


def MSSC(client, model="o4-mini", data=None, id=0, save_root="images/ARC-AGI/", answer=None, max_round=3):
    for round in range(max_round):
        judge = MSSC_Check(client, model, data, id ,save_root, answer)
        if judge is None:
            return answer
        if judge == "True":
            return answer
        answer_new = MSSC_Correction(client, model, data, id, answer, judge)
        answer = answer_new


    return answer_new

def Reason(client, model="o4-mini", data=None, id=0, save_root="images/ARC-AGI/", max_round=3):
    answer, rule = VLSR(client, model, data, id, save_root)
    answer = MSSC(client, model, data, id, save_root, answer, max_round)
    return answer, rule