convert_datasheet.py

import json
from openai import OpenAI
import argparse
import os
import pandas as pd
import glob
import base64
import re
import math
from concurrent.futures import ThreadPoolExecutor, as_completed

client = OpenAI()

# Configuration
MODEL = "gpt-5"
debug = False


def build_prompt(entry):
    hops = list(entry["original_annotation"].items())

    hop_descriptions = []
    for i, (hop_type, hop_value) in enumerate(hops):
        hop_descriptions.append(
            f"Hop {i+1} ({hop_type}): {hop_value}"
        )

    prompt = f"""
Using attached image, generate queries for this example.

Main question:
"{entry['question']}"
Punchline:
"{entry['answer']}"

Ground-truth reasoning steps (do NOT restate these, only use them to guide question creation):
{chr(10).join(hop_descriptions)}

Tasks:
1. Generate ONE question per hop that logically leads from one hop to the next.
   - Question 1 leads to hop 1
   - Question 2 leads from hop 1 to hop 2
   - etc.

Rules:
- Output questions only.
- Do NOT mention hop names or types.
- Do NOT invent new facts.

Return JSON only in this format:
{{
  "multihop_queries": ["...", "..."]
}}
"""
    if debug:
        print(prompt)
    return prompt.strip()

def process_file(filepath, outpath, imgpath, max_workers=10):
    """
    Load CSV, infer queries in parallel, and save JSON. Parallel GPT inference determined by max_workers.
    """
    assert isinstance(filepath, str)
    assert isinstance(outpath, str)
    assert isinstance(imgpath, str)

    print(filepath)
    data = pd.read_csv(filepath, on_bad_lines="warn", header=0)
    data.columns = [re.sub(r"\s+", "_", col) for col in data.columns]

    # Image lookup
    extensions = ["*.jpg", "*.jpeg", "*.png", "*.bmp", "*.webp", "*.tiff"]
    image_map = {}
    for pattern in extensions:
        for full_path in glob.glob(os.path.join(imgpath, pattern)):
            key = os.path.splitext(os.path.basename(full_path))[0]
            image_map[key.strip()] = full_path

    cleaned_entries = []
    inference_items = []

    # ---- deterministic preprocessing ----
    for idx, entry in enumerate(data.itertuples(index=False)):
        entry = entry._asdict()

        img_name = entry["Image_name"].strip()
        if img_name not in image_map:
            print(f"Missing image for {img_name}")
            continue

        img_b64, mime = encode_image(image_map[img_name])
        cleaned = clean_entry(entry)

        cleaned_entries.append(cleaned)
        inference_items.append((len(cleaned_entries) - 1, cleaned, img_b64, mime))

    # ---- parallel GPT calls ----
    results = run_all_query_inference(inference_items, max_workers=max_workers)

    # ---- deterministic postprocessing ----
    processed_entries = []
    for idx, cleaned in enumerate(cleaned_entries):
        query_output = results.get(idx, {"error": "missing result"})

        if "error" in query_output:
            cleaned["error"] = query_output["error"]
            processed_entries.append(cleaned)
            continue

        final_entry = attach_reasoning(cleaned, query_output)
        processed_entries.append(final_entry)

    with open(outpath, "w", encoding="utf-8") as f:
        json.dump(processed_entries, f, indent=4, ensure_ascii=False)

    print(f"Processed {filepath}, response recorded to {outpath}")

def infer_queries(entry, img_b64, mime, client, model):
    prompt = build_prompt(entry)

    if debug:
        output = {
            "singlehop_query": True,
            "multihop_queries": [True, True],
            "answer": True
                  }
        return output
    else:
        try:
            completion = client.chat.completions.create(
                model=model,
                messages=[
                    {
                        "role": "system",
                        "content": "You generate dataset queries only. Do not explain. Output JSON only."
                    },
                    {
                        "role": "user",
                        "content": [
                            {
                                "type": "image_url",
                                "image_url": {
                                    "url": f"data:image/{mime};base64,{img_b64}"
                                }
                            },
                            {
                                "type": "text",
                                "text": prompt
                            }
                        ]
                    }
                ],
                response_format={"type": "json_object"},
            )
            output = json.loads(completion.choices[0].message.content)
            # print(output)
            return output
        except Exception as e:
            print(str(e))
            assert 'Error code: 429' not in str(e), str(e)

def clean_entry(entry):
    """
    entry: dict from dataframe row
    subject_id == image_name
    """

    image_name = entry["Image_name"].strip()

    cleaned = {
        "subject_id": image_name,   # EXACT match
        "question": entry.get("question"),
        "answer": entry.get("answer"),
        "original_annotation": {}
    }

    hop_pattern = re.compile(r"hop(\d+)$")

    def is_valid(val):
        return val is not None and not (isinstance(val, float) and math.isnan(val))

    for key, value in entry.items():
        match = hop_pattern.match(key)
        if not match or not is_valid(value):
            continue

        idx = match.group(1)
        hop_type = entry.get(f"hop{idx}_type")

        if not is_valid(hop_type):
            continue

        cleaned["original_annotation"][hop_type] = value

    return cleaned

def attach_reasoning(entry, query_output):
    subject_id = entry["subject_id"]
    hops = list(entry["original_annotation"].items())

    if debug:
        print(query_output)
    # Single-hop
    entry["singlehop_reasoning"] = [
        {
            "test_queries": [
                {
                    # "prompt": query_output["singlehop_query"],
                    "prompt": entry["question"],
                    "answers": [
                        {
                            "value": entry["answer"],
                            "aliases": []
                        }
                    ]
                }
            ]
        }
    ]

    # Multi-hop
    multihop_queries = []
    for i, ((hop_type, hop_value), hop_out) in enumerate(
        zip(hops, query_output["multihop_queries"])
    ):
        multihop_queries.append(
            {
                "prompt": hop_out,
                "answers": [
                    {
                        "value": hop_value,
                        "aliases": []
                    }
                ],
                "hop_type": hop_type,
                "hop_id": f"{subject_id}_hop{i+1}"
            }
        )
    else:
        multihop_queries[-1]["answers"][0]["aliases"] = [entry["answer"]]

    entry["multihop_reasoning"] = [
        {
            "test_queries": multihop_queries
        }
    ]

    return entry

def encode_image(path):
    ext = os.path.splitext(path)[1].lower()

    # Figure out which image type it is
    if ext in [".jpg", ".jpeg"]:
        mime = "jpeg"
    elif ext == ".png":
        mime = "png"
    elif ext == ".webp":
        mime = "webp"
    elif ext == ".bmp":
        mime = "bmp"
    elif ext in [".tif", ".tiff"]:
        mime = "tiff"
    else:
        raise ValueError(f"Unsupported image format: {ext}")

    with open(path, "rb") as f:
        b64 = base64.b64encode(f.read()).decode("utf-8")

    return b64, mime

def run_single_query_inference(args):
    """
    args = (idx, cleaned_entry, img_b64, mime)
    """
    idx, cleaned, img_b64, mime = args
    try:
        query_output = infer_queries(
            cleaned,
            img_b64,
            mime,
            client,
            MODEL
        )

        return idx, query_output

    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 idx, {"error": str(e)}

def run_all_query_inference(items, max_workers=10):
    """
    items: list of (idx, cleaned_entry, img_b64, mime)
    """
    results = {}

    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        future_to_idx = {
            executor.submit(run_single_query_inference, item): item[0]
            for item in items
        }

        for future in as_completed(future_to_idx):
            idx = future_to_idx[future]
            try:
                entry_idx, out = future.result()
            except Exception as e:
                out = {"error": str(e)}
            results[idx] = out

    return results

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Script to automate multi-hop query generation from a .csv file. Columns should match round1_run.")
    parser.add_argument("--data", type=str, help="Location of .csv file")
    parser.add_argument("--images", type=str, help="Name of image directory")
    args = parser.parse_args()

    cur_dir = os.getcwd()

    imgpath = os.path.join(cur_dir, args.images)
    filepath = os.path.join(cur_dir, args.data)
    outpath = os.path.join(cur_dir, args.data.split('.')[0] + "_processed.json")
    process_file(filepath, outpath, imgpath)