yazidbou/IslamicHallucination

from __future__ import annotations

import argparse
import json
import re
import time
from pathlib import Path
from typing import Any

import torch
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer


DEFAULT_BASE_MODEL = "humain-ai/ALLaM-7B-Instruct-preview"
DEFAULT_ADAPTER_ID = "Lamine92500/IslamicHallucination"
DEFAULT_DTYPE = "bfloat16"
DEFAULT_DEVICE_MAP = "auto"
DEFAULT_BATCH_SIZE = 1
DEFAULT_MAX_INPUT_TOKENS = 2048
DEFAULT_MAX_NEW_TOKENS = 384
DEFAULT_SAVE_EVERY = 2
VALID_OPTIONS = ("A", "B", "C", "D", "E", "F")

SYSTEM_PROMPT = """أنت نظام للتحقق من صحة إجابات النماذج اللغوية باللغة العربية.
مهمتك:
1) تحديد هل تحتوي الإجابة المولدة على هلوسة أم لا.
2) اختيار الإجابة الصحيحة من الخيارات A-F.
3) عند وجود هلوسة، تحديد المقاطع الخاطئة وشرح سبب الخطأ.

أجب بصيغة JSON فقط بدون أي نص إضافي."""


AR_DIACRITICS = re.compile(r"[\u0610-\u061A\u064B-\u065F\u0670\u06D6-\u06ED]")


def looks_like_examples(payload: Any) -> bool:
    if not isinstance(payload, list) or not payload or not isinstance(payload[0], dict):
        return False
    first = payload[0]
    return "id" in first and "question" in first and "generated_answer" in first


def load_json_if_examples(path: Path) -> list[dict[str, Any]] | None:
    try:
        payload = json.loads(path.read_text(encoding="utf-8"))
    except Exception:
        return None
    if not looks_like_examples(payload):
        return None
    return payload


def candidate_score(path: Path, examples: list[dict[str, Any]]) -> tuple[int, int]:
    text = str(path).lower()
    score = 0
    if "test" in text:
        score += 1000
    if "islamic" in text:
        score += 500
    if "prediction" in text or "output" in text or "result" in text:
        score -= 1000
    if "gold" in text or "train" in text or "dev" in text:
        score -= 100
    return score, len(examples)


def find_examples_file(prediction_path: Path, full_prediction_path: Path) -> tuple[Path, list[dict[str, Any]]]:
    excluded = {
        prediction_path.resolve(),
        full_prediction_path.resolve(),
    }
    candidates: list[tuple[tuple[int, int], Path, list[dict[str, Any]]]] = []
    for path in sorted(Path.cwd().rglob("*.json")):
        if path.resolve() in excluded:
            continue
        examples = load_json_if_examples(path)
        if examples is not None:
            candidates.append((candidate_score(path, examples), path, examples))

    if candidates:
        _, path, examples = max(candidates, key=lambda item: item[0])
        return path, examples

    raise FileNotFoundError(
        "No input JSON examples found in the current folder. "
        "Run this script from the folder that contains the test JSON."
    )


def save_json(path: str | Path, payload: Any) -> None:
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")


def dtype_from_name(name: str) -> torch.dtype:
    normalized = name.strip().lower().replace("torch.", "")
    if normalized in {"bf16", "bfloat16"}:
        return torch.bfloat16 if torch.cuda.is_available() else torch.float32
    if normalized in {"fp16", "float16", "half"}:
        return torch.float16 if torch.cuda.is_available() else torch.float32
    if normalized in {"fp32", "float32", "float"}:
        return torch.float32
    raise ValueError(f"Unsupported dtype: {name!r}")


def normalize_arabic(text: Any) -> str:
    text = "" if text is None else str(text)
    text = AR_DIACRITICS.sub("", text)
    text = text.replace("ـ", "")
    text = re.sub("[إأآٱ]", "ا", text)
    text = text.replace("ى", "ي")
    return re.sub(r"\s+", " ", text).strip()


def format_options(options: dict[str, Any]) -> str:
    lines = []
    for key in sorted(options):
        option = str(key).strip().upper()
        if option in VALID_OPTIONS:
            lines.append(f"{option}. {str(options[key]).strip()}")
    return "\n".join(lines)


def build_user_prompt(example: dict[str, Any]) -> str:
    gold_answer = str(example.get("gold_answer", "")).strip()
    reference_block = f"\n\nالإجابة المرجعية الصحيحة:\n{gold_answer}" if gold_answer else ""
    options_text = format_options(dict(example.get("options", {})))
    return f"""السؤال:
{str(example.get("question", "")).strip()}

الإجابة المولدة:
{str(example.get("generated_answer", "")).strip()}{reference_block}

الخيارات:
{options_text}

أعد JSON فقط بالحقول التالية:
- predicted_label: no-hallucinate أو hallucination
- predicted_answer: الخيار الصحيح من A-F
- selected_option: نفس الخيار الصحيح
- justification: شرح مختصر ودقيق
- hallucinations: قائمة المقاطع الخاطئة إن وجدت"""


def fallback_chat_template(messages: list[dict[str, str]], *, add_generation_prompt: bool) -> str:
    role_names = {
        "system": "النظام",
        "user": "المستخدم",
        "assistant": "المساعد",
    }
    parts = []
    for message in messages:
        role = role_names.get(message["role"], message["role"])
        parts.append(f"{role}:\n{message['content'].strip()}")
    if add_generation_prompt:
        parts.append("المساعد:\n")
    return "\n\n".join(parts).rstrip() + "\n"


def render_prompt(tokenizer, example: dict[str, Any]) -> str:
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": build_user_prompt(example)},
    ]
    try:
        return tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True,
        )
    except Exception:
        return fallback_chat_template(messages, add_generation_prompt=True)


def normalize_label(value: Any) -> str:
    text = str(value or "").strip().lower().replace("-", "_").replace(" ", "_")
    if text in {"hallucination", "hallucinated", "yes", "true", "1"}:
        return "hallucination"
    return "no_hallucination"


def normalize_option(value: Any, option_keys: list[str]) -> str | None:
    if value is None:
        return None
    text = str(value).strip().upper()
    if text in option_keys:
        return text
    match = re.search(r"\b([A-F])\b", text)
    if match and match.group(1) in option_keys:
        return match.group(1)
    return None


def option_from_reference(example: dict[str, Any]) -> str:
    reference = normalize_arabic(example.get("gold_answer", ""))
    options = example.get("options", {})
    if not isinstance(options, dict) or not options:
        return "A"

    reference_tokens = set(reference.split())
    scored: list[tuple[float, str]] = []
    for option in VALID_OPTIONS:
        option_text = normalize_arabic(options.get(option, ""))
        option_tokens = set(option_text.split())
        if not option_tokens:
            score = 0.0
        elif reference and reference in option_text:
            score = 1.0
        else:
            score = len(reference_tokens & option_tokens) / max(1, len(reference_tokens | option_tokens))
        scored.append((score, option))
    scored.sort(reverse=True)
    return scored[0][1]


def extract_json_object(text: str) -> dict[str, Any]:
    match = re.search(r"\{.*\}", text, flags=re.DOTALL)
    if not match:
        raise ValueError("No JSON object found in model output")
    return json.loads(match.group(0))


def parse_prediction(example: dict[str, Any], raw_output: str) -> dict[str, Any]:
    options = example.get("options", {})
    option_keys = [str(key).strip().upper() for key in options] if isinstance(options, dict) else list(VALID_OPTIONS)
    option_keys = [key for key in option_keys if key in VALID_OPTIONS] or list(VALID_OPTIONS)

    try:
        parsed = extract_json_object(raw_output)
    except Exception as exc:
        parsed = {
            "predicted_label": "hallucination",
            "selected_option": None,
            "justification": f"Could not parse model output: {exc}",
            "hallucinations": [],
        }

    selected_option = (
        normalize_option(parsed.get("selected_option"), option_keys)
        or normalize_option(parsed.get("predicted_answer"), option_keys)
        or normalize_option(parsed.get("pred_option"), option_keys)
        or normalize_option(parsed.get("answer"), option_keys)
        or normalize_option(raw_output, option_keys)
        or option_from_reference(example)
    )
    pred_label = normalize_label(parsed.get("predicted_label", parsed.get("pred_label")))

    return {
        "id": str(example.get("id", "")),
        "pred_label": pred_label,
        "pred_option": selected_option,
        "prediction_label": pred_label,
        "answer": selected_option,
        "justification": str(parsed.get("justification", "")),
        "hallucinations": parsed.get("hallucinations", []),
        "raw_output": raw_output,
    }


def model_device(model) -> torch.device:
    return next(model.parameters()).device


def generate_batch(
    model,
    tokenizer,
    prompts: list[str],
    *,
    max_input_tokens: int,
    max_new_tokens: int,
) -> list[str]:
    encoded = tokenizer(
        prompts,
        return_tensors="pt",
        padding=True,
        truncation=True,
        max_length=max_input_tokens,
    )
    encoded.pop("token_type_ids", None)
    encoded = {key: value.to(model_device(model)) for key, value in encoded.items()}
    with torch.inference_mode():
        generated = model.generate(
            **encoded,
            max_new_tokens=max_new_tokens,
            do_sample=False,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id,
            use_cache=True,
        )
    generated = generated[:, encoded["input_ids"].shape[1] :]
    return tokenizer.batch_decode(generated, skip_special_tokens=True)


def load_existing(path: Path) -> dict[str, dict[str, Any]]:
    if not path.exists():
        return {}
    try:
        data = json.loads(path.read_text(encoding="utf-8"))
    except Exception:
        return {}
    if not isinstance(data, list):
        return {}
    return {
        str(row["id"]): row
        for row in data
        if isinstance(row, dict) and row.get("id") is not None
    }


def save_outputs(short_path: Path, full_path: Path, rows: list[dict[str, Any]]) -> None:
    short_rows = [
        {
            "id": row["id"],
            "pred_label": row["pred_label"],
            "pred_option": row["pred_option"],
        }
        for row in rows
    ]
    save_json(short_path, short_rows)
    save_json(full_path, rows)


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Generate predictions with IslamicHallu Badr v1.1 LoRA from Hugging Face.",
    )
    parser.add_argument("prediction_file", help="Output prediction JSON path.")
    args = parser.parse_args()

    short_path = Path(args.prediction_file)
    full_path = short_path.with_name(f"{short_path.stem}_full{short_path.suffix}")
    input_path, examples = find_examples_file(short_path, full_path)

    print(f"Adapter: {DEFAULT_ADAPTER_ID}")
    print(f"Base model: {DEFAULT_BASE_MODEL}")
    print(f"Input: {input_path} ({len(examples)} examples)")

    tokenizer = AutoTokenizer.from_pretrained(DEFAULT_ADAPTER_ID, trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "left"
    tokenizer.truncation_side = "left"

    base = AutoModelForCausalLM.from_pretrained(
        DEFAULT_BASE_MODEL,
        torch_dtype=dtype_from_name(DEFAULT_DTYPE),
        device_map=DEFAULT_DEVICE_MAP if torch.cuda.is_available() else None,
        trust_remote_code=True,
    )
    model = PeftModel.from_pretrained(base, DEFAULT_ADAPTER_ID)
    if hasattr(model.config, "use_cache"):
        model.config.use_cache = True
    model.eval()

    existing = load_existing(full_path)
    allowed_ids = {str(example.get("id", "")) for example in examples}
    row_by_id = {row_id: row for row_id, row in existing.items() if row_id in allowed_ids}
    pending = [example for example in examples if str(example.get("id", "")) not in row_by_id]

    start_time = time.time()
    batch_size = max(1, DEFAULT_BATCH_SIZE)
    for start in range(0, len(pending), batch_size):
        batch = pending[start : start + batch_size]
        ids = ", ".join(str(example.get("id", "")) for example in batch)
        print(f"[{start + 1}-{start + len(batch)}/{len(pending)}] Predicting {ids}", flush=True)

        prompts = [render_prompt(tokenizer, example) for example in batch]
        outputs = generate_batch(
            model,
            tokenizer,
            prompts,
            max_input_tokens=DEFAULT_MAX_INPUT_TOKENS,
            max_new_tokens=DEFAULT_MAX_NEW_TOKENS,
        )
        for example, raw_output in zip(batch, outputs):
            row_by_id[str(example.get("id", ""))] = parse_prediction(example, raw_output)

        if len(row_by_id) % max(1, DEFAULT_SAVE_EVERY) == 0:
            ordered_rows = [
                row_by_id[str(example.get("id", ""))]
                for example in examples
                if str(example.get("id", "")) in row_by_id
            ]
            save_outputs(short_path, full_path, ordered_rows)

    rows = [row_by_id[str(example.get("id", ""))] for example in examples]
    save_outputs(short_path, full_path, rows)
    print(f"Saved short predictions: {short_path}")
    print(f"Saved full predictions: {full_path}")
    print(f"Elapsed prediction time: {time.time() - start_time:.1f}s")


if __name__ == "__main__":
    main()