judge.py

import json
from openai import OpenAI
import argparse
import os
import re
from tqdm import tqdm
from concurrent.futures import ThreadPoolExecutor, as_completed

# print("OPENAI_API_KEY" in os.environ)
# print(os.environ.get("OPENAI_API_KEY"))
client = OpenAI(api_key=os.environ["OPENAI_API_KEY"])

# Configuration
INPUT_FILE = "result.json"
OUTPUT_FILE = "judged_result.json"
HOP_FILE = "datasheet/matthew_sheet_processed.json"

ANSWER_ONLY_PROMPTS = [1, 2, 4, 5]  # Prompts with tag _p1/2/4/5 do not use multi-hop system
MULTI_HOP_PROMPTS = [3, 6, 7, 8]  # Prompts with tag _p3/6/7/8 do use multi-hop and are graded on it

MODEL = "gpt-5"
debug = False

# Functions
def build_judge_prompt(entry, prompt_id=0):
    """
    Helper function containing the LLM as judge prompt. Also formats each json entry to be readable by LLM.

    :param entry: (dict) Converted json file containing all the experiment results
    :param prompt_id: (int) Which prompt to use. 0 means answer only, 1 means answer + multi-hop reasoning
    :return: (str) Judge prompt
    """

    assert isinstance(entry, dict)
    assert isinstance(prompt_id, int)
    assert re.search(r"_p[1-8]$", entry['index'])  # Check if a correct index is in entry (each have the _pX tag)
    assert 0 <= prompt_id <= 1

    if prompt_id == 0:
        prompt = f"""
        Here is a response and the corresponding ground truth answer. Grade it following the questions and rubric given below. 

        Prompt: {entry['prompt']}
        Ground Truth Answer: {entry['answer']}
        Model Response: {entry['response']}

        Answer the following questions about the **overall quality of the model's answer**. Use a 1-5 range and a 
        0/1 binary label, respectively, with better answers corresponding to higher score values.

        1. **Does the model’s final answer correctly and clearly answer the question?**  

        2. **Does the response make sense overall?**  

        Format the final grades in the following JSON format:

        {{
          "whole_response": {{
            "overall_answer_quality": x
            "overall_answer_quality_bin": x
          }},
          "whole_reasoning_chain": {{
            "makes_sense": x,
            "makes_sense_bin": x,
          }}
        }}
        """
    else:
        prompt = f"""
        Here is a multi-hop reasoning response and the corresponding ground truth answer. First organize the response's reasoning into
        a multi-hop reasoning chain following the hops: (object) recognition, action recognition, spatial, context, punch linguistic. 
        The first two are hops recognizing parts of the image, spatial is a hop relating parts of the image in space, context is bringing in specific
        outside knowledge, punch linguistic is connecting previous parts to form the punchline.
        
        Then grade it following the questions and rubric given below. Use a 1-5 range and a 
        0/1 binary label, respectively, on questions 1 and 2, while the rest are 0/1 binary. Better answers correspond to higher score values.
    
        ---
        
        Prompt: {entry['prompt']}
        Ground Truth Answer: {entry['answer']}
        Ground Truth Hops: {entry['original_annotation'] if entry['original_annotation'] else 'N/A'}
        Model Response: {entry['response']}
        
        ---
        
        Whole Response Evaluation:
        
        **1. Does the whole generated response answer the question in a well-justified manner?**  
        (1 = incomprehensible and wrong, 5 = clear and correct)
        
        ---
        
        Questions for the Entire Reasoning Chain:
        
        **2. Does the whole generated response make sense?**  
        (1 = sounds like nonsense, 5 = easy to parse)
        
        **3. Is the content of the generated reasoning incomplete and lacking required information to produce the correct answer?**  
        If missing steps were added, the model could produce the correct answer.  
        (Example: “A has 3 apples, B has 2. How much more apples A has?”  
        Chain: “A has 3 apples. The answer is 3.” Missing step: “So A has 3 - 2 = 1 more.”)
        
        ---
        
        Questions for Each Individual Hop from the Response's Reasoning Chain:
        
        **4. Does this step use reasoning of the same type as one of the ground truth hops?**  
        (e.g., if the model hop is "spatial," does ground truth involve spatial reasoning?)
        
        **5. Is this step redundant or unnecessary for reaching the next hops or the final answer?**
        (If this is an extra hop added between two hops that match with ground truth, it would be considered unnecessary.)
        
        **6. Does the hop explanation match the ground truth explanation?**  
        ---
        
        Return the results in the following JSON format:
        
        {{
          "whole_response": {{
            "overall_answer_quality": x
            "overall_answer_quality_bin": x
          }},
          "whole_reasoning_chain": {{
            "makes_sense": x,
            "makes_sense_bin": x,
            "content_incomplete": x,
          }},
          "<first hop name>": {{
            "hop_type_correct": x,
            "redundant_hop": x,
            "correct_reasoning": x,
          }},
          "<second hop name>": {{
            ...
          }}
          ... (additional hops as needed)
        }}
        """

    return prompt

def normalize_hops_from_sheet(entry, max_hops=10):
    """
    Extract hops from flattened sheet-style fields and attach
    a normalized hop structure to the entry.

    Produces:
        entry["_hops"] = [
            {
                "hop_index": 1,
                "hop_id": "...",
                "hop_type": "...",
                "original_annotation": "...",
                "prompt": "...",
                "model_answer": "..."
            },
            ...
        ]
    """
    hops = []

    for i in range(1, max_hops + 1):
        hop_id = entry.get(f"hop{i}_id")
        hop_type = entry.get(f"hop{i}_type")
        hop_ann = entry.get(f"hop{i}_original_annotation")
        hop_prompt = entry.get(f"hop{i}_prompt")
        hop_ans = entry.get(f"hop{i}_answers_value")

        # Skip empty / non-existent hops
        if not any(
            isinstance(x, str) and x.strip()
            for x in [hop_id, hop_type, hop_ann, hop_prompt, hop_ans]
        ):
            continue

        hops.append({
            "hop_index": i,
            "hop_id": hop_id,
            "hop_type": hop_type,
            "original_annotation": hop_ann,
            "prompt": hop_prompt,
            "model_answer": hop_ans,
        })

    entry["_hops"] = hops
    entry["_num_hops"] = len(hops)

# def add_hops(data, hoppath):
#     assert isinstance(hoppath, str)
#     with open(hoppath, "r", encoding="utf-8") as f:
#         text = f.read()
#
#     decoder = json.JSONDecoder()
#     hop_entries = []
#     idx = 0
#
#     while idx < len(text):
#         text = text.lstrip()
#         if not text:
#             break
#         obj, end = decoder.raw_decode(text)
#         hop_entries.extend(obj)
#         text = text[end:]
#
#     hop_map = {
#         hop["subject_id"]: hop.get("original_annotation")
#         for hop in hop_entries
#         if "subject_id" in hop
#     }
#     # print(hop_map)
#
#     for entry in data:
#         subject_id = entry["index"].split("_")[0]
#         if subject_id in hop_map:
#             original_ann = hop_map[subject_id]
#             # entry["_hop_order"] = list(original_ann.keys())
#             entry["original_annotation"] = original_ann
#         else:
#             if debug:
#                 print('not working')
#             else:
#                 assert NotImplementedError, 'entry not matching with hops'
#
#     return data

def remap_judge_response_hops(entry):
    """
    Normalize judge_response to hop1, hop2, ... order
    using the extracted _hops structure.
    """
    jr = entry.get("judge_response")
    hops = entry.get("_hops")

    if not isinstance(jr, dict) or not isinstance(hops, list):
        return

    new_jr = {}
    used_keys = set()

    for idx, hop in enumerate(hops, start=1):
        hop_type = hop["hop_type"]

        if hop_type in jr:
            new_jr[f"hop{idx}"] = {
                "hop_type": hop_type,
                **jr[hop_type]
            }
            used_keys.add(hop_type)

    # Preserve non-hop judge fields
    for k, v in jr.items():
        if k not in used_keys:
            new_jr[k] = v

    entry["judge_response"] = new_jr


def get_active_hops(entry):
    """
    Return hop indices actually present in the entry.
    """
    return [hop["hop_index"] for hop in entry.get("_hops", [])]

def build_hop_prompt(entry, hop_idx):
    """
    Build LLM inference prompt for a single hop.
    """
    hop_prompt = entry.get(f"hop{hop_idx}_prompt")
    question = entry.get("question")

    if not isinstance(hop_prompt, str):
        raise ValueError(f"Missing hop{hop_idx}_prompt")

    return hop_prompt

def process_file(filepath, outpath, hoppath=None):
    """
    Load model outputs, grade them using judge prompts,
    and save graded JSON.
    """
    assert isinstance(filepath, str)
    assert isinstance(outpath, str)

    # ----------------------------
    # Load model inference outputs
    # ----------------------------
    with open(filepath, "r", encoding="utf-8") as f:
        print(f"Loading {filepath}")
        data = json.load(f)

    assert isinstance(data, list), "Input JSON must be a list of entries"

    # ----------------------------
    # Attach gold hop annotations
    # ----------------------------
    if hoppath is not None:
        data = add_hops(data, hoppath)

    prompt_items = []

    # ----------------------------
    # Build judge prompts ONLY
    # ----------------------------
    for i, entry in enumerate(data):
        try:
            # Judge prompt always assumes hops exist
            prompt = build_judge_prompt(entry, prompt_id=1)
            prompt_items.append((i, prompt))
        except Exception as e:
            entry["judge_response"] = {"error": str(e)}

    # ----------------------------
    # Run grading in parallel
    # ----------------------------
    results = run_all(prompt_items)

    # ----------------------------
    # Attach judge responses
    # ----------------------------
    for entry_idx, out in results.items():
        data[entry_idx]["judge_response"] = out

        # Optional: remap judge hops to hop1, hop2, ...
        if "original_annotation" in data[entry_idx]:
            hop_order = list(data[entry_idx]["original_annotation"].keys())
            remap_judge_response_hops(data[entry_idx], hop_order)

    # ----------------------------
    # Save graded results
    # ----------------------------
    with open(outpath, "w", encoding="utf-8") as f:
        json.dump(data, f, indent=4, ensure_ascii=False)

    print(f"Grading complete. Saved to {outpath}")


def run_single_eval(prompt: str):
    """Make a single GPT API request."""
    try:
        if debug:
            return {'debug': True}

        completion = client.chat.completions.create(
            model=MODEL,
            messages=[
                {"role": "user", "content": prompt}
            ],
            response_format={"type": "json_object"},
        )
        return json.loads(completion.choices[0].message.content)

    except Exception as e:
        assert 'Error code: 429' not in str(e), str(e)
        assert 'Error code: 401' not in str(e), str(e)
        return {"error": str(e), "prompt": prompt}


def run_all(prompt_items, max_workers=10):
    """Multithread GPT evaluation."""
    results = {}

    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        future_to_idx = {
            executor.submit(run_single_eval, prompt): idx
            for idx, prompt in prompt_items
        }

        for future in as_completed(future_to_idx):
            entry_idx = future_to_idx[future]
            try:
                result = future.result()
            except Exception as e:
                assert 'Error code: 429' not in str(e), str(e)
                assert 'Error code: 401' not in str(e), str(e)
                result = {"error": str(e)}
            results[entry_idx] = result

    return results

def extract_boxed(text: str):
    """
    Extract content inside \\boxed{...}, stripping optional \\text{}.
    """
    if not isinstance(text, str):
        return None

    match = re.search(
        r"\\boxed\s*\{([\s\S]*?)\}",
        text
    )
    if not match:
        return None

    inner = match.group(1).strip()

    # Strip \text{...} if present
    text_match = re.fullmatch(
        r"\\text\s*\{([\s\S]*?)\}",
        inner
    )
    if text_match:
        inner = text_match.group(1).strip()

    return inner

def normalize_response(entry):
    """
    Join response list into a single string.
    """
    resp = entry.get("response")
    if isinstance(resp, list):
        return "\n".join(x for x in resp if isinstance(x, str))
    if isinstance(resp, str):
        return resp
    return None

def judge_direct_question(entry):
    pred_text = normalize_response(entry)
    ref_text = entry.get("answer")

    if not pred_text or not ref_text:
        return {"error": "missing prediction or reference"}

    boxed = extract_boxed(pred_text)
    if boxed is None:
        return {"error": "missing boxed answer"}

    return {
        "direct": compute_rouge(boxed, ref_text)
    }


def main():
    description = """
    This is a script that takes some arguments (configure at top of judge.py) pointing to the location of a json file and
    parse it in a way that lets us pass it into GPT using a predetermined judge prompt. Add to said json the
    outputs of the LLM as a judge, following whatever scores (most should be binary labels) are created. For multi-hop prompts, 

    How to run: >>> python judge.py --name <name of folder>.
        Only include --name if you want to look at a single .json, otherwise it looks for result.json in
         all directories inside the working directory.
        e.g. >>> python judge.py --name test_folder
    """

    parser = argparse.ArgumentParser(description=description)
    parser.add_argument("--name", type=str, help="Name of folder if looking at only one model")
    parser.add_argument("--exclude", type=str, nargs='+', help="Name of folder(s) to exclude, include rest")
    parser.add_argument("--include", type=str, nargs='+', help="Name of folder(s) to include, exclude rest")
    args = parser.parse_args()

    assert not (args.include and args.exclude)

    if args.name:  # Specific model
        folder = args.name
        filepath = os.path.join(folder, INPUT_FILE)
        outpath = os.path.join(folder, OUTPUT_FILE)
        
        if os.path.exists(filepath):
            process_file(filepath, outpath, HOP_FILE)
        else:
            print(f"Could not find {INPUT_FILE} in {folder}")
            
    else:  # All model folders in working directory, with exclusion check
        for folder in os.listdir("."):
            if folder in args.exclude and args.exclude is not None:
                continue
            elif folder not in args.include and args.include is not None:
                continue
            if os.path.isdir(folder):
                filepath = os.path.join(folder, INPUT_FILE)
                outpath = os.path.join(folder, OUTPUT_FILE)
                
                if os.path.exists(filepath):
                    process_file(filepath, outpath)


        

if __name__ == "__main__":
    main()