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