#!/usr/bin/env python3 """Run constrained one-token inference with a staged DRU-RE-Yehia adapter.""" from __future__ import annotations import argparse import json import os from pathlib import Path from typing import Any, Dict, Iterable, List, Mapping os.environ.setdefault("TOKENIZERS_PARALLELISM", "false") os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0" import torch from dotenv import load_dotenv from huggingface_hub import snapshot_download from peft import PeftModel from tqdm.auto import tqdm from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig from re_sft_common import ( OPTION_CODES, env_int, env_str, load_jsonl, validate_local_hf_revision, ) ROOT = Path(__file__).resolve().parent load_dotenv(ROOT / ".env") BASE_MODEL_ID = env_str("YEHIA_BASE_MODEL_ID", "Navid-AI/Yehia-7B-preview") BASE_MODEL_REVISION = env_str( "YEHIA_BASE_MODEL_REVISION", "b9dda4715eafee7e8090d2c83cfe078d75f4ebb8" ) def repository_path(value: str) -> Path: path = Path(value) return path if path.is_absolute() else ROOT / path def require_token() -> str: token = os.environ.get("HF_TOKENONE", "").strip() if not token or token.startswith("hf_your_"): raise RuntimeError("HF_TOKENONE is required because the local Yehia snapshot is absent") return token def resolve_base_model() -> Path: model_dir = repository_path( env_str("LOCAL_YEHIA_MODEL_DIR", "models/Yehia-7B-preview") ) if (model_dir / "config.json").is_file(): validate_local_hf_revision(model_dir, BASE_MODEL_REVISION) return model_dir offline = os.environ.get("HF_HUB_OFFLINE", "").strip().lower() if offline in {"1", "true", "yes", "on"}: raise RuntimeError("Yehia is missing locally while HF_HUB_OFFLINE is enabled") model_dir.mkdir(parents=True, exist_ok=True) snapshot_download( repo_id=BASE_MODEL_ID, revision=BASE_MODEL_REVISION, local_dir=str(model_dir), token=require_token(), ) if not (model_dir / "config.json").is_file(): raise RuntimeError(f"Incomplete Yehia snapshot: {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) return model_dir def resolve_adapter() -> Path: configured = repository_path(env_str("ADAPTER_DIR", ".")) candidates = [configured, ROOT, ROOT / "runs" / "DRU-RE-Yehia" / "best_adapter"] for candidate in candidates: if (candidate / "adapter_config.json").is_file() and ( candidate / "adapter_model.safetensors" ).is_file(): return candidate raise FileNotFoundError( "No staged adapter found. Run tools/stage_release.py after training completes " "or set ADAPTER_DIR to a PEFT adapter directory." ) def selected_bias(adapter_dir: Path, override: float | None) -> float: if override is not None: return override for path in (adapter_dir / "inference_config.json", ROOT / "inference_config.json"): if path.is_file(): return float(json.loads(path.read_text(encoding="utf-8"))["no_relation_logit_bias"]) print("WARNING: no inference_config.json found; using no-relation bias 0.0") return 0.0 def batches(rows: List[Dict[str, Any]], size: int) -> Iterable[List[Dict[str, Any]]]: for start in range(0, len(rows), size): yield rows[start : start + size] def main() -> None: parser = argparse.ArgumentParser() parser.add_argument( "--input", default=env_str("CHOICE_DATASET_DIR", "data/Yehia-RE-SFT") + "/official.jsonl", ) parser.add_argument("--output", default="predictions/official_predictions.jsonl") parser.add_argument("--batch-size", type=int, default=env_int("INFERENCE_BATCH_SIZE", 16)) parser.add_argument("--limit", type=int, default=0) parser.add_argument("--no-relation-bias", type=float, default=None) args = parser.parse_args() input_path = repository_path(args.input) output_path = repository_path(args.output) rows = load_jsonl(input_path) if args.limit > 0: rows = rows[: args.limit] if not rows: raise RuntimeError(f"No rows found in {input_path}") base_dir = resolve_base_model() adapter_dir = resolve_adapter() bias = selected_bias(adapter_dir, args.no_relation_bias) 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" 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])) quantization = BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_use_double_quant=True, bnb_4bit_compute_dtype=torch.bfloat16, ) base = AutoModelForCausalLM.from_pretrained( str(base_dir), local_files_only=True, quantization_config=quantization, torch_dtype=torch.bfloat16, device_map={"": 0}, attn_implementation=env_str("ATTENTION_IMPLEMENTATION", "sdpa"), ) model = PeftModel.from_pretrained(base, str(adapter_dir), is_trainable=False) model.eval() output_path.parent.mkdir(parents=True, exist_ok=True) with output_path.open("w", encoding="utf-8") as handle, torch.inference_mode(): for group in tqdm(list(batches(rows, args.batch_size)), desc="predict"): sequences = [ list( tokenizer.apply_chat_template( row["prompt_messages"], tokenize=True, add_generation_prompt=True ) ) for row in group ] lengths = torch.tensor([len(sequence) for sequence in sequences], dtype=torch.long) max_length = int(lengths.max().item()) input_ids = torch.full( (len(group), max_length), tokenizer.pad_token_id, dtype=torch.long ) attention_mask = torch.zeros((len(group), max_length), dtype=torch.long) for index, sequence in enumerate(sequences): input_ids[index, : len(sequence)] = torch.tensor(sequence, dtype=torch.long) attention_mask[index, : len(sequence)] = 1 input_ids = input_ids.to(model.device) attention_mask = attention_mask.to(model.device) logits = model(input_ids=input_ids, attention_mask=attention_mask).logits next_logits = logits[ torch.arange(len(group), device=logits.device), lengths.to(logits.device) - 1, ] for index, row in enumerate(group): options = list(row["allowed_options_ar"]) labels = list(row["allowed_relation_full_labels"]) codes = list(row["option_codes"]) if not options or options[-1] != "لا توجد علاقة": raise RuntimeError(f"Malformed option list for {row.get('id')}") candidate_ids = torch.tensor( code_token_ids[: len(options)], device=logits.device ) scores = next_logits[index, candidate_ids].float().clone() scores[-1] += bias chosen = int(torch.argmax(scores).item()) record: Dict[str, Any] = { "id": row.get("id"), "sentence_id": row.get("sentence_id"), "triple_id": row.get("triple_id"), "predicted_option_index": chosen, "predicted_code": codes[chosen], "predicted_option_ar": options[chosen], "predicted_relation_full": labels[chosen], "no_relation_logit_bias": bias, } ontology_ids = row.get("allowed_relation_ontology_ids") if isinstance(ontology_ids, list) and chosen < len(ontology_ids): record["predicted_relation_ontology_id"] = ontology_ids[chosen] handle.write(json.dumps(record, ensure_ascii=False) + "\n") print(json.dumps({"rows": len(rows), "output": str(output_path), "bias": bias}, indent=2)) if __name__ == "__main__": main()