export_text_rejections.py

import json
import csv
from openpyxl import Workbook
from openpyxl.utils import get_column_letter
from openpyxl.styles import Font

# Load JSON
with open("finalized_dataset.json") as f:
    data = json.load(f)
lookup = {r["prompt_hash"]: r for r in data}

# Load rejection IDs from CSV
ids = []
with open("Grained AI Rejection.xlsx - Sheet1.csv") as f:
    reader = csv.DictReader(f)
    for row in reader:
        ids.append(row["task ID"].strip())

# Filter: in CSV + found in JSON + no image (pure text)
found = [lookup[i] for i in ids if i in lookup]
text_only = [r for r in found if r.get("image") == "null"]
print(f"Total CSV IDs: {len(ids)}")
print(f"Found in JSON: {len(found)}")
print(f"Pure text (no image): {len(text_only)}")

# Build workbook
wb = Workbook()
ws = wb.active
ws.title = "Text Rejections"

# Headers - all fields, rollouts serialized as JSON string
headers = ["prompt_hash", "question", "answer", "image", "subject", "format", "avg_at_8", "rollouts"]
for col, header in enumerate(headers, 1):
    cell = ws.cell(row=1, column=col, value=header)
    cell.font = Font(bold=True)

# Write rows
for row_idx, record in enumerate(text_only, 2):
    for col_idx, field in enumerate(headers, 1):
        value = record.get(field)
        if field == "rollouts" and isinstance(value, (list, dict)):
            value = json.dumps(value, ensure_ascii=False)
        ws.cell(row=row_idx, column=col_idx, value=value)

# Auto-width (capped at 100)
for col in ws.columns:
    max_len = max((len(str(cell.value or "")) for cell in col), default=10)
    ws.column_dimensions[get_column_letter(col[0].column)].width = min(max_len + 2, 100)

output_path = "text_rejections_export.xlsx"
wb.save(output_path)
print(f"Saved: {output_path}")