model.py

#!/usr/bin/env python3
"""Train and evaluate DRU-RE-Yehia with one-token multiple-choice QLoRA.

The model sees Arabic relation options labeled by one-token Arabic symbols and
is supervised only on the selected symbol. Evaluation performs constrained
greedy choice among symbols valid for each row, so formatting cannot become an
invalid output. The transformed dataset is read locally. Yehia is also loaded
locally; when its snapshot is absent, this entrypoint can download it once with
HF_TOKENONE and then continue from the local snapshot.
"""

from __future__ import annotations

import inspect
import json
import logging
import math
import os
import random
import shutil
from collections import Counter, defaultdict
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, Iterable, List, Mapping, Optional, Sequence, Tuple

os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0"

import bitsandbytes as bnb
import numpy as np
import torch
from dotenv import load_dotenv
from huggingface_hub import snapshot_download
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
from sklearn.metrics import accuracy_score, f1_score
from torch.nn.utils.rnn import pad_sequence
from torch.utils.data import Dataset, WeightedRandomSampler
from tqdm.auto import tqdm
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    Trainer,
    TrainerCallback,
    TrainingArguments,
    set_seed,
)
from transformers.pytorch_utils import Conv1D
from transformers.trainer_utils import get_last_checkpoint

from re_sft_common import (
    NO_RELATION_FULL,
    OPTION_CODES,
    PROMPT_VERSION,
    env_bool,
    env_float,
    env_int,
    env_str,
    load_jsonl,
    validate_local_hf_revision,
)


load_dotenv()


def utc_now() -> str:
    return datetime.now(timezone.utc).replace(microsecond=0).isoformat()


def dump_json(value: Any, path: Path) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as handle:
        json.dump(value, handle, ensure_ascii=False, indent=2)
        handle.write("\n")


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


DATASET_DIR = Path(env_str("CHOICE_DATASET_DIR", "data/Yehia-RE-SFT"))
BASE_MODEL_ID = env_str("YEHIA_BASE_MODEL_ID", "Navid-AI/Yehia-7B-preview")
BASE_MODEL_REVISION = env_str(
    "YEHIA_BASE_MODEL_REVISION", "b9dda4715eafee7e8090d2c83cfe078d75f4ebb8"
)
MODEL_DIR = Path(env_str("LOCAL_YEHIA_MODEL_DIR", "models/Yehia-7B-preview"))
ALLOW_MODEL_DOWNLOAD = env_bool("ALLOW_MODEL_DOWNLOAD", True)
OUTPUT_DIR = Path(env_str("CHOICE_OUTPUT_DIR", "runs/DRU-RE-Yehia"))
CHECKPOINT_DIR = OUTPUT_DIR / "checkpoints"
ARTIFACT_DIR = OUTPUT_DIR / "artifacts"
METRICS_DIR = OUTPUT_DIR / "metrics"
PREDICTIONS_DIR = OUTPUT_DIR / "predictions"
LOG_DIR = Path(env_str("LOG_DIR", "logs"))
LOG_FILE = LOG_DIR / env_str("CHOICE_TRAINING_LOG_FILE", "train.log")

SEED = env_int("SEED", 42)
MAX_LENGTH = env_int("MAX_LENGTH", 1024)
NUM_EPOCHS = env_float("NUM_EPOCHS", 3.0)
MAX_STEPS = env_int("CHOICE_MAX_STEPS", -1)
TRAIN_BATCH = env_int("TRAIN_BATCH_SIZE", 4)
EVAL_BATCH = env_int("CHOICE_EVAL_BATCH_SIZE", 16)
GRAD_ACCUM = env_int("GRAD_ACCUM_STEPS", 4)
LEARNING_RATE = env_float("LEARNING_RATE", 5e-5)
WARMUP_RATIO = env_float("WARMUP_RATIO", 0.05)
WEIGHT_DECAY = env_float("WEIGHT_DECAY", 0.0)
MAX_GRAD_NORM = env_float("MAX_GRAD_NORM", 1.0)
LOGGING_STEPS = env_int("LOGGING_STEPS", 20)
EVAL_STEPS = env_int("SAVE_EVAL_STEPS", 250)
SAVE_TOTAL_LIMIT = env_int("SAVE_TOTAL_LIMIT", 2)
CLASS_SAMPLING_ALPHA = env_float("CLASS_SAMPLING_ALPHA", 0.25)
CALIBRATE_NO_RELATION = env_bool("CALIBRATE_NO_RELATION", True)
LORA_R = env_int("LORA_R", 16)
LORA_ALPHA = env_int("LORA_ALPHA", 32)
LORA_DROPOUT = env_float("LORA_DROPOUT", 0.05)
USE_RSLORA = env_bool("USE_RSLORA", True)
RAW_TARGETS = env_str("LORA_TARGET_MODULES", "all-linear").strip()
TARGET_MODULES: str | List[str] = (
    "all-linear"
    if RAW_TARGETS == "all-linear"
    else [item.strip() for item in RAW_TARGETS.split(",") if item.strip()]
)
ATTN_IMPL = env_str("ATTENTION_IMPLEMENTATION", "sdpa")
RESUME = env_bool("CHOICE_RESUME", True)

for path in (OUTPUT_DIR, CHECKPOINT_DIR, ARTIFACT_DIR, METRICS_DIR, PREDICTIONS_DIR, LOG_DIR):
    path.mkdir(parents=True, exist_ok=True)

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s | %(levelname)s | %(message)s",
    handlers=[logging.FileHandler(LOG_FILE, encoding="utf-8"), logging.StreamHandler()],
)
logger = logging.getLogger("yehia_choice_qlora")


def stage(name: str) -> None:
    message = f"\n{'#' * 96}\n# {name}\n{'#' * 96}"
    print(message, flush=True)
    logger.info(message)


set_seed(SEED)
random.seed(SEED)
np.random.seed(SEED)

stage("CONFIGURATION")
config = {
    "created_at_utc": utc_now(),
    "dataset_dir": str(DATASET_DIR),
    "base_model_id": BASE_MODEL_ID,
    "base_model_revision": BASE_MODEL_REVISION,
    "model_dir": str(MODEL_DIR),
    "allow_model_download": ALLOW_MODEL_DOWNLOAD,
    "output_dir": str(OUTPUT_DIR),
    "prompt_version": PROMPT_VERSION,
    "objective": "one_token_full_vocab_multiple_choice",
    "loss_mode": "full_vocab",
    "decision_tokens": list(OPTION_CODES),
    "max_length": MAX_LENGTH,
    "num_epochs": NUM_EPOCHS,
    "max_steps": MAX_STEPS,
    "train_batch": TRAIN_BATCH,
    "eval_batch": EVAL_BATCH,
    "gradient_accumulation": GRAD_ACCUM,
    "effective_batch": TRAIN_BATCH * GRAD_ACCUM,
    "learning_rate": LEARNING_RATE,
    "warmup_ratio": WARMUP_RATIO,
    "weight_decay": WEIGHT_DECAY,
    "max_grad_norm": MAX_GRAD_NORM,
    "eval_steps": EVAL_STEPS,
    "class_sampling_alpha": CLASS_SAMPLING_ALPHA,
    "calibrate_no_relation": CALIBRATE_NO_RELATION,
    "lora_r": LORA_R,
    "lora_alpha": LORA_ALPHA,
    "lora_dropout": LORA_DROPOUT,
    "use_rslora": USE_RSLORA,
    "target_modules": TARGET_MODULES,
    "seed": SEED,
}
dump_json(config, OUTPUT_DIR / "run_config.json")
print(json.dumps(config, ensure_ascii=False, indent=2), flush=True)


stage("LOAD AND VALIDATE LOCAL DATASET")
required_files = {
    "train": DATASET_DIR / "train.jsonl",
    "validation": DATASET_DIR / "validation.jsonl",
    "official": DATASET_DIR / "official.jsonl",
}
for path in required_files.values():
    if not path.exists():
        raise FileNotFoundError(path)
rows = {split: load_jsonl(path) for split, path in required_files.items()}
expected_counts = {"train": 15686, "validation": 1687, "official": 2074}
for split, expected in expected_counts.items():
    if len(rows[split]) != expected:
        raise ValueError(f"{split} count {len(rows[split])} != {expected}")

for split in ("train", "validation"):
    for index, row in enumerate(rows[split]):
        options = row.get("allowed_options_ar") or []
        codes = row.get("option_codes") or []
        gold_index = row.get("gold_option_index")
        gold_code = row.get("gold_answer_code")
        if row.get("prompt_version") != PROMPT_VERSION:
            raise ValueError(f"Prompt mismatch at {split}:{index}")
        if not options or options[-1] != "لا توجد علاقة":
            raise ValueError(f"No-relation is not last at {split}:{index}")
        if codes != list(OPTION_CODES[: len(options)]):
            raise ValueError(f"Option-code mismatch at {split}:{index}")
        if not isinstance(gold_index, int) or not 0 <= gold_index < len(options):
            raise ValueError(f"Gold option index mismatch at {split}:{index}")
        if gold_code != codes[gold_index]:
            raise ValueError(f"Gold code mismatch at {split}:{index}")
        if row["messages"][-1] != {"role": "assistant", "content": gold_code}:
            raise ValueError(f"Assistant target mismatch at {split}:{index}")
for index, row in enumerate(rows["official"]):
    if row.get("gold_answer_code") is not None or len(row.get("messages", [])) != 2:
        raise ValueError(f"Official label leakage at row {index}")

all_labels = sorted(
    {str(row["gold_relation_full"]) for row in rows["train"] + rows["validation"]}
)
if len(all_labels) != 41 or NO_RELATION_FULL not in all_labels:
    raise ValueError(f"Expected 41 labels; found {len(all_labels)}")
class_counts = Counter(str(row["gold_relation_full"]) for row in rows["train"])
dump_json(dict(sorted(class_counts.items())), ARTIFACT_DIR / "train_class_counts.json")
print({split: len(value) for split, value in rows.items()}, flush=True)


def require_yehia_token() -> str:
    """Read the gated-model token without ever logging its value."""
    token = os.environ.get("HF_TOKENONE", "").strip()
    if not token or token.startswith("hf_your_"):
        raise RuntimeError(
            "Yehia is not available locally and HF_TOKENONE is missing. "
            "Add the gated-model token to .env or pre-populate LOCAL_YEHIA_MODEL_DIR."
        )
    return token


def ensure_local_model_snapshot() -> None:
    """Use an existing local snapshot or download Yehia into MODEL_DIR once."""
    if (MODEL_DIR / "config.json").is_file():
        revision = validate_local_hf_revision(MODEL_DIR, BASE_MODEL_REVISION)
        print(f"Using existing local Yehia snapshot: {MODEL_DIR}", flush=True)
        print(f"Verified local Yehia revision: {revision}", flush=True)
        return
    if not ALLOW_MODEL_DOWNLOAD:
        raise FileNotFoundError(
            f"Local Yehia snapshot is missing at {MODEL_DIR} and "
            "ALLOW_MODEL_DOWNLOAD=false"
        )
    offline = os.environ.get("HF_HUB_OFFLINE", "").strip().lower()
    if offline in {"1", "true", "yes", "on"}:
        raise RuntimeError(
            "Local Yehia snapshot is missing but HF_HUB_OFFLINE is enabled. "
            "Disable offline mode for the first run or download the model separately."
        )
    MODEL_DIR.mkdir(parents=True, exist_ok=True)
    print(f"Downloading gated base model {BASE_MODEL_ID} into {MODEL_DIR}", flush=True)
    snapshot_download(
        repo_id=BASE_MODEL_ID,
        revision=BASE_MODEL_REVISION,
        local_dir=str(MODEL_DIR),
        token=require_yehia_token(),
    )
    if not (MODEL_DIR / "config.json").is_file():
        raise RuntimeError(f"Downloaded Yehia snapshot is incomplete: {MODEL_DIR}")
    (MODEL_DIR / ".dru_hf_revision").write_text(BASE_MODEL_REVISION + "\n", encoding="utf-8")
    validate_local_hf_revision(MODEL_DIR, BASE_MODEL_REVISION)


stage("RESOLVE LOCAL YEHIA SNAPSHOT AND BUILD DECISION-TOKEN LABELS")
ensure_local_model_snapshot()
tokenizer = AutoTokenizer.from_pretrained(
    str(MODEL_DIR), local_files_only=True, use_fast=True
)
if not tokenizer.chat_template:
    raise RuntimeError("Local Yehia tokenizer has no chat template")
if tokenizer.pad_token_id is None:
    tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right"

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]))
dump_json(
    {"codes": list(OPTION_CODES), "token_ids": code_token_ids},
    ARTIFACT_DIR / "decision_token_inventory.json",
)


def encode_row(row: Mapping[str, Any]) -> Dict[str, Any]:
    prompt_ids = list(
        tokenizer.apply_chat_template(
            row["prompt_messages"], tokenize=True, add_generation_prompt=True
        )
    )
    full_ids = list(
        tokenizer.apply_chat_template(
            row["messages"], tokenize=True, add_generation_prompt=False
        )
    )
    common = 0
    for first, second in zip(prompt_ids, full_ids):
        if first != second:
            break
        common += 1
    if common != len(prompt_ids):
        raise ValueError(
            f"Prompt/full boundary mismatch for {row.get('id')}: "
            f"prompt={len(prompt_ids)}, common={common}"
        )
    if len(full_ids) > MAX_LENGTH:
        raise ValueError(f"Sequence exceeds {MAX_LENGTH} tokens for {row.get('id')}")
    expected_token = code_token_ids[int(row["gold_option_index"])]
    if common >= len(full_ids) or full_ids[common] != expected_token:
        raise ValueError(
            f"Decision token mismatch for {row.get('id')}: "
            f"expected={expected_token}, observed={full_ids[common] if common < len(full_ids) else None}"
        )
    # Pure classification objective: only the one decision token contributes.
    labels = [-100] * len(full_ids)
    labels[common] = expected_token
    return {
        "input_ids": full_ids,
        "attention_mask": [1] * len(full_ids),
        "labels": labels,
        "prompt_token_count": common,
        "decision_token_index": common,
        "decision_token_id": expected_token,
    }


encoded = {
    split: [encode_row(row) for row in tqdm(rows[split], desc=f"tokenize:{split}")]
    for split in ("train", "validation")
}
for split, values in encoded.items():
    lengths = [len(item["input_ids"]) for item in values]
    print(
        f"{split}: min={min(lengths)} mean={np.mean(lengths):.1f} "
        f"p95={np.percentile(lengths, 95):.1f} max={max(lengths)}",
        flush=True,
    )

audit_indices = [
    0,
    next(i for i, row in enumerate(rows["train"]) if row["gold_relation_full"] == NO_RELATION_FULL),
    max(range(len(encoded["train"])), key=lambda i: len(encoded["train"][i]["input_ids"])),
]
masking_audit: List[Dict[str, Any]] = []
for index in dict.fromkeys(audit_indices):
    row = rows["train"][index]
    item = encoded["train"][index]
    supervised_positions = [i for i, label in enumerate(item["labels"]) if label != -100]
    if supervised_positions != [item["decision_token_index"]]:
        raise RuntimeError(f"Non-decision tokens are supervised for train row {index}")
    masking_audit.append(
        {
            "index": index,
            "id": row["id"],
            "gold_relation_full": row["gold_relation_full"],
            "gold_answer_code": row["gold_answer_code"],
            "decoded_conversation": tokenizer.decode(
                item["input_ids"], skip_special_tokens=False
            ),
            "prompt_token_count": item["prompt_token_count"],
            "masked_prompt_token_count": item["prompt_token_count"],
            "supervised_token_count": 1,
            "supervised_token_id": item["decision_token_id"],
            "decoded_supervised_token": tokenizer.decode([item["decision_token_id"]]),
            "all_nondecision_tokens_masked": True,
        }
    )
dump_json(
    {"status": "passed", "objective": "decision_token_only", "rows": masking_audit},
    ARTIFACT_DIR / "decision_token_masking_audit.json",
)


class EncodedDataset(Dataset):
    def __init__(self, items: Sequence[Mapping[str, Any]]):
        self.items = list(items)

    def __len__(self) -> int:
        return len(self.items)

    def __getitem__(self, index: int) -> Dict[str, torch.Tensor]:
        item = self.items[index]
        return {
            key: torch.tensor(item[key], dtype=torch.long)
            for key in ("input_ids", "attention_mask", "labels")
        }


class CompletionCollator:
    def __call__(
        self, features: Sequence[Mapping[str, torch.Tensor]]
    ) -> Dict[str, torch.Tensor]:
        return {
            "input_ids": pad_sequence(
                [item["input_ids"] for item in features],
                batch_first=True,
                padding_value=tokenizer.pad_token_id,
            ),
            "attention_mask": pad_sequence(
                [item["attention_mask"] for item in features],
                batch_first=True,
                padding_value=0,
            ),
            "labels": pad_sequence(
                [item["labels"] for item in features],
                batch_first=True,
                padding_value=-100,
            ),
        }


train_dataset = EncodedDataset(encoded["train"])
validation_dataset = EncodedDataset(encoded["validation"])
collator = CompletionCollator()


def normalize_module_name(name: str) -> str:
    prefixes = ("base_model.model.", "model.", "module.")
    changed = True
    while changed:
        changed = False
        for prefix in prefixes:
            if name.startswith(prefix):
                name = name[len(prefix) :]
                changed = True
    return name


def supported_linear_classes() -> Tuple[type, ...]:
    classes: List[type] = [torch.nn.Linear, Conv1D]
    for attr in ("Linear4bit", "Linear8bitLt"):
        cls = getattr(getattr(bnb, "nn", object()), attr, None)
        if isinstance(cls, type):
            classes.append(cls)
    return tuple(classes)


def module_class_name(module: torch.nn.Module) -> str:
    return f"{module.__class__.__module__}.{module.__class__.__name__}"


def enumerate_eligible_linear_modules(active_model: torch.nn.Module) -> Dict[str, Any]:
    """Inventory every PEFT-supported projection before attaching adapters."""
    supported = supported_linear_classes()
    output_embedding = active_model.get_output_embeddings()
    input_embedding = active_model.get_input_embeddings()
    output_ids = {id(output_embedding)} if output_embedding is not None else set()
    input_ids = {id(input_embedding)} if input_embedding is not None else set()
    records: List[Dict[str, Any]] = []
    output_head_names: List[str] = []
    embedding_names: List[str] = []

    for name, module in active_model.named_modules():
        if not name:
            continue
        normalized = normalize_module_name(name)
        is_output_head = id(module) in output_ids
        is_embedding = id(module) in input_ids or isinstance(module, torch.nn.Embedding)
        if is_output_head:
            output_head_names.append(normalized)
        if is_embedding:
            embedding_names.append(normalized)
        if isinstance(module, supported):
            records.append(
                {
                    "name": normalized,
                    "original_name": name,
                    "class": module_class_name(module),
                    "suffix": normalized.rsplit(".", 1)[-1],
                    "is_output_head": bool(is_output_head),
                    "is_embedding": bool(is_embedding),
                }
            )

    expected_wrapped = [
        item["name"]
        for item in records
        if not item["is_output_head"] and not item["is_embedding"]
    ]
    inventory = {
        "target_modules": TARGET_MODULES,
        "supported_classes": [f"{cls.__module__}.{cls.__name__}" for cls in supported],
        "total_eligible_linear_or_conv1d_modules": len(records),
        "expected_lora_wrapped_modules": len(expected_wrapped),
        "output_head_names": sorted(output_head_names),
        "embedding_names": sorted(embedding_names),
        "modules": sorted(records, key=lambda item: item["name"]),
        "expected_wrapped_module_names": sorted(expected_wrapped),
    }
    dump_json(inventory, ARTIFACT_DIR / "all_linear_eligible_modules.json")
    return inventory


def trainable_parameter_summary(active_model: torch.nn.Module) -> Dict[str, Any]:
    if hasattr(active_model, "get_nb_trainable_parameters"):
        trainable, total = active_model.get_nb_trainable_parameters()
    else:
        total = sum(parameter.numel() for parameter in active_model.parameters())
        trainable = sum(
            parameter.numel()
            for parameter in active_model.parameters()
            if parameter.requires_grad
        )
    trainable_names = [
        name for name, parameter in active_model.named_parameters() if parameter.requires_grad
    ]
    unintended = [name for name in trainable_names if "lora_" not in name]
    if unintended:
        raise RuntimeError(f"Unexpected non-LoRA trainable parameters: {unintended[:20]}")
    return {
        "total_parameters": int(total),
        "trainable_parameters": int(trainable),
        "trainable_percentage": float(100.0 * trainable / total if total else 0.0),
        "trainable_tensor_count": len(trainable_names),
        "trainable_parameter_names": trainable_names,
    }


def validate_lora_wrapping(
    active_model: torch.nn.Module, eligible_inventory: Mapping[str, Any]
) -> Dict[str, Any]:
    """Prove true all-linear coverage and that the frozen head/base stayed frozen."""
    wrapped: List[Dict[str, Any]] = []
    for name, module in active_model.named_modules():
        if hasattr(module, "lora_A") and hasattr(module, "lora_B"):
            normalized = normalize_module_name(name)
            wrapped.append(
                {
                    "name": normalized,
                    "original_name": name,
                    "class": module_class_name(module),
                    "suffix": normalized.rsplit(".", 1)[-1],
                }
            )

    wrapped_names = {item["name"] for item in wrapped}
    expected_names = set(eligible_inventory["expected_wrapped_module_names"])
    output_head_names = set(eligible_inventory["output_head_names"])
    if TARGET_MODULES == "all-linear":
        missing = sorted(expected_names - wrapped_names)
        unexpected = sorted((wrapped_names - expected_names) - output_head_names)
        if missing:
            raise RuntimeError(f"Incomplete all-linear LoRA coverage: {missing[:50]}")
        if unexpected:
            raise RuntimeError(f"Unexpected all-linear LoRA modules: {unexpected[:50]}")
    if wrapped_names & output_head_names:
        raise RuntimeError(
            "Output head was unintentionally adapted: "
            f"{sorted(wrapped_names & output_head_names)}"
        )

    required = {"q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"}
    wrapped_suffixes = {item["suffix"] for item in wrapped}
    missing_required = sorted(required - wrapped_suffixes)
    if missing_required:
        raise RuntimeError(f"Missing attention/MLP adapters: {missing_required}")

    by_suffix = Counter(item["suffix"] for item in wrapped)
    by_layer: Counter[str] = Counter()
    for item in wrapped:
        parts = item["name"].split(".")
        layer = next(
            (
                f"{parts[index]}.{parts[index + 1]}"
                for index in range(len(parts) - 1)
                if parts[index] == "layers"
            ),
            "non_layer",
        )
        by_layer[layer] += 1

    audit = {
        "target_modules": TARGET_MODULES,
        "eligible_linear_or_conv1d_count": eligible_inventory[
            "total_eligible_linear_or_conv1d_modules"
        ],
        "expected_lora_wrapped_count": len(expected_names),
        "lora_wrapped_count": len(wrapped),
        "wrapped_by_suffix": dict(sorted(by_suffix.items())),
        "wrapped_by_transformer_layer": dict(sorted(by_layer.items())),
        "output_head_names": sorted(output_head_names),
        "output_head_adapted": False,
        "attention_projection_adapters_present": all(
            suffix in wrapped_suffixes for suffix in ("q_proj", "k_proj", "v_proj", "o_proj")
        ),
        "mlp_projection_adapters_present": all(
            suffix in wrapped_suffixes for suffix in ("gate_proj", "up_proj", "down_proj")
        ),
        "base_quantized_weights_trainable": False,
        "only_lora_parameters_trainable": True,
        "wrapped_modules": sorted(wrapped, key=lambda item: item["name"]),
        **trainable_parameter_summary(active_model),
    }
    dump_json(audit, ARTIFACT_DIR / "lora_wrapped_modules.json")
    return audit


def positive_scores(gold: Sequence[str], pred: Sequence[str]) -> Dict[str, Any]:
    tp = fp = fn = 0
    for truth, guess in zip(gold, pred):
        truth_positive = truth != NO_RELATION_FULL
        guess_positive = guess != NO_RELATION_FULL
        if truth_positive and guess_positive and truth == guess:
            tp += 1
        else:
            fp += int(guess_positive)
            fn += int(truth_positive)
    precision = tp / (tp + fp) if tp + fp else 0.0
    recall = tp / (tp + fn) if tp + fn else 0.0
    f1 = 2 * precision * recall / (precision + recall) if precision + recall else 0.0
    return {
        "positive_micro_precision": precision,
        "positive_micro_recall": recall,
        "positive_micro_f1": f1,
        "positive_tp": tp,
        "positive_fp": fp,
        "positive_fn": fn,
    }


def relation_metrics(gold: Sequence[str], pred: Sequence[str]) -> Dict[str, Any]:
    return {
        "rows": len(gold),
        "accuracy": float(accuracy_score(gold, pred)),
        "macro_f1_41": float(
            f1_score(gold, pred, labels=all_labels, average="macro", zero_division=0)
        ),
        "weighted_f1_41": float(
            f1_score(gold, pred, labels=all_labels, average="weighted", zero_division=0)
        ),
        **positive_scores(gold, pred),
        "invalid_outputs": 0,
        "invalid_output_rate": 0.0,
    }


def choose_from_scores(
    records: Sequence[Mapping[str, Any]], score_rows: Sequence[Sequence[float]], bias: float
) -> Tuple[List[str], List[int]]:
    predictions: List[str] = []
    positions: List[int] = []
    for row, raw_scores in zip(records, score_rows):
        scores = list(raw_scores)
        scores[-1] += bias
        position = int(np.argmax(scores))
        positions.append(position)
        predictions.append(str(row["allowed_relation_full_labels"][position]))
    return predictions, positions


def calibrate_bias(
    records: Sequence[Mapping[str, Any]], score_rows: Sequence[Sequence[float]]
) -> Tuple[float, Dict[str, Any]]:
    gold = [str(row["gold_relation_full"]) for row in records]
    if not CALIBRATE_NO_RELATION:
        pred, _ = choose_from_scores(records, score_rows, 0.0)
        return 0.0, relation_metrics(gold, pred)
    best_bias = 0.0
    best_metrics: Optional[Dict[str, Any]] = None
    for bias in np.arange(-12.0, 12.0001, 0.1):
        pred, _ = choose_from_scores(records, score_rows, float(bias))
        current = relation_metrics(gold, pred)
        key = (current["positive_micro_f1"], current["accuracy"], -abs(float(bias)))
        old_key = (
            (-1.0, -1.0, -math.inf)
            if best_metrics is None
            else (
                best_metrics["positive_micro_f1"],
                best_metrics["accuracy"],
                -abs(best_bias),
            )
        )
        if key > old_key:
            best_bias, best_metrics = float(bias), current
    assert best_metrics is not None
    for bias in np.arange(best_bias - 0.1, best_bias + 0.1001, 0.01):
        pred, _ = choose_from_scores(records, score_rows, float(bias))
        current = relation_metrics(gold, pred)
        key = (current["positive_micro_f1"], current["accuracy"], -abs(float(bias)))
        old_key = (
            best_metrics["positive_micro_f1"],
            best_metrics["accuracy"],
            -abs(best_bias),
        )
        if key > old_key:
            best_bias, best_metrics = float(bias), current
    return best_bias, best_metrics


def constrained_choice_evaluation(
    active_model: torch.nn.Module,
    records: Sequence[Mapping[str, Any]],
    batch_size: int,
    step: int,
    prefix: str,
) -> Dict[str, Any]:
    old_padding = tokenizer.padding_side
    was_training = active_model.training
    tokenizer.padding_side = "left"
    active_model.eval()
    # Length bucketing preserves row order in the artifacts while avoiding the
    # large padding penalty caused by mixing short and long Arabic prompts.
    indexed_prompts: List[Tuple[int, Mapping[str, Any], List[int]]] = []
    for original_index, row in enumerate(records):
        prompt_ids = list(
            tokenizer.apply_chat_template(
                row["prompt_messages"], tokenize=True, add_generation_prompt=True
            )
        )
        if len(prompt_ids) > MAX_LENGTH:
            raise ValueError("Evaluation prompt exceeds max length")
        indexed_prompts.append((original_index, row, prompt_ids))
    indexed_prompts.sort(key=lambda item: len(item[2]))

    score_rows_by_index: List[Optional[List[float]]] = [None] * len(records)
    generated_token_ids_by_index: List[Optional[int]] = [None] * len(records)
    gold_nll: List[float] = []
    try:
        for start in tqdm(
            range(0, len(indexed_prompts), batch_size),
            desc=f"choice-eval:{prefix}",
            leave=False,
        ):
            current = indexed_prompts[start : start + batch_size]
            batch = tokenizer.pad(
                {"input_ids": [item[2] for item in current]},
                padding=True,
                return_tensors="pt",
            )
            batch = {key: value.to(active_model.device) for key, value in batch.items()}
            with torch.inference_mode():
                decoder, lm_head = decoder_and_lm_head(active_model)
                decoder_output = decoder(
                    **batch, use_cache=False, return_dict=True
                )
                decision_hidden = decoder_output.last_hidden_state[:, -1, :]
                full_next_logits = lm_head(decision_hidden).float()
                next_code_logits = full_next_logits.index_select(
                    1, torch.as_tensor(code_token_ids, device=decision_hidden.device)
                )
                generated_token_ids = full_next_logits.argmax(dim=-1)
            for batch_index, (original_index, row, _) in enumerate(current):
                option_count = len(row["option_codes"])
                candidates = next_code_logits[batch_index, :option_count]
                score_rows_by_index[original_index] = candidates.detach().cpu().tolist()
                generated_token_ids_by_index[original_index] = int(
                    generated_token_ids[batch_index].detach().cpu()
                )
                if row.get("gold_option_index") is not None:
                    gold_position = int(row["gold_option_index"])
                    gold_nll.append(
                        float(
                            (
                                torch.logsumexp(candidates, dim=-1)
                                - candidates[gold_position]
                            ).cpu()
                        )
                    )
    finally:
        tokenizer.padding_side = old_padding
        if was_training:
            active_model.train()

    if any(scores is None for scores in score_rows_by_index) or any(
        token_id is None for token_id in generated_token_ids_by_index
    ):
        raise RuntimeError("Constrained evaluation failed to score every row")
    score_rows = [scores for scores in score_rows_by_index if scores is not None]
    generated_token_ids_complete = [
        int(token_id) for token_id in generated_token_ids_by_index if token_id is not None
    ]

    gold = [str(row["gold_relation_full"]) for row in records]
    code_id_to_position = {token_id: index for index, token_id in enumerate(code_token_ids)}
    generated_pred: List[str] = []
    generated_valid: List[bool] = []
    for row, generated_token_id in zip(records, generated_token_ids_complete):
        position = code_id_to_position.get(generated_token_id)
        valid = position is not None and position < len(row["option_codes"])
        generated_valid.append(valid)
        generated_pred.append(
            str(row["allowed_relation_full_labels"][position])
            if valid and position is not None
            else NO_RELATION_FULL
        )
    generated_metrics = relation_metrics(gold, generated_pred)
    generated_metrics["invalid_outputs"] = int(sum(not valid for valid in generated_valid))
    generated_metrics["invalid_output_rate"] = float(
        generated_metrics["invalid_outputs"] / len(records)
    )
    raw_pred, raw_positions = choose_from_scores(records, score_rows, 0.0)
    raw_metrics = relation_metrics(gold, raw_pred)
    bias, calibrated_metrics = calibrate_bias(records, score_rows)
    final_pred, final_positions = choose_from_scores(records, score_rows, bias)
    predictions: List[Dict[str, Any]] = []
    for row, raw_scores, raw_position, final_position, generated_token_id, generated_is_valid in zip(
        records,
        score_rows,
        raw_positions,
        final_positions,
        generated_token_ids_complete,
        generated_valid,
    ):
        predictions.append(
            {
                "id": row["id"],
                "triple_id": row.get("triple_id"),
                "gold_relation_full": row.get("gold_relation_full"),
                "predicted_relation_full": row["allowed_relation_full_labels"][final_position],
                "predicted_code": row["option_codes"][final_position],
                "predicted_answer_ar": row["allowed_options_ar"][final_position],
                "raw_predicted_relation_full": row["allowed_relation_full_labels"][raw_position],
                "raw_predicted_code": row["option_codes"][raw_position],
                "candidate_logits": raw_scores,
                "no_relation_logit_bias": bias,
                "allowed_relation_full_labels": row["allowed_relation_full_labels"],
                "option_codes": row["option_codes"],
                "allowed_options_ar": row["allowed_options_ar"],
                "unconstrained_generated_token_id": generated_token_id,
                "unconstrained_generated_token": tokenizer.decode([generated_token_id]),
                "unconstrained_generated_valid": generated_is_valid,
            }
        )
    prediction_path = PREDICTIONS_DIR / f"{prefix}_choice_step_{step}.jsonl"
    dump_jsonl(predictions, prediction_path)
    result = {
        **calibrated_metrics,
        "decision_nll": float(np.mean(gold_nll)),
        "raw_accuracy": raw_metrics["accuracy"],
        "raw_macro_f1_41": raw_metrics["macro_f1_41"],
        "raw_weighted_f1_41": raw_metrics["weighted_f1_41"],
        "raw_positive_micro_precision": raw_metrics["positive_micro_precision"],
        "raw_positive_micro_recall": raw_metrics["positive_micro_recall"],
        "raw_positive_micro_f1": raw_metrics["positive_micro_f1"],
        "unconstrained_accuracy": generated_metrics["accuracy"],
        "unconstrained_macro_f1_41": generated_metrics["macro_f1_41"],
        "unconstrained_weighted_f1_41": generated_metrics["weighted_f1_41"],
        "unconstrained_positive_micro_precision": generated_metrics[
            "positive_micro_precision"
        ],
        "unconstrained_positive_micro_recall": generated_metrics[
            "positive_micro_recall"
        ],
        "unconstrained_positive_micro_f1": generated_metrics["positive_micro_f1"],
        "unconstrained_invalid_outputs": generated_metrics["invalid_outputs"],
        "unconstrained_invalid_output_rate": generated_metrics["invalid_output_rate"],
        "no_relation_logit_bias": bias,
        "global_step": step,
        "evaluated_rows": len(records),
        "inference": "constrained_one_token_greedy",
        "prediction_file": str(prediction_path),
    }
    dump_json(result, METRICS_DIR / f"{prefix}_choice_metrics_step_{step}.json")
    return result


stage("LOAD 4-BIT LOCAL YEHIA AND ATTACH QLORA")
quantization = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_use_double_quant=True,
    bnb_4bit_compute_dtype=torch.bfloat16,
)
model = AutoModelForCausalLM.from_pretrained(
    str(MODEL_DIR),
    local_files_only=True,
    quantization_config=quantization,
    torch_dtype=torch.bfloat16,
    device_map={"": 0},
    attn_implementation=ATTN_IMPL,
)
model.config.use_cache = False
model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=True)
try:
    model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
except TypeError:
    model.gradient_checkpointing_enable()
eligible_linear_inventory = enumerate_eligible_linear_modules(model)
model = get_peft_model(
    model,
    LoraConfig(
        r=LORA_R,
        lora_alpha=LORA_ALPHA,
        lora_dropout=LORA_DROPOUT,
        bias="none",
        task_type="CAUSAL_LM",
        target_modules=TARGET_MODULES,
        use_rslora=USE_RSLORA,
    ),
)
model.print_trainable_parameters()
lora_audit = validate_lora_wrapping(model, eligible_linear_inventory)
print(
    json.dumps(
        {
            "eligible_linear_or_conv1d": lora_audit["eligible_linear_or_conv1d_count"],
            "lora_wrapped": lora_audit["lora_wrapped_count"],
            "trainable_parameters": lora_audit["trainable_parameters"],
            "trainable_percentage": lora_audit["trainable_percentage"],
        },
        indent=2,
    ),
    flush=True,
)


def decoder_and_lm_head(active_model: torch.nn.Module):
    """Resolve the adapted decoder and frozen output projection from PEFT."""
    causal_model = (
        active_model.get_base_model()
        if hasattr(active_model, "get_base_model")
        else active_model
    )
    if not hasattr(causal_model, "model") or not hasattr(causal_model, "lm_head"):
        raise RuntimeError(f"Unexpected causal model structure: {type(causal_model)}")
    return causal_model.model, causal_model.lm_head


stage("FINITE FORWARD/BACKWARD SMOKE TEST")
long_indices = sorted(
    range(len(train_dataset)),
    key=lambda index: len(encoded["train"][index]["input_ids"]),
    reverse=True,
)[:TRAIN_BATCH]
torch.cuda.empty_cache()
torch.cuda.reset_peak_memory_stats()
smoke = collator([train_dataset[index] for index in long_indices])
if any(int((item["labels"] != -100).sum()) != 1 for item in [train_dataset[i] for i in long_indices]):
    raise RuntimeError("Smoke rows do not contain exactly one supervised token")
smoke = {key: value.to(model.device) for key, value in smoke.items()}
model.train()
smoke_output = model(
    input_ids=smoke["input_ids"],
    attention_mask=smoke["attention_mask"],
    labels=smoke["labels"],
)
smoke_loss = smoke_output.loss
if not torch.isfinite(smoke_loss):
    raise RuntimeError(f"Non-finite smoke loss: {smoke_loss}")
(smoke_loss / GRAD_ACCUM).backward()
finite_nonzero_grads = 0
for name, parameter in model.named_parameters():
    if parameter.grad is not None:
        if "lora_" not in name:
            raise RuntimeError(f"Frozen base parameter received a gradient: {name}")
        if torch.isfinite(parameter.grad).all() and float(parameter.grad.abs().sum()) > 0:
            finite_nonzero_grads += 1
if finite_nonzero_grads < 4:
    raise RuntimeError("Too few finite nonzero LoRA gradients")
smoke_artifact = {
    "status": "passed",
    "loss": float(smoke_loss.detach().cpu()),
    "loss_mode": "full_vocab",
    "finite_nonzero_lora_gradient_tensors": finite_nonzero_grads,
    "physical_batch": TRAIN_BATCH,
    "gradient_accumulation": GRAD_ACCUM,
    "peak_allocated_gib": torch.cuda.max_memory_allocated() / (1024 ** 3),
    "peak_reserved_gib": torch.cuda.max_memory_reserved() / (1024 ** 3),
}
dump_json(smoke_artifact, ARTIFACT_DIR / "smoke_test.json")
model.zero_grad(set_to_none=True)
del smoke, smoke_loss
del smoke_output
torch.cuda.empty_cache()
print(json.dumps(smoke_artifact, indent=2), flush=True)


stage("PRETRAINING FULL CONSTRAINED VALIDATION")
baseline_metrics = constrained_choice_evaluation(
    model, rows["validation"], EVAL_BATCH, step=0, prefix="validation"
)
print(json.dumps(baseline_metrics, ensure_ascii=False, indent=2), flush=True)


sample_weights = [
    float(class_counts[str(row["gold_relation_full"])]) ** (-CLASS_SAMPLING_ALPHA)
    for row in rows["train"]
]
sampling_summary: Dict[str, Any] = {
    "alpha": CLASS_SAMPLING_ALPHA,
    "minimum_weight": min(sample_weights),
    "maximum_weight": max(sample_weights),
    "class_weights": {
        label: count ** (-CLASS_SAMPLING_ALPHA)
        for label, count in sorted(class_counts.items())
    },
}
dump_json(sampling_summary, ARTIFACT_DIR / "sampling_strategy.json")


class ChoiceTrainer(Trainer):
    def __init__(
        self,
        *args: Any,
        generation_records: Sequence[Mapping[str, Any]],
        sampler_weights: Sequence[float],
        **kwargs: Any,
    ):
        super().__init__(*args, **kwargs)
        self.generation_records = list(generation_records)
        self.sampler_weights = torch.as_tensor(sampler_weights, dtype=torch.double)

    def _get_train_sampler(self, train_dataset=None):
        dataset = self.train_dataset if train_dataset is None else train_dataset
        generator = torch.Generator()
        generator.manual_seed(self.args.data_seed)
        return WeightedRandomSampler(
            self.sampler_weights,
            num_samples=len(dataset),
            replacement=True,
            generator=generator,
        )

    def evaluate(self, eval_dataset=None, ignore_keys=None, metric_key_prefix="eval"):
        step = int(self.state.global_step)
        choice_metrics = constrained_choice_evaluation(
            self.model,
            self.generation_records,
            EVAL_BATCH,
            step=step,
            prefix="validation",
        )
        metrics = {
            f"{metric_key_prefix}_{key}": value
            for key, value in choice_metrics.items()
            if isinstance(value, (int, float))
        }
        # This NLL is exactly the decision-token objective used in training.
        metrics[f"{metric_key_prefix}_loss"] = choice_metrics["decision_nll"]
        self.log(metrics)
        self.control = self.callback_handler.on_evaluate(
            self.args, self.state, self.control, metrics
        )
        print("CHOICE_METRICS " + json.dumps(metrics, ensure_ascii=False), flush=True)
        return metrics


class NamedCheckpointCallback(TrainerCallback):
    """Maintain complete resumable named copies after every scheduled save."""

    def on_save(self, args, state, control, **kwargs):
        checkpoint = Path(args.output_dir) / f"checkpoint-{state.global_step}"
        if not checkpoint.is_dir():
            raise RuntimeError(f"Expected checkpoint was not written: {checkpoint}")

        latest = OUTPUT_DIR / "latest_checkpoint"
        if latest.exists():
            shutil.rmtree(latest)
        shutil.copytree(checkpoint, latest)

        if state.best_model_checkpoint:
            best_source = Path(state.best_model_checkpoint)
            if not best_source.is_dir():
                raise RuntimeError(f"Best checkpoint is missing: {best_source}")
            best_named = OUTPUT_DIR / "best_checkpoint"
            if best_named.exists():
                shutil.rmtree(best_named)
            shutil.copytree(best_source, best_named)
        return control


stage("TRAIN")
training_kwargs: Dict[str, Any] = {
    "output_dir": str(CHECKPOINT_DIR),
    "num_train_epochs": NUM_EPOCHS,
    "max_steps": MAX_STEPS,
    "per_device_train_batch_size": TRAIN_BATCH,
    "per_device_eval_batch_size": EVAL_BATCH,
    "gradient_accumulation_steps": GRAD_ACCUM,
    "learning_rate": LEARNING_RATE,
    "warmup_ratio": WARMUP_RATIO,
    "weight_decay": WEIGHT_DECAY,
    "lr_scheduler_type": "cosine",
    "optim": "paged_adamw_8bit",
    "logging_steps": LOGGING_STEPS,
    "save_strategy": "steps",
    "save_steps": EVAL_STEPS,
    "eval_steps": EVAL_STEPS,
    "save_total_limit": SAVE_TOTAL_LIMIT,
    "load_best_model_at_end": True,
    "metric_for_best_model": "positive_micro_f1",
    "greater_is_better": True,
    "bf16": True,
    "fp16": False,
    "max_grad_norm": MAX_GRAD_NORM,
    "report_to": "none",
    "remove_unused_columns": False,
    "seed": SEED,
    "data_seed": SEED,
    "gradient_checkpointing": True,
    "ddp_find_unused_parameters": False,
}
signature = inspect.signature(TrainingArguments.__init__).parameters
training_kwargs["eval_strategy" if "eval_strategy" in signature else "evaluation_strategy"] = "steps"
training_args = TrainingArguments(**training_kwargs)
trainer_kwargs: Dict[str, Any] = {
    "model": model,
    "args": training_args,
    "train_dataset": train_dataset,
    "eval_dataset": validation_dataset,
    "data_collator": collator,
    "generation_records": rows["validation"],
    "sampler_weights": sample_weights,
    "callbacks": [NamedCheckpointCallback()],
}
trainer_signature = inspect.signature(Trainer.__init__).parameters
trainer_kwargs["processing_class" if "processing_class" in trainer_signature else "tokenizer"] = tokenizer
trainer = ChoiceTrainer(**trainer_kwargs)

last_checkpoint = get_last_checkpoint(str(CHECKPOINT_DIR)) if RESUME else None
print(f"Resume checkpoint: {last_checkpoint}", flush=True)
train_result = trainer.train(resume_from_checkpoint=last_checkpoint)
trainer.save_state()


stage("SAVE BEST ADAPTER AND FINAL METRICS")
best_dir = OUTPUT_DIR / "best_adapter"
best_dir.mkdir(parents=True, exist_ok=True)
trainer.model.save_pretrained(best_dir, safe_serialization=True)
tokenizer.save_pretrained(best_dir)
final_metrics = constrained_choice_evaluation(
    trainer.model,
    rows["validation"],
    EVAL_BATCH,
    step=int(trainer.state.global_step),
    prefix="final_validation",
)
summary = {
    "status": "complete",
    "finished_at_utc": utc_now(),
    "global_step": int(trainer.state.global_step),
    "best_metric": trainer.state.best_metric,
    "best_model_checkpoint": trainer.state.best_model_checkpoint,
    "train_metrics": train_result.metrics,
    "baseline_metrics": baseline_metrics,
    "final_metrics": final_metrics,
    "run_config": config,
    "lora_audit": {
        key: value for key, value in lora_audit.items() if key != "wrapped_modules"
    },
    "best_adapter": str(best_dir),
}
dump_json(summary, OUTPUT_DIR / "run_summary.json")
print(json.dumps(summary, ensure_ascii=False, indent=2), flush=True)