#!/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()