prepare_input.py

#!/usr/bin/env python3
"""Prepare raw entity-pair records for DRU-RE-Yehia inference.

The relation model does not perform named-entity recognition. Callers supply an
Arabic sentence, the two mention strings, exact character spans, and canonical
coarse entity types. This utility applies the same deterministic marker,
ontology-option, Arabic-template, option-code, and prompt logic used to build
the official split.
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Any, Dict, List, Mapping, Tuple

from dotenv import load_dotenv

from re_sft_common import (
    PROMPT_VERSION,
    dump_jsonl,
    env_bool,
    env_int,
    env_str,
    load_jsonl,
    load_resources,
    transform_source_row,
)


ROOT = Path(__file__).resolve().parent
load_dotenv(ROOT / ".env")


def repository_path(value: str) -> Path:
    path = Path(value)
    return path if path.is_absolute() else ROOT / path


def unique_span(sentence: str, mention: str, role: str) -> Tuple[int, int]:
    starts: List[int] = []
    offset = 0
    while True:
        start = sentence.find(mention, offset)
        if start < 0:
            break
        starts.append(start)
        offset = start + 1
    if len(starts) != 1:
        raise ValueError(
            f"Cannot infer {role} span: mention occurs {len(starts)} times. "
            f"Supply explicit --{role}-start and --{role}-end offsets."
        )
    return starts[0], starts[0] + len(mention)


def resolved_span(
    sentence: str,
    mention: str,
    role: str,
    start: int | None,
    end: int | None,
) -> Tuple[int, int]:
    if start is None and end is None:
        return unique_span(sentence, mention, role)
    if start is None:
        raise ValueError(f"--{role}-start is required when --{role}-end is supplied")
    if end is None:
        end = start + len(mention)
    if not 0 <= start <= end <= len(sentence):
        raise ValueError(f"Invalid {role} span [{start}, {end}) for sentence length {len(sentence)}")
    observed = sentence[start:end]
    if observed != mention:
        raise ValueError(
            f"{role} span recovers {observed!r}, not the supplied mention {mention!r}"
        )
    return start, end


def direct_row(args: argparse.Namespace) -> Dict[str, Any]:
    required = {
        "sentence": args.sentence,
        "subject": args.subject,
        "object": args.object,
        "subject_type": args.subject_type,
        "object_type": args.object_type,
    }
    missing = [name for name, value in required.items() if value is None]
    if missing:
        raise ValueError(
            "Direct mode is missing: " + ", ".join("--" + name.replace("_", "-") for name in missing)
        )
    subject_start, subject_end = resolved_span(
        args.sentence,
        args.subject,
        "subject",
        args.subject_start,
        args.subject_end,
    )
    object_start, object_end = resolved_span(
        args.sentence,
        args.object,
        "object",
        args.object_start,
        args.object_end,
    )
    return {
        "id": args.id,
        "sentence_id": args.sentence_id or args.id,
        "triple_id": args.triple_id or args.id,
        "sentence": args.sentence,
        "subject": args.subject,
        "object": args.object,
        "subject_start": subject_start,
        "subject_end": subject_end,
        "object_start": object_start,
        "object_end": object_end,
        "subject_type": args.subject_type,
        "object_type": args.object_type,
        "relation": "",
    }


def source_rows(args: argparse.Namespace) -> List[Dict[str, Any]]:
    if args.input:
        rows = load_jsonl(repository_path(args.input))
        for index, row in enumerate(rows):
            row.setdefault("id", f"inference_{index}")
            row.setdefault("sentence_id", row["id"])
            row.setdefault("triple_id", row["id"])
            # Inference preparation must never propagate an accidental label.
            row["relation"] = ""
        return rows
    return [direct_row(args)]


def validate_prepared(row: Mapping[str, Any]) -> None:
    if row["prompt_version"] != PROMPT_VERSION:
        raise RuntimeError("Prepared prompt version mismatch")
    if len(row["messages"]) != 2 or row["messages"] != row["prompt_messages"]:
        raise RuntimeError("Prepared inference row unexpectedly contains an assistant message")
    for key in (
        "gold_relation_full",
        "gold_relation_ontology_id",
        "gold_answer_ar",
        "gold_option_index",
        "gold_answer_code",
    ):
        if row.get(key) is not None:
            raise RuntimeError(f"Prepared inference row unexpectedly contains {key}")
    options = row["allowed_options_ar"]
    if not options or options[-1] != "لا توجد علاقة":
        raise RuntimeError("Prepared no-relation option is not last")
    if len(options) != len(row["option_codes"]):
        raise RuntimeError("Prepared options and codes are misaligned")


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", help="Source-shaped JSONL; direct arguments are ignored")
    parser.add_argument("--output", default="predictions/prepared_inputs.jsonl")
    parser.add_argument("--list-types", action="store_true")
    parser.add_argument("--id", default="inference_0")
    parser.add_argument("--sentence-id")
    parser.add_argument("--triple-id")
    parser.add_argument("--sentence")
    parser.add_argument("--subject")
    parser.add_argument("--object")
    parser.add_argument("--subject-start", type=int)
    parser.add_argument("--subject-end", type=int)
    parser.add_argument("--object-start", type=int)
    parser.add_argument("--object-end", type=int)
    parser.add_argument("--subject-type")
    parser.add_argument("--object-type")
    args = parser.parse_args()

    resource_dir = repository_path(env_str("RESOURCE_DIR", "resources"))
    resources = load_resources(resource_dir)
    if args.list_types:
        print(json.dumps(resources["type_to_ar"], ensure_ascii=False, indent=2))
        return

    seed = env_int("SEED", 42)
    context_chars = env_int("PROMPT_CONTEXT_CHARS", 500)
    shuffle_options = env_bool("SHUFFLE_OPTIONS", True)
    prepared: List[Dict[str, Any]] = []
    for row in source_rows(args):
        transformed = transform_source_row(
            row=row,
            split_name="official",
            resources=resources,
            seed=seed,
            context_chars=context_chars,
            shuffle_options=shuffle_options,
        )
        validate_prepared(transformed)
        prepared.append(transformed)

    output = repository_path(args.output)
    dump_jsonl(prepared, output)
    print(
        json.dumps(
            {
                "status": "passed",
                "rows": len(prepared),
                "output": str(output),
                "prompt_version": PROMPT_VERSION,
                "labels_present": False,
                "option_counts": [len(row["allowed_options_ar"]) for row in prepared],
            },
            ensure_ascii=False,
            indent=2,
        )
    )


if __name__ == "__main__":
    main()