test_codabench/runner/run_test.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""Run the complete span-free DRU-RE-Yehia pipeline on /workspace/test.jsonl.

For every test row, the script:
  1. finds every subject and object occurrence;
  2. predicts one coarse type for every unique occurrence;
  3. forms every directed subject-occurrence × object-occurrence candidate;
  4. delegates prompt construction to the released DRU-RE-Yehia prepare_input.py;
  5. scores only the row-local one-token Arabic option codes;
  6. applies the released frozen no-relation logit bias;
  7. aggregates candidate decisions;
  8. creates a validated /workspace/submission.zip.

The released prompt, ontology, Arabic templates, option shuffle, code inventory,
adapter, tokenizer, and inference bias are downloaded from the public
U4RASD/DRU-RE-Yehia repository.  The gated Yehia base remains pinned to the
revision stated in the released inference_config.json.
"""

from __future__ import annotations

import argparse
import gc
import hashlib
import importlib.util
import json
import math
import os
import re
import shutil
import subprocess
import sys
import time
import unicodedata
import zipfile
from collections import Counter, defaultdict
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any, Dict, Iterable, Iterator, List, Mapping, Optional, Sequence, Tuple

os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0"
# vLLM 0.10.2's V1 sampler has two LoRA shape bugs for this model: its
# allowed-token mask uses tokenizer.vocab_size instead of the padded LM-head
# size, and its prompt-logprob path misaligns LoRA token mappings.  The mature
# V0 engine supports the same offline LoRA API without those V1-only paths.
os.environ.setdefault("VLLM_USE_V1", "0")

import torch
from dotenv import load_dotenv
from huggingface_hub import HfApi, snapshot_download
from tqdm.auto import tqdm
from transformers import AutoTokenizer
from vllm import LLM, SamplingParams, __version__ as vllm_version
from vllm.lora.request import LoRARequest

from type_predictor import TypePredictorEngine


Span = Tuple[int, int]
ROOT = Path(__file__).resolve().parent
WORKSPACE = ROOT.parent

DEFAULT_BASE_MODEL_ID = "Navid-AI/Yehia-7B-preview"
DEFAULT_BASE_MODEL_REVISION = "b9dda4715eafee7e8090d2c83cfe078d75f4ebb8"
DEFAULT_ADAPTER_ID = "U4RASD/DRU-RE-Yehia"
DEFAULT_TYPE_MODEL_ID = "U4RASD/TypePredictor"
NO_RELATION = "no_relation"
NO_RELATION_SUBMISSION = "no-relation"
ZIP_TIMESTAMP = (1980, 1, 1, 0, 0, 0)


# ---------------------------------------------------------------------------
# Generic helpers
# ---------------------------------------------------------------------------


def env_bool(name: str, default: bool) -> bool:
    value = os.getenv(name)
    if value is None or not value.strip():
        return default
    normalized = value.strip().lower()
    if normalized in {"1", "true", "yes", "y", "on"}:
        return True
    if normalized in {"0", "false", "no", "n", "off"}:
        return False
    raise ValueError(f"{name} must be a boolean, got {value!r}")


def env_int(name: str, default: int) -> int:
    value = os.getenv(name)
    return default if value is None or not value.strip() else int(value)


def env_float(name: str, default: float) -> float:
    value = os.getenv(name)
    return default if value is None or not value.strip() else float(value)


def batched(items: Sequence[Any], batch_size: int) -> Iterator[Sequence[Any]]:
    if batch_size <= 0:
        raise ValueError("batch_size must be positive")
    for start in range(0, len(items), batch_size):
        yield items[start : start + batch_size]


def read_jsonl(path: Path) -> List[Dict[str, Any]]:
    rows: List[Dict[str, Any]] = []
    with path.open("r", encoding="utf-8-sig") as handle:
        for line_number, line in enumerate(handle, start=1):
            if not line.strip():
                continue
            try:
                value = json.loads(line)
            except json.JSONDecodeError as exc:
                raise ValueError(f"Invalid JSON at {path}:{line_number}: {exc}") from exc
            if not isinstance(value, dict):
                raise ValueError(f"Expected a JSON object at {path}:{line_number}")
            rows.append(value)
    if not rows:
        raise ValueError(f"No rows found in {path}")
    return rows


def write_jsonl(path: Path, rows: Iterable[Mapping[str, Any]]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8", newline="\n") as handle:
        for row in rows:
            handle.write(json.dumps(dict(row), ensure_ascii=False, separators=(",", ":")) + "\n")


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def row_field(row: Mapping[str, Any], names: Sequence[str], *, required: bool = True) -> Any:
    for name in names:
        if name in row:
            return row[name]
    if required:
        raise KeyError(f"Missing required field. Tried: {', '.join(names)}")
    return None


def scalar_text(value: Any, field_name: str) -> str:
    if isinstance(value, str):
        return value
    if isinstance(value, Mapping):
        for key in ("text", "mention", "surface", "word", "name", "value", "entity"):
            nested = value.get(key)
            if isinstance(nested, str):
                return nested
    raise ValueError(f"Could not read {field_name} as text: {value!r}")


def optional_token(*names: str) -> Optional[str]:
    for name in names:
        value = os.getenv(name, "").strip()
        if value and not value.startswith("hf_your_"):
            return value
    return None


def require_token(name: str) -> str:
    value = os.getenv(name, "").strip()
    if not value or value.startswith("hf_your_"):
        raise RuntimeError(
            f"{name} is required. It must belong to an account with access to the gated Yehia base."
        )
    return value


# ---------------------------------------------------------------------------
# Occurrence finding
# ---------------------------------------------------------------------------


ARABIC_CHAR_MAP = str.maketrans(
    {
        "أ": "ا",
        "إ": "ا",
        "آ": "ا",
        "ٱ": "ا",
        "ى": "ي",
    }
)
SURROUNDING_CHARS = " \t\r\n\ufeff\u200f\u200e\"'“”‘’`()[]{}<>،,.;:؛!?؟"


def exact_overlapping_occurrences(sentence: str, mention: str) -> List[Span]:
    if not mention:
        return []
    spans: List[Span] = []
    cursor = 0
    while True:
        found = sentence.find(mention, cursor)
        if found < 0:
            break
        spans.append((found, found + len(mention)))
        cursor = found + 1
    return spans


def normalize_with_offset_map(text: str) -> Tuple[str, List[int], List[int]]:
    chars: List[str] = []
    starts: List[int] = []
    ends: List[int] = []
    previous_was_space = False

    for original_index, original_char in enumerate(text):
        expanded = unicodedata.normalize("NFKC", original_char)
        for char in expanded:
            if char == "ـ" or unicodedata.combining(char):
                continue
            char = char.translate(ARABIC_CHAR_MAP)
            if char.isspace():
                if previous_was_space and chars:
                    ends[-1] = original_index + 1
                    continue
                chars.append(" ")
                starts.append(original_index)
                ends.append(original_index + 1)
                previous_was_space = True
            else:
                chars.append(char)
                starts.append(original_index)
                ends.append(original_index + 1)
                previous_was_space = False

    return "".join(chars), starts, ends


def robust_occurrences(sentence: str, mention: str) -> Tuple[List[Span], str, str]:
    variants: List[Tuple[str, str]] = [(mention, "exact")]
    stripped = mention.strip()
    if stripped and stripped != mention:
        variants.append((stripped, "stripped"))
    edge_stripped = mention.strip(SURROUNDING_CHARS)
    if edge_stripped and edge_stripped not in {text for text, _ in variants}:
        variants.append((edge_stripped, "edge_stripped"))

    for variant, mode in variants:
        spans = exact_overlapping_occurrences(sentence, variant)
        if spans:
            return spans, mode, variant

    normalized_sentence, start_map, end_map = normalize_with_offset_map(sentence)
    for variant, _ in variants:
        normalized_mention, _, _ = normalize_with_offset_map(variant)
        normalized_mention = normalized_mention.strip()
        if not normalized_mention:
            continue
        normalized_spans = exact_overlapping_occurrences(
            normalized_sentence, normalized_mention
        )
        if normalized_spans:
            original_spans = [
                (start_map[start], end_map[end - 1])
                for start, end in normalized_spans
            ]
            return list(dict.fromkeys(original_spans)), "arabic_normalized", variant

    return [], "not_found", mention


# ---------------------------------------------------------------------------
# Records
# ---------------------------------------------------------------------------


@dataclass
class InputRow:
    row_index: int
    triple_id: str
    sentence_id: str
    sentence: str
    subject: str
    object: str
    original: Dict[str, Any]


@dataclass
class Candidate:
    uid: str
    row_index: int
    candidate_index: int
    source_triple_id: str
    subject_span: Span
    object_span: Span
    subject_text: str
    object_text: str
    subject_type: str = ""
    object_type: str = ""
    subject_type_confidence: Optional[float] = None
    object_type_confidence: Optional[float] = None
    subject_match_mode: str = ""
    object_match_mode: str = ""
    predicted_relation: Optional[str] = None
    predicted_option_index: Optional[int] = None
    predicted_code: Optional[str] = None
    predicted_option_ar: Optional[str] = None
    predicted_probability: Optional[float] = None
    no_relation_probability: Optional[float] = None
    positive_probability: Optional[float] = None
    score_margin: Optional[float] = None
    option_relations: Optional[List[str]] = None
    option_codes: Optional[List[str]] = None
    option_probabilities: Optional[List[float]] = None
    option_adjusted_logits: Optional[List[float]] = None


# ---------------------------------------------------------------------------
# Download and load the exact released Yehia resources
# ---------------------------------------------------------------------------


def resolve_adapter_snapshot(output_dir: Path) -> Tuple[Path, str]:
    model_id = os.getenv("YEHIA_ADAPTER_MODEL_ID", DEFAULT_ADAPTER_ID)
    requested_revision = os.getenv("YEHIA_ADAPTER_REVISION", "main")
    token = optional_token("HF_TOKENTWO", "HF_TOKEN", "HUGGINGFACE_TOKEN")
    api = HfApi(token=token)
    resolved_revision = api.model_info(model_id, revision=requested_revision).sha

    local_dir = Path(
        os.getenv("LOCAL_YEHIA_ADAPTER_DIR", str(WORKSPACE / "models/DRU-RE-Yehia"))
    ).expanduser()
    local_dir.mkdir(parents=True, exist_ok=True)

    required = [
        "adapter_config.json",
        "adapter_model.safetensors",
        "prepare_input.py",
        "re_sft_common.py",
        "inference_config.json",
        "resources/relation_mapper.json",
    ]
    marker = local_dir / ".dru_adapter_revision"
    reusable = (
        all((local_dir / relative).is_file() for relative in required)
        and marker.is_file()
        and marker.read_text(encoding="utf-8").strip() == resolved_revision
    )

    if not reusable:
        print(f"Downloading {model_id}@{resolved_revision} to {local_dir}")
        snapshot_download(
            repo_id=model_id,
            revision=resolved_revision,
            repo_type="model",
            token=token,
            local_dir=str(local_dir),
            allow_patterns=[
                "adapter_config.json",
                "adapter_model.safetensors",
                "inference_config.json",
                "tokenizer*",
                "special_tokens_map.json",
                "added_tokens.json",
                "chat_template*",
                "*.model",
                "prepare_input.py",
                "re_sft_common.py",
                "resources/*",
            ],
        )
        marker.write_text(resolved_revision + "\n", encoding="utf-8")

    for relative in required:
        if not (local_dir / relative).is_file():
            raise FileNotFoundError(f"Released Yehia snapshot is missing {relative}")

    (output_dir / "adapter_revision.txt").write_text(
        resolved_revision + "\n", encoding="utf-8"
    )
    return local_dir, resolved_revision


def load_release_common(adapter_dir: Path):
    module_name = "dru_re_yehia_release_common"
    spec = importlib.util.spec_from_file_location(
        module_name, adapter_dir / "re_sft_common.py"
    )
    if spec is None or spec.loader is None:
        raise RuntimeError("Could not import the released re_sft_common.py")
    module = importlib.util.module_from_spec(spec)
    sys.modules[module_name] = module
    spec.loader.exec_module(module)
    return module


def selected_inference_config(adapter_dir: Path) -> Dict[str, Any]:
    config = json.loads(
        (adapter_dir / "inference_config.json").read_text(encoding="utf-8")
    )
    required = {
        "base_model_id",
        "base_model_revision",
        "prompt_version",
        "no_relation_logit_bias",
    }
    missing = sorted(required - config.keys())
    if missing:
        raise RuntimeError(f"inference_config.json is missing: {missing}")
    return config


def resolve_base_model(
    model_id: str,
    revision: str,
) -> Path:
    local_dir = Path(
        os.getenv("LOCAL_YEHIA_MODEL_DIR", str(WORKSPACE / "models/Yehia-7B-preview"))
    ).expanduser()
    marker = local_dir / ".dru_hf_revision"

    reusable = (
        (local_dir / "config.json").is_file()
        and marker.is_file()
        and marker.read_text(encoding="utf-8").strip() == revision
    )
    if reusable:
        return local_dir

    if env_bool("HF_HUB_OFFLINE", False):
        raise RuntimeError(
            f"Yehia base is absent or unverified at {local_dir}, while HF_HUB_OFFLINE is enabled."
        )

    token = require_token("HF_TOKENONE")
    local_dir.mkdir(parents=True, exist_ok=True)
    print(f"Downloading gated base {model_id}@{revision} to {local_dir}")
    snapshot_download(
        repo_id=model_id,
        revision=revision,
        repo_type="model",
        token=token,
        local_dir=str(local_dir),
    )
    if not (local_dir / "config.json").is_file():
        raise RuntimeError(f"Incomplete Yehia base snapshot at {local_dir}")
    marker.write_text(revision + "\n", encoding="utf-8")
    return local_dir


def prepare_candidates_with_release(
    adapter_dir: Path,
    raw_path: Path,
    prepared_path: Path,
) -> None:
    environment = os.environ.copy()
    environment["RESOURCE_DIR"] = str(adapter_dir / "resources")
    environment.setdefault("SEED", "42")
    environment.setdefault("PROMPT_CONTEXT_CHARS", "500")
    environment.setdefault("SHUFFLE_OPTIONS", "true")

    command = [
        sys.executable,
        str(adapter_dir / "prepare_input.py"),
        "--input",
        str(raw_path),
        "--output",
        str(prepared_path),
    ]
    print("Preparing exact released prompts:", " ".join(command))
    subprocess.run(
        command,
        check=True,
        cwd=str(adapter_dir),
        env=environment,
    )
    if not prepared_path.is_file():
        raise RuntimeError("Released prepare_input.py did not create its output")


# ---------------------------------------------------------------------------
# Constrained Yehia inference
# ---------------------------------------------------------------------------


def run_yehia_candidates(
    adapter_dir: Path,
    prepared_rows: List[Dict[str, Any]],
    candidate_by_uid: Mapping[str, Candidate],
    common: Any,
    inference_config: Mapping[str, Any],
    batch_size: int,
) -> Dict[str, Any]:
    base_id = str(inference_config["base_model_id"])
    base_revision = str(inference_config["base_model_revision"])
    bias = float(inference_config["no_relation_logit_bias"])
    base_dir = resolve_base_model(base_id, base_revision)

    tokenizer_source = (
        adapter_dir if (adapter_dir / "tokenizer_config.json").is_file() else base_dir
    )
    tokenizer = AutoTokenizer.from_pretrained(
        str(tokenizer_source),
        local_files_only=True,
        use_fast=True,
    )
    if not tokenizer.chat_template:
        raise RuntimeError("Yehia tokenizer has no native chat template")
    if tokenizer.pad_token_id is None:
        tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "right"

    option_codes = tuple(common.OPTION_CODES)
    code_token_ids: List[int] = []
    for code in option_codes:
        ids = tokenizer.encode(" " + code, add_special_tokens=False)
        if len(ids) != 1:
            raise RuntimeError(f"Decision code {code!r} is not one token: {ids}")
        code_token_ids.append(int(ids[0]))

    # vLLM performs the causal-LM inference.  Transformers is retained only for
    # deterministic tokenizer/chat-template handling and for TypePredictor,
    # whose custom encoder is not a generative model supported by vLLM.
    sequences = [
        list(
            tokenizer.apply_chat_template(
                row["prompt_messages"],
                tokenize=True,
                add_generation_prompt=True,
            )
        )
        for row in prepared_rows
    ]
    max_observed_tokens = max((len(sequence) for sequence in sequences), default=0)
    configured_max_model_len = env_int("VLLM_MAX_MODEL_LEN", 0)
    max_model_len = configured_max_model_len or max_observed_tokens + 8
    if max_model_len <= max_observed_tokens:
        raise ValueError(
            f"VLLM_MAX_MODEL_LEN={max_model_len} must exceed the longest prompt "
            f"({max_observed_tokens} tokens)"
        )

    adapter_config = json.loads(
        (adapter_dir / "adapter_config.json").read_text(encoding="utf-8")
    )
    adapter_rank = int(adapter_config.get("r", 64))
    supported_lora_ranks = (1, 8, 16, 32, 64, 128, 256, 320, 512)
    max_lora_rank = next(
        (rank for rank in supported_lora_ranks if rank >= adapter_rank),
        adapter_rank,
    )
    llm = LLM(
        model=str(base_dir),
        tokenizer=str(tokenizer_source),
        dtype="bfloat16",
        trust_remote_code=False,
        enable_lora=True,
        max_lora_rank=max_lora_rank,
        max_model_len=max_model_len,
        max_num_seqs=env_int("VLLM_MAX_NUM_SEQS", 256),
        gpu_memory_utilization=env_float("VLLM_GPU_MEMORY_UTILIZATION", 0.90),
        enable_prefix_caching=env_bool("VLLM_ENABLE_PREFIX_CACHING", True),
    )
    lora_request = LoRARequest("dru_re_yehia", 1, str(adapter_dir))

    output_records: List[Dict[str, Any]] = []
    request_batch_size = env_int("VLLM_REQUEST_BATCH_SIZE", max(batch_size, 512))
    indexed = list(zip(prepared_rows, sequences))
    for group in tqdm(
        list(batched(indexed, request_batch_size)),
        desc="vLLM constrained candidate inference",
    ):
        prompts: List[Dict[str, List[int]]] = []
        sampling_params: List[SamplingParams] = []
        for row, sequence in group:
            option_count = len(row["option_codes"])
            allowed_ids = code_token_ids[:option_count]
            prompts.append({"prompt_token_ids": sequence})
            sampling_params.append(
                SamplingParams(
                    temperature=0.0,
                    max_tokens=1,
                    detokenize=False,
                    allowed_token_ids=allowed_ids,
                    logit_bias={allowed_ids[-1]: bias},
                    logprobs=len(allowed_ids),
                )
            )
        results = llm.generate(
            prompts,
            sampling_params,
            lora_request=lora_request,
            use_tqdm=False,
        )
        if len(results) != len(group):
            raise RuntimeError(
                f"vLLM returned {len(results)} results for {len(group)} prompts"
            )
        for (row, _), result in zip(group, results):
                options = list(row["allowed_options_ar"])
                relations = list(row["allowed_relation_full_labels"])
                codes = list(row["option_codes"])
                if not options or options[-1] != "لا توجد علاقة":
                    raise RuntimeError(f"Malformed option list for candidate {row.get('id')}")
                if relations[-1] != NO_RELATION:
                    raise RuntimeError(f"no_relation is not last for candidate {row.get('id')}")
                if not (len(options) == len(relations) == len(codes)):
                    raise RuntimeError(f"Unaligned row-local options for candidate {row.get('id')}")
                if codes != list(option_codes[: len(codes)]):
                    raise RuntimeError(
                        f"Unexpected code inventory for candidate {row.get('id')}: {codes}"
                    )

                candidate_ids = code_token_ids[: len(options)]
                if not result.outputs or not result.outputs[0].token_ids:
                    raise RuntimeError(f"vLLM produced no decision for {row.get('id')}")
                generated = result.outputs[0]
                chosen_token_id = int(generated.token_ids[0])
                if chosen_token_id not in candidate_ids:
                    raise RuntimeError(
                        f"vLLM chose disallowed token {chosen_token_id} for {row.get('id')}"
                    )
                chosen = candidate_ids.index(chosen_token_id)
                if not generated.logprobs or generated.logprobs[0] is None:
                    raise RuntimeError(f"vLLM returned no decision logprobs for {row.get('id')}")
                token_logprobs = generated.logprobs[0]
                missing_ids = [token_id for token_id in candidate_ids if token_id not in token_logprobs]
                if missing_ids:
                    raise RuntimeError(
                        f"vLLM omitted requested logprobs {missing_ids} for {row.get('id')}"
                    )
                adjusted_scores = [
                    float(token_logprobs[token_id].logprob) for token_id in candidate_ids
                ]
                probabilities = [math.exp(value) for value in adjusted_scores]
                total_probability = sum(probabilities)
                probabilities = [value / total_probability for value in probabilities]
                sorted_scores = sorted(adjusted_scores, reverse=True)
                margin = (
                    float(sorted_scores[0] - sorted_scores[1])
                    if len(sorted_scores) > 1
                    else float("inf")
                )

                uid = str(row.get("id"))
                if uid not in candidate_by_uid:
                    raise KeyError(f"Prepared candidate ID is unknown: {uid}")
                candidate = candidate_by_uid[uid]
                candidate.predicted_relation = relations[chosen]
                candidate.predicted_option_index = chosen
                candidate.predicted_code = codes[chosen]
                candidate.predicted_option_ar = options[chosen]
                candidate.predicted_probability = probabilities[chosen]
                candidate.no_relation_probability = probabilities[-1]
                candidate.positive_probability = 1.0 - probabilities[-1]
                candidate.score_margin = margin
                candidate.option_relations = relations
                candidate.option_codes = codes
                candidate.option_probabilities = [
                    float(value) for value in probabilities
                ]
                candidate.option_adjusted_logits = [
                    float(value) for value in adjusted_scores
                ]

                output_records.append(
                    {
                        "id": uid,
                        "source_triple_id": candidate.source_triple_id,
                        "row_index": candidate.row_index,
                        "candidate_index": candidate.candidate_index,
                        "predicted_relation_full": relations[chosen],
                        "predicted_option_index": chosen,
                        "predicted_code": codes[chosen],
                        "predicted_option_ar": options[chosen],
                        "predicted_probability": candidate.predicted_probability,
                        "no_relation_probability": candidate.no_relation_probability,
                        "positive_probability": candidate.positive_probability,
                        "score_margin": margin,
                        "allowed_relation_full_labels": relations,
                        "option_codes": codes,
                        "option_probabilities": candidate.option_probabilities,
                        "option_adjusted_logits": candidate.option_adjusted_logits,
                        "no_relation_logit_bias": bias,
                    }
                )

    del llm
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    return {
        "records": output_records,
        "no_relation_logit_bias": bias,
        "base_model_id": base_id,
        "base_model_revision": base_revision,
        "max_prompt_tokens": max_observed_tokens,
        "inference_backend": "vllm",
        "vllm_version": vllm_version,
    }


# ---------------------------------------------------------------------------
# Candidate aggregation
# ---------------------------------------------------------------------------


def canonical_relation_order(adapter_dir: Path) -> List[str]:
    mapper = json.loads(
        (adapter_dir / "resources" / "relation_mapper.json").read_text(encoding="utf-8")
    )
    labels = [str(label) for label in mapper.keys()]
    return [label for label in labels if label != NO_RELATION]


def probability_for(candidate: Candidate, relation: str) -> float:
    if not candidate.option_relations or not candidate.option_probabilities:
        return 0.0
    try:
        index = candidate.option_relations.index(relation)
    except ValueError:
        return 0.0
    return float(candidate.option_probabilities[index])


def aggregate_majority_vote(
    candidates: Sequence[Candidate],
    relation_order: Sequence[str],
    positive_fraction: float,
) -> Tuple[str, Dict[str, Any]]:
    positive = [
        candidate
        for candidate in candidates
        if candidate.predicted_relation not in {None, NO_RELATION}
    ]
    required = max(1, math.ceil(positive_fraction * len(candidates) - 1e-12))
    gate_passed = len(positive) >= required
    if not gate_passed:
        return NO_RELATION, {
            "gate": "candidate_positive_fraction",
            "positive_candidates": len(positive),
            "required_positive_candidates": required,
            "gate_passed": False,
        }

    votes = Counter(str(candidate.predicted_relation) for candidate in positive)
    max_votes = max(votes.values())
    tied = [label for label, count in votes.items() if count == max_votes]

    summed_support = {
        label: sum(probability_for(candidate, label) for candidate in candidates)
        for label in tied
    }
    best_support = max(summed_support.values())
    tied = [
        label for label in tied
        if math.isclose(summed_support[label], best_support, rel_tol=1e-12, abs_tol=1e-12)
    ]

    best_margin = {
        label: max(
            (
                float(candidate.score_margin or 0.0)
                for candidate in positive
                if candidate.predicted_relation == label
            ),
            default=0.0,
        )
        for label in tied
    }
    max_margin = max(best_margin.values())
    tied = [
        label for label in tied
        if math.isclose(best_margin[label], max_margin, rel_tol=1e-12, abs_tol=1e-12)
    ]

    order_index = {label: index for index, label in enumerate(relation_order)}
    winner = min(tied, key=lambda label: order_index.get(label, 10**9))
    return winner, {
        "gate": "candidate_positive_fraction",
        "positive_candidates": len(positive),
        "required_positive_candidates": required,
        "gate_passed": True,
        "votes": dict(votes),
        "tied_after_votes": [label for label, count in votes.items() if count == max_votes],
        "summed_probability_support": summed_support,
        "winning_margin": best_margin.get(winner),
    }


def aggregate_soft_pool(
    candidates: Sequence[Candidate],
    relation_order: Sequence[str],
    positive_threshold: float,
) -> Tuple[str, Dict[str, Any]]:
    if not candidates:
        return NO_RELATION, {"gate_passed": False, "fallback": "no_candidates"}

    mean_no = sum(float(candidate.no_relation_probability or 0.0) for candidate in candidates) / len(candidates)
    mean_positive = 1.0 - mean_no
    positive_scores = {
        relation: sum(probability_for(candidate, relation) for candidate in candidates)
        / len(candidates)
        for relation in relation_order
    }
    if mean_positive < positive_threshold or not positive_scores:
        return NO_RELATION, {
            "gate": "mean_positive_probability",
            "mean_positive_probability": mean_positive,
            "positive_threshold": positive_threshold,
            "gate_passed": False,
        }

    winner = max(
        relation_order,
        key=lambda relation: (positive_scores[relation], -relation_order.index(relation)),
    )
    return winner, {
        "gate": "mean_positive_probability",
        "mean_positive_probability": mean_positive,
        "positive_threshold": positive_threshold,
        "gate_passed": True,
        "winning_mean_probability": positive_scores[winner],
    }


def aggregate_max_positive(
    candidates: Sequence[Candidate],
    relation_order: Sequence[str],
) -> Tuple[str, Dict[str, Any]]:
    positive = [
        candidate
        for candidate in candidates
        if candidate.predicted_relation not in {None, NO_RELATION}
    ]
    if not positive:
        return NO_RELATION, {
            "gate": "any_candidate_positive",
            "gate_passed": False,
        }
    winner = max(
        positive,
        key=lambda candidate: (
            float(candidate.predicted_probability or 0.0),
            float(candidate.score_margin or 0.0),
            -candidate.candidate_index,
        ),
    )
    return str(winner.predicted_relation), {
        "gate": "any_candidate_positive",
        "gate_passed": True,
        "winning_candidate_index": winner.candidate_index,
        "winning_candidate_probability": winner.predicted_probability,
        "winning_candidate_margin": winner.score_margin,
    }


def aggregate_top_confidence(
    candidates: Sequence[Candidate],
    relation_order: Sequence[str],
) -> Tuple[str, Dict[str, Any]]:
    if not candidates:
        return NO_RELATION, {
            "gate": "top_candidate_confidence",
            "gate_passed": False,
        }
    winner = max(
        candidates,
        key=lambda candidate: (
            float(candidate.predicted_probability or 0.0),
            float(candidate.score_margin or 0.0),
            -candidate.candidate_index,
        ),
    )
    return str(winner.predicted_relation), {
        "gate": "top_candidate_confidence",
        "gate_passed": winner.predicted_relation != NO_RELATION,
        "winning_candidate_index": winner.candidate_index,
        "winning_candidate_probability": winner.predicted_probability,
        "winning_candidate_margin": winner.score_margin,
    }


AGGREGATORS = {
    "majority_vote": aggregate_majority_vote,
    "soft_pool": aggregate_soft_pool,
    "max_positive": aggregate_max_positive,
    "top_confidence": aggregate_top_confidence,
}


# ---------------------------------------------------------------------------
# Submission
# ---------------------------------------------------------------------------


def validate_submission(
    zip_path: Path,
    expected_ids: Sequence[str],
    legal_relations: Sequence[str],
) -> None:
    with zipfile.ZipFile(zip_path, "r") as archive:
        if archive.namelist() != ["predictions.txt"]:
            raise RuntimeError(
                f"{zip_path} must contain only predictions.txt at root; "
                f"found {archive.namelist()}"
            )
        content = archive.read("predictions.txt").decode("utf-8")

    lines = content.splitlines()
    if len(lines) != len(expected_ids):
        raise RuntimeError(
            f"Expected {len(expected_ids)} predictions, found {len(lines)}"
        )

    allowed = set(legal_relations) | {NO_RELATION_SUBMISSION}
    observed_ids: List[str] = []
    for line_number, line in enumerate(lines, start=1):
        pieces = line.split("\t")
        if len(pieces) != 2:
            raise RuntimeError(f"Malformed predictions line {line_number}: {line!r}")
        triple_id, relation = pieces
        if relation not in allowed:
            raise RuntimeError(
                f"Illegal relation {relation!r} at predictions line {line_number}"
            )
        observed_ids.append(triple_id)

    if observed_ids != list(expected_ids):
        raise RuntimeError("Submission IDs/order do not exactly match test.jsonl")
    if len(set(observed_ids)) != len(observed_ids):
        raise RuntimeError("Submission contains duplicate triple IDs")


def write_submission(
    output_dir: Path,
    method: str,
    triple_ids: Sequence[str],
    labels: Sequence[str],
    legal_relations: Sequence[str],
) -> Tuple[Path, Path]:
    if len(triple_ids) != len(labels):
        raise ValueError("ID and prediction counts differ")

    text_path = output_dir / f"predictions_{method}.txt"
    zip_path = output_dir / f"submission_{method}.zip"
    with text_path.open("w", encoding="utf-8", newline="\n") as handle:
        for triple_id, label in zip(triple_ids, labels):
            output_label = NO_RELATION_SUBMISSION if label == NO_RELATION else label
            handle.write(f"{triple_id}\t{output_label}\n")

    zip_info = zipfile.ZipInfo("predictions.txt", date_time=ZIP_TIMESTAMP)
    zip_info.compress_type = zipfile.ZIP_DEFLATED
    zip_info.external_attr = 0o644 << 16
    with zipfile.ZipFile(zip_path, "w") as archive:
        archive.writestr(zip_info, text_path.read_bytes())

    validate_submission(zip_path, triple_ids, legal_relations)
    return text_path, zip_path


# ---------------------------------------------------------------------------
# CLI and main
# ---------------------------------------------------------------------------


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--input",
        type=Path,
        default=Path(os.getenv("TEST_JSONL", str(WORKSPACE / "test.jsonl"))),
    )
    parser.add_argument(
        "--output-dir",
        type=Path,
        default=Path(os.getenv("OUTPUT_DIR", str(WORKSPACE / "yehia_test_output"))),
    )
    parser.add_argument(
        "--aggregation",
        choices=sorted(AGGREGATORS),
        default=os.getenv("AGGREGATION_METHOD", "majority_vote"),
    )
    parser.add_argument(
        "--validate-input-only",
        action="store_true",
        help="Validate schema and enumerate candidate spans without loading any model.",
    )
    return parser.parse_args()


def main() -> None:
    load_dotenv(ROOT / ".env", override=False)
    args = parse_args()
    input_path = args.input.expanduser().resolve()
    output_dir = args.output_dir.expanduser().resolve()
    output_dir.mkdir(parents=True, exist_ok=True)

    if not input_path.is_file():
        raise FileNotFoundError(
            f"Test file not found: {input_path}\n"
            "Put the blind test at /workspace/test.jsonl or pass --input."
        )

    raw_rows = read_jsonl(input_path)
    rows: List[InputRow] = []
    seen_ids: set[str] = set()
    for row_index, raw in enumerate(raw_rows):
        triple_id = str(row_field(raw, ("triple_id", "tripleId", "id")))
        if not triple_id:
            raise ValueError(f"Row {row_index} has an empty triple_id")
        if triple_id in seen_ids:
            raise ValueError(f"Duplicate triple_id: {triple_id}")
        seen_ids.add(triple_id)

        sentence = scalar_text(row_field(raw, ("sentence", "text")), "sentence")
        subject = scalar_text(
            row_field(raw, ("subject", "first_entity", "entity1")), "subject"
        )
        object_ = scalar_text(
            row_field(raw, ("object", "second_entity", "entity2")), "object"
        )
        sentence_id_value = row_field(
            raw, ("sentence_id", "sentenceId"), required=False
        )
        sentence_id = str(sentence_id_value if sentence_id_value is not None else triple_id)
        rows.append(
            InputRow(
                row_index=row_index,
                triple_id=triple_id,
                sentence_id=sentence_id,
                sentence=sentence,
                subject=subject,
                object=object_,
                original=dict(raw),
            )
        )

    print(f"Loaded {len(rows)} rows from {input_path}")
    print("Input keys:", sorted(set().union(*(row.original.keys() for row in rows))))

    max_candidates = env_int("MAX_CANDIDATES_PER_ROW", 4096)
    allow_same = env_bool("ALLOW_SAME_SPAN_PAIRS", False)
    fail_fast = env_bool("FAIL_FAST", False)

    occurrence_inputs: Dict[Tuple[str, int, int], Tuple[str, int, int]] = {}
    candidates: List[Candidate] = []
    candidate_by_uid: Dict[str, Candidate] = {}
    candidates_by_row: Dict[int, List[Candidate]] = defaultdict(list)
    unresolved_rows: Dict[int, str] = {}
    occurrence_audit: List[Dict[str, Any]] = []

    for row in rows:
        subject_spans, subject_mode, subject_variant = robust_occurrences(
            row.sentence, row.subject
        )
        object_spans, object_mode, object_variant = robust_occurrences(
            row.sentence, row.object
        )

        audit: Dict[str, Any] = {
            "row_index": row.row_index,
            "triple_id": row.triple_id,
            "subject": row.subject,
            "object": row.object,
            "subject_match_mode": subject_mode,
            "object_match_mode": object_mode,
            "subject_match_variant": subject_variant,
            "object_match_variant": object_variant,
            "subject_spans": [list(span) for span in subject_spans],
            "object_spans": [list(span) for span in object_spans],
        }

        if not subject_spans or not object_spans:
            reason = (
                f"mention_not_found: subject={bool(subject_spans)}, "
                f"object={bool(object_spans)}"
            )
            unresolved_rows[row.row_index] = reason
            audit["error"] = reason
            occurrence_audit.append(audit)
            if fail_fast:
                raise ValueError(f"Row {row.row_index} ({row.triple_id}): {reason}")
            continue

        span_pairs = [
            (subject_span, object_span)
            for subject_span in subject_spans
            for object_span in object_spans
            if allow_same or subject_span != object_span
        ]
        if not span_pairs and subject_spans and object_spans:
            span_pairs = [(subject_spans[0], object_spans[0])]
            audit["same_span_fallback_used"] = True

        if len(span_pairs) > max_candidates:
            reason = (
                f"candidate_count={len(span_pairs)} exceeds "
                f"MAX_CANDIDATES_PER_ROW={max_candidates}"
            )
            unresolved_rows[row.row_index] = reason
            audit["error"] = reason
            occurrence_audit.append(audit)
            if fail_fast:
                raise RuntimeError(f"Row {row.row_index} ({row.triple_id}): {reason}")
            continue

        for span in dict.fromkeys(subject_spans + object_spans):
            key = (row.sentence, span[0], span[1])
            occurrence_inputs.setdefault(key, (row.sentence, span[0], span[1]))

        for candidate_index, (subject_span, object_span) in enumerate(span_pairs):
            uid = f"row{row.row_index:06d}_cand{candidate_index:04d}"
            candidate = Candidate(
                uid=uid,
                row_index=row.row_index,
                candidate_index=candidate_index,
                source_triple_id=row.triple_id,
                subject_span=subject_span,
                object_span=object_span,
                subject_text=row.sentence[subject_span[0] : subject_span[1]],
                object_text=row.sentence[object_span[0] : object_span[1]],
                subject_match_mode=subject_mode,
                object_match_mode=object_mode,
            )
            candidate_by_uid[uid] = candidate
            candidates.append(candidate)
            candidates_by_row[row.row_index].append(candidate)

        audit["candidate_count"] = len(span_pairs)
        occurrence_audit.append(audit)

    write_jsonl(output_dir / "input_occurrence_audit.jsonl", occurrence_audit)
    print("Unique entity occurrences:", len(occurrence_inputs))
    print("Candidate pairs:", len(candidates))
    print("Unresolved rows:", len(unresolved_rows))
    print(
        "Candidate-count distribution:",
        Counter(len(candidates_by_row.get(row.row_index, [])) for row in rows),
    )
    print("Subject match modes:", Counter(x["subject_match_mode"] for x in occurrence_audit))
    print("Object match modes:", Counter(x["object_match_mode"] for x in occurrence_audit))

    if args.validate_input_only:
        print("Input-only validation passed; no models were loaded.")
        return

    # TypePredictor stage.
    public_token = optional_token("HF_TOKENTWO", "HF_TOKEN", "HUGGINGFACE_TOKEN")
    type_engine = TypePredictorEngine(
        model_id=os.getenv("TYPE_MODEL_ID", DEFAULT_TYPE_MODEL_ID),
        revision=os.getenv("TYPE_MODEL_REVISION", "main"),
        token=public_token,
        device=os.getenv("TYPE_DEVICE", "auto"),
        use_fp16=env_bool("TYPE_USE_FP16", True),
    )
    started = time.time()
    type_results = type_engine.predict_many(
        occurrence_inputs,
        batch_size=env_int("TYPE_BATCH_SIZE", 64),
    )
    print(f"Type prediction finished in {(time.time() - started) / 60:.2f} minutes")
    type_model_revision = type_engine.resolved_revision

    type_audit: List[Dict[str, Any]] = []
    for key, result in type_results.items():
        sentence, start, end = occurrence_inputs[key]
        type_audit.append(
            {
                "sentence": sentence,
                "start": start,
                "end": end,
                "mention": sentence[start:end],
                **result,
            }
        )
    write_jsonl(output_dir / "type_predictions.jsonl", type_audit)

    for candidate in candidates:
        row = rows[candidate.row_index]
        subject_key = (
            row.sentence,
            candidate.subject_span[0],
            candidate.subject_span[1],
        )
        object_key = (
            row.sentence,
            candidate.object_span[0],
            candidate.object_span[1],
        )
        subject_result = type_results[subject_key]
        object_result = type_results[object_key]
        candidate.subject_type = str(subject_result["predicted_type"])
        candidate.object_type = str(object_result["predicted_type"])
        candidate.subject_type_confidence = float(subject_result["confidence"])
        candidate.object_type_confidence = float(object_result["confidence"])

    # Free the 0.3B encoder before loading the 7B model.
    del type_engine
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    adapter_dir, adapter_revision = resolve_adapter_snapshot(output_dir)
    common = load_release_common(adapter_dir)
    inference_config = selected_inference_config(adapter_dir)

    raw_candidate_rows: List[Dict[str, Any]] = []
    for candidate in candidates:
        row = rows[candidate.row_index]
        raw_candidate_rows.append(
            {
                "id": candidate.uid,
                "sentence_id": row.sentence_id,
                "triple_id": candidate.uid,
                "source_triple_id": row.triple_id,
                "row_index": row.row_index,
                "candidate_index": candidate.candidate_index,
                "sentence": row.sentence,
                "subject": candidate.subject_text,
                "object": candidate.object_text,
                "subject_start": candidate.subject_span[0],
                "subject_end": candidate.subject_span[1],
                "object_start": candidate.object_span[0],
                "object_end": candidate.object_span[1],
                "subject_type": candidate.subject_type,
                "object_type": candidate.object_type,
                "relation": "",
            }
        )

    raw_candidates_path = output_dir / "candidate_inputs_raw.jsonl"
    prepared_candidates_path = output_dir / "candidate_inputs_prepared.jsonl"
    write_jsonl(raw_candidates_path, raw_candidate_rows)
    prepare_candidates_with_release(
        adapter_dir,
        raw_candidates_path,
        prepared_candidates_path,
    )
    prepared_rows = read_jsonl(prepared_candidates_path)
    if len(prepared_rows) != len(candidates):
        raise RuntimeError(
            f"Prepared candidate count {len(prepared_rows)} != generated count {len(candidates)}"
        )
    if str(prepared_rows[0]["prompt_version"]) != str(inference_config["prompt_version"]):
        raise RuntimeError("Prepared prompt version differs from released inference config")

    started = time.time()
    yehia_summary = run_yehia_candidates(
        adapter_dir=adapter_dir,
        prepared_rows=prepared_rows,
        candidate_by_uid=candidate_by_uid,
        common=common,
        inference_config=inference_config,
        batch_size=env_int("YEHIA_BATCH_SIZE", 16),
    )
    print(f"Yehia inference finished in {(time.time() - started) / 60:.2f} minutes")
    write_jsonl(output_dir / "candidate_predictions.jsonl", yehia_summary["records"])

    relation_order = canonical_relation_order(adapter_dir)
    legal_relations = relation_order
    requested_method = args.aggregation
    methods = [requested_method]
    if env_bool("WRITE_ALTERNATIVE_SUBMISSIONS", True):
        methods = list(dict.fromkeys([requested_method, *AGGREGATORS.keys()]))

    positive_fraction = env_float("CANDIDATE_POSITIVE_FRACTION", 0.5)
    soft_pool_threshold = env_float("SOFT_POOL_POSITIVE_THRESHOLD", 0.5)
    labels_by_method: Dict[str, List[str]] = {method: [] for method in methods}
    row_debug: List[Dict[str, Any]] = []

    for row in rows:
        row_candidates = candidates_by_row.get(row.row_index, [])
        debug: Dict[str, Any] = {
            "row_index": row.row_index,
            "triple_id": row.triple_id,
            "sentence": row.sentence,
            "subject": row.subject,
            "object": row.object,
            "candidate_count": len(row_candidates),
            "unresolved_error": unresolved_rows.get(row.row_index),
            "methods": {},
        }

        for method in methods:
            if not row_candidates:
                label = NO_RELATION
                details = {
                    "gate_passed": False,
                    "fallback_reason": unresolved_rows.get(
                        row.row_index, "no_candidates"
                    ),
                }
            elif method == "majority_vote":
                label, details = aggregate_majority_vote(
                    row_candidates,
                    relation_order,
                    positive_fraction,
                )
            elif method == "soft_pool":
                label, details = aggregate_soft_pool(
                    row_candidates,
                    relation_order,
                    soft_pool_threshold,
                )
            else:
                label, details = AGGREGATORS[method](
                    row_candidates,
                    relation_order,
                )
            labels_by_method[method].append(label)
            debug["methods"][method] = {"label": label, **details}

        if env_bool("INCLUDE_CANDIDATES_IN_DEBUG", True):
            debug["candidates"] = [asdict(candidate) for candidate in row_candidates]
        row_debug.append(debug)

    write_jsonl(output_dir / "row_predictions_debug.jsonl", row_debug)
    write_jsonl(
        output_dir / "predictions_primary.jsonl",
        [
            {
                **row.original,
                "relation": labels_by_method[requested_method][row.row_index],
            }
            for row in rows
        ],
    )

    triple_ids = [row.triple_id for row in rows]
    written: Dict[str, Tuple[Path, Path]] = {}
    for method in methods:
        written[method] = write_submission(
            output_dir,
            method,
            triple_ids,
            labels_by_method[method],
            legal_relations,
        )
        print(f"{method} counts:", Counter(labels_by_method[method]).most_common())

    primary_text, primary_zip = written[requested_method]
    final_text = output_dir / "predictions.txt"
    final_zip = output_dir / "submission.zip"
    shutil.copy2(primary_text, final_text)
    shutil.copy2(primary_zip, final_zip)
    validate_submission(final_zip, triple_ids, legal_relations)

    workspace_submission = Path(
        os.getenv("FINAL_SUBMISSION_PATH", str(WORKSPACE / "submission.zip"))
    ).expanduser()
    workspace_submission.parent.mkdir(parents=True, exist_ok=True)
    if workspace_submission.resolve() != final_zip.resolve():
        shutil.copy2(final_zip, workspace_submission)
        validate_submission(workspace_submission, triple_ids, legal_relations)

    manifest = {
        "status": "passed",
        "input": str(input_path),
        "input_sha256": sha256(input_path),
        "row_count": len(rows),
        "unique_occurrence_count": len(occurrence_inputs),
        "candidate_count": len(candidates),
        "unresolved_row_count": len(unresolved_rows),
        "type_model_id": os.getenv("TYPE_MODEL_ID", DEFAULT_TYPE_MODEL_ID),
        "type_model_revision": type_model_revision,
        "yehia_adapter_model_id": os.getenv(
            "YEHIA_ADAPTER_MODEL_ID", DEFAULT_ADAPTER_ID
        ),
        "yehia_adapter_revision": adapter_revision,
        "base_model_id": yehia_summary["base_model_id"],
        "base_model_revision": yehia_summary["base_model_revision"],
        "prompt_version": inference_config["prompt_version"],
        "no_relation_logit_bias": yehia_summary["no_relation_logit_bias"],
        "max_prompt_tokens": yehia_summary["max_prompt_tokens"],
        "inference_backend": yehia_summary["inference_backend"],
        "vllm_version": yehia_summary["vllm_version"],
        "primary_aggregation_method": requested_method,
        "candidate_positive_fraction": positive_fraction,
        "soft_pool_positive_threshold": soft_pool_threshold,
        "alternative_methods": methods,
        "output_submission": str(final_zip),
        "output_submission_sha256": sha256(final_zip),
        "workspace_submission": str(workspace_submission),
        "workspace_submission_sha256": sha256(workspace_submission),
    }
    (output_dir / "run_manifest.json").write_text(
        json.dumps(manifest, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )

    print("\nDONE")
    print("Primary aggregation:", requested_method)
    print("Validated output:", final_zip)
    print("RunPod convenience copy:", workspace_submission)
    print("Debug rows:", output_dir / "row_predictions_debug.jsonl")
    print("Manifest:", output_dir / "run_manifest.json")


if __name__ == "__main__":
    main()