dataset/datasets.py

"""Per-baseline sample loading for the baseline-intervention sweep.

Both datasets answer the same question at every baseline; only the image
changes. A `Sample` therefore carries one question/answer plus one image path
per baseline, and `Task` bundles the per-dataset prompt suffix, the candidate
answer set (used by likelihood scoring) and the answer normaliser.

VSR
    Questions come from V7_testing.json. Baseline 0.00 is the untouched
    original image; the other baselines are the GenStereo views recorded in
    each sample's "stereo" block.

Youtube_self_depth_QA
    Levels 1-3, each a jsonl of QA over a rendered visualisation. Every sample
    already stores one visualisation per baseline ("visualizations"), rendered
    with the annotation points warped into that view, so baseline 0.00 is the
    original view with the same annotation.
"""

import itertools
import json
import os
import re

from .baselines import BASELINES, by_baseline, canonical

REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))

VSR_TEST_JSON = os.path.join(REPO_ROOT, "02_data/VSR/V7_testing.json")
VSR_ORIGINAL_DIR = os.path.join(REPO_ROOT, "00_data/gen_output_VSR/original_images")
VSR_GENERATED_ROOT = os.path.join(REPO_ROOT, "00_data/gen_output_VSR/generate_image")
YOUTUBE_ROOT = os.path.join(REPO_ROOT, "02_data/Youtube_self_depth_QA")

IMAGE_EXTENSIONS = (".jpg", ".jpeg", ".png", ".webp", ".bmp")

LETTERS = ["A", "B", "C", "D"]
ORDERINGS = [" > ".join(p) for p in itertools.permutations(LETTERS)]


class Sample:
    """One question, one answer, one image per baseline."""

    __slots__ = ("sample_id", "question", "answer", "images", "meta")

    def __init__(self, sample_id, question, answer, images, meta=None):
        self.sample_id = sample_id
        self.question = question
        self.answer = answer
        self.images = images  # canonical baseline -> absolute image path
        self.meta = meta or {}

    def image(self, baseline):
        try:
            return self.images[canonical(baseline)]
        except KeyError:
            raise KeyError(
                f"Sample {self.sample_id} has no view for baseline {baseline}; "
                f"available: {sorted(self.images)}"
            ) from None


class Task:
    """A dataset split plus everything needed to prompt and score it."""

    def __init__(self, name, samples, answer_instruction, choices, normalize):
        self.name = name
        self.samples = samples
        self.answer_instruction = answer_instruction
        self.choices = choices
        self._normalize = normalize

    def normalize(self, text, sample=None):
        """Parse a model answer (or a reference) into a canonical form.

        `sample` is optional: the level-2 multiple-choice normaliser uses it to
        recognise a model that answered with the ordering instead of the option
        letter, and every other normaliser ignores it.
        """
        return self._normalize(text, sample)

    def prompt(self, sample):
        if self.answer_instruction:
            return f"{sample.question}\n{self.answer_instruction}"
        return sample.question


# --------------------------------------------------------------------------
# answer normalisation
# --------------------------------------------------------------------------

def normalize_true_false(text, sample=None):
    """First standalone true/false in the text, lowercased."""
    if not text:
        return ""
    match = re.search(r"\b(true|false)\b", text.lower())
    return match.group(1) if match else text.strip().lower()


def normalize_letter(text, sample=None):
    """First standalone A-D in the text, uppercased."""
    if not text:
        return ""
    match = re.search(r"\b([ABCD])\b", text.upper())
    return match.group(1) if match else text.strip().upper()


def normalize_ordering(text, sample=None):
    """The first four distinct A-D letters, rendered as "B > C > A > D".

    Accepts any separator the model happens to use (">", ",", "then", ...) and
    tolerates a preamble, so a correct ranking is not scored wrong purely on
    formatting.
    """
    if not text:
        return ""
    letters = []
    for letter in re.findall(r"\b([ABCD])\b", text.upper()):
        if letter not in letters:
            letters.append(letter)
        if len(letters) == 4:
            break
    if len(letters) != 4:
        return text.strip().upper()
    return " > ".join(letters)


ORDERING_AFTER_LETTER = re.compile(r"\b([ABCD])\b\s*>")


def normalize_choice_letter(text, sample=None):
    """Option letter for a multiple-choice answer.

    Level 2 labels its four OPTIONS A-D while the image labels its four POINTS
    A-D too, so a model may answer either with the option letter ("C") or by
    writing out the ordering that option holds ("B > C > A > D"). Only the
    former is the expected format, but the latter is unambiguous and is mapped
    back through the sample's own option table rather than scored wrong.

    The two are told apart by what follows the first standalone letter: a ">"
    means the letter opens an ordering, anything else ("C", "(C)", "C) B > C >
    A > D", "Answer: C") means it is the option label.
    """
    if not text:
        return ""
    match = re.search(r"\b([ABCD])\b", text.upper())
    options = (sample.meta.get("options") if sample is not None else None) or {}
    if match and options and ORDERING_AFTER_LETTER.match(text.upper(), match.start()):
        ordering = normalize_ordering(text)
        for label, option in options.items():
            if option == ordering:
                return label
    return normalize_letter(text)


# --------------------------------------------------------------------------
# VSR
# --------------------------------------------------------------------------

def _load_records(path):
    """Read either a JSON list or a jsonl file."""
    with open(path, "r", encoding="utf-8") as f:
        if path.endswith(".jsonl"):
            return [json.loads(line) for line in f if line.strip()]
        return json.load(f)


def _resolve_image(directory, stem):
    for ext in IMAGE_EXTENSIONS:
        candidate = os.path.join(directory, f"{stem}{ext}")
        if os.path.isfile(candidate):
            return candidate
    return os.path.join(directory, f"{stem}.jpg")


def load_vsr(test_json=VSR_TEST_JSON, original_dir=VSR_ORIGINAL_DIR,
             generated_root=VSR_GENERATED_ROOT, answer_instruction=None):
    samples = []
    for record in _load_records(test_json):
        sample_id = str(record["id"])
        question = record["conversations"][0]["value"].replace("<image>", "").strip()
        answer = next(
            (turn.get("value", "") for turn in record["conversations"] if turn.get("from") == "gpt"),
            "",
        )
        images = {"0.00": _resolve_image(original_dir, sample_id)}
        for baseline, entry in by_baseline(record.get("stereo", {})).items():
            images[baseline] = os.path.join(generated_root, entry["image"])
        samples.append(Sample(sample_id, question, answer, images))

    if answer_instruction is None:
        # The VSR question text already ends with its own instruction.
        answer_instruction = ""
    return Task("vsr", samples, answer_instruction, ["True", "False"], normalize_true_false)


# --------------------------------------------------------------------------
# Youtube_self_depth_QA
# --------------------------------------------------------------------------

LETTER_INSTRUCTION = "Answer with a single letter: A, B, C, or D."


def already_instructed(records):
    """True when the annotations write their own answer instruction into the question.

    Which levels do this has changed with every regeneration of the data (level 2
    started doing it when it became multiple choice, level 3 when it was rebuilt),
    so it is detected per file instead of hard-coded per level. Appending a second
    copy is not harmless: the duplicated sentence is exactly the kind of prompt
    noise that shifts a borderline answer, and it would differ between splits.
    """
    return all("Answer with" in record["question"] for record in records)


def load_youtube(level, root=YOUTUBE_ROOT, split="test", answer_instruction=None):
    """One Youtube level.

    Levels 1 and 3 ask for one of the four labelled points/trajectories, so the
    answer space is A-D and the prompt gets a format instruction appended.

    Level 2 comes in two formats. In the current one each record carries an
    `options` table and the question already embeds the four candidate
    orderings as (A)-(D) plus its own answer instruction, which makes it an A-D
    single choice like the other two levels. The older format asked the model
    to write the ordering out, so the answer space was the 24 permutations.
    Which one is in the file is detected from the records themselves rather
    than assumed, so the older annotations (and the `.ordering_bak` copies)
    still load correctly.
    """
    level = int(level)
    if level not in (1, 2, 3):
        raise ValueError(f"Youtube_self_depth_QA has levels 1-3, got {level}")
    path = os.path.join(root, f"level{level}", f"{split}.jsonl")
    records = _load_records(path)

    multiple_choice = level == 2 and all("options" in record for record in records)
    if level == 2 and not multiple_choice:
        # Legacy level 2: the model wrote the ordering out, so the answer space
        # is the 24 permutations rather than A-D.
        choices, normalize = ORDERINGS, normalize_ordering
        default_instruction = "Answer with the ranking only."
    else:
        choices = LETTERS
        normalize = normalize_choice_letter if multiple_choice else normalize_letter
        default_instruction = "" if already_instructed(records) else LETTER_INSTRUCTION

    samples = []
    for record in records:
        images = {
            baseline: os.path.join(root, relative)
            for baseline, relative in by_baseline(record["visualizations"]).items()
        }
        meta = {"task_type": record.get("task_type"), "level": record.get("level")}
        if record.get("options"):
            meta["options"] = record["options"]
            meta["answer_type"] = record.get("answer_type")
        samples.append(
            Sample(record["sample_id"], record["question"], record["answer"], images, meta)
        )

    if answer_instruction is None:
        answer_instruction = default_instruction
    return Task(f"youtube_level{level}", samples, answer_instruction, choices, normalize)


def load_task(dataset, level=None, **kwargs):
    """dataset in {vsr, youtube}; `level` selects the Youtube split."""
    if dataset == "vsr":
        return load_vsr(**kwargs)
    if dataset == "youtube":
        if level is None:
            raise ValueError("--level is required for the youtube dataset")
        return load_youtube(level, **kwargs)
    raise ValueError(f"Unknown dataset {dataset!r}; expected 'vsr' or 'youtube'")


def missing_images(task, baseline):
    return [s.sample_id for s in task.samples if not os.path.isfile(s.image(baseline))]


__all__ = [
    "BASELINES",
    "ORDERINGS",
    "Sample",
    "Task",
    "canonical",
    "load_task",
    "load_vsr",
    "load_youtube",
    "missing_images",
    "normalize_choice_letter",
    "normalize_letter",
    "normalize_ordering",
    "normalize_true_false",
]