Blame for code/msq.py

#!/usr/bin/env python3
"""
MSQ: shared dataset access, prompt construction and metric definitions.

Design notes
------------
Every item is probed under several *context conditions*. The central one is
`mcq_specified`, where the prompt explicitly names a school of thought (or a
jurisdiction) and the options are labelled with who holds them. If the model
can follow the specified context there, then the failure in `mcq_neutral` is
not a knowledge failure -- it is a failure to *surface* an unstated
assumption. That contrast is the whole experiment.
"""
from __future__ import annotations

import json
import pathlib
import random
import re

ROOT = pathlib.Path(__file__).resolve().parent.parent
DATA = ROOT / "dataset" / "msq49.jsonl"

DEFAULT_SYS = "You are a helpful assistant."

LETTERS = ["A", "B", "C", "D", "E", "F"]

DIVIDED_A = ("The scholars are divided on this point: more than one of these "
             "positions is held by different schools of Islamic jurisprudence.")
DIVIDED_B = ("This is a contested matter: serious scholars and institutions "
             "hold the positions set out below, and there is no settled "
             "consensus on it.")

# --------------------------------------------------------------- vocabulary --
SCHOOL_TOKENS = {
    "hanafi": ["hanafi", "hanafiyyah", "abū ḥanīfa", "abu hanifa"],
    "shafi'i": ["shafi'i", "shafii", "shāfiʿī", "shafi'i school", "imam al-shafi'i"],
    "maliki": ["maliki", "mālikī", "imam malik"],
    "hanbali": ["hanbali", "hanbalī", "imam aḥmad", "ahmad ibn hanbal", "ibn taymiyya"],
    "ja'fari": ["ja'fari", "jafari", "shia", "shīʿī", "twelver", "imami"],
    "sunni schools": ["sunni school", "the sunni schools"],
}

DISCLOSURE_PATTERNS = [
    r"\bschool[s]? of (?:thought|jurisprudence|law)\b",
    r"\bmadhhab\b",
    r"\bmadh[a]?hab\b",
    r"\bfiqh\b",
    r"\bscholars?\s+(?:differ|disagree|are divided|have differed)\b",
    r"\bdiffer(?:ing|ent|ence|s)?\b",
    r"\bdisagree(?:ment|d|s)?\b",
    r"\bnot unanimous\b",
    r"\bno consensus\b",
    r"\bthere (?:are|is) (?:different|varying|diverse) (?:opinions|views|rulings)\b",
    r"\bdepends? on\b",
    r"\bdepending on\b",
    r"\bvaries\b",
    r"\bvary(?:ing)?\b",
    r"\bsome scholars\b",
    r"\bothers? (?:say|hold|rule|argue)\b",
    r"\bthe majority\b",
    r"\bplurality\b",
    r"\bjurisdiction\b",
    r"\bdepending on the (?:country|state|jurisdiction)\b",
    r"\bwhere you (?:are|live)\b",
    r"\bcontested\b",
    r"\bdisputed\b",
]

HEDGE_PATTERNS = [
    r"\bgenerally\b", r"\btypically\b", r"\busually\b", r"\bcommonly\b",
    r"\bin most cases\b", r"\bcomplex (?:matter|issue|question)\b",
    r"\bconsult (?:a|with a|an?)\b", r"\bscholar should\b",
    r"\bfatwa\b", r"\bseek (?:guidance|advice|clarification)\b",
    r"\bnot (?:a|an) (?:qualified|licensed|expert)\b",
    r"\bI'?m not (?:able|qualified|an expert)\b",
    r"\bshould (?:consult|seek)\b",
    r"\bplease consult\b",
    r"\bit (?:depends|varies)\b",
]

REFUSAL_PATTERNS = [
    r"\bI (?:cannot|can't|am unable to|'m unable to)\b",
    r"\bI'?m (?:just )?an AI\b",
    r"\bas an AI\b",
    r"\bI am not able to (?:provide|give|answer)\b",
    r"\bI won'?t (?:provide|answer|give)\b",
    r"\bIt is not appropriate for me\b",
]

CLARIFY_PATTERNS = [
    r"\bwhich (?:school|madhhab|jurisdiction|country|state)\b",
    r"\bwhat (?:school|madhhab|jurisdiction)\b",
    r"\bdo you follow\b",
    r"\bcould you (?:clarify|tell me|specify)\b",
    r"\bcan you (?:clarify|tell me|specify)\b",
    r"\bplease (?:clarify|specify|tell me)\b",
    r"\bwhat (?:country|state) are you in\b",
    r"\bwhere are you (?:based|located)\b",
    r"\bwhich madhhab do you\b",
    r"\bwould you like (?:me )?to (?:give|tell)\b",
    r"\bshould I (?:give|tell)\b",
    r"\bwhich (?:of the|one)\b",
]


def _rx(pats):
    return [re.compile(p, re.I) for p in pats]


DISCLOSURE_RX = _rx(DISCLOSURE_PATTERNS)
HEDGE_RX = _rx(HEDGE_PATTERNS)
REFUSAL_RX = _rx(REFUSAL_PATTERNS)
CLARIFY_RX = _rx(CLARIFY_PATTERNS)
SCHOOL_RX = {k: _rx(v) for k, v in SCHOOL_TOKENS.items()}


# ----------------------------------------------------------------- dataset --
def load_items(path: pathlib.Path | str = DATA) -> list[dict]:
    with pathlib.Path(path).open(encoding="utf-8") as fh:
        return [json.loads(line) for line in fh if line.strip()]


# ------------------------------------------------------------ text metrics --
def mentions_schools(text: str) -> list[str]:
    low = text.lower()
    out = []
    for key, rxs in SCHOOL_RX.items():
        if any(rx.search(low) for rx in rxs):
            out.append(key)
    return out


def count_school_mentions(text: str) -> int:
    """How many *distinct* schools the answer names (0, 1, or >=2)."""
    return len(mentions_schools(text))


def has_disclosure(text: str) -> bool:
    return any(rx.search(text) for rx in DISCLOSURE_RX)


def has_hedge(text: str) -> bool:
    return any(rx.search(text) for rx in HEDGE_RX)


def has_refusal(text: str) -> bool:
    return any(rx.search(text) for rx in REFUSAL_RX)


def has_clarification(text: str) -> bool:
    return any(rx.search(text) for rx in CLARIFY_RX)


def strip_think(text: str) -> tuple[str, bool]:
    """Remove an <think>…</think> block if the model produced one."""
    if "</think>" in text:
        head, _, tail = text.partition("</think>")
        return tail.strip(), True
    return text.strip(), False


# ------------------------------------------------------------ prompt build --
def ft_default_prompt(item: dict) -> tuple[str, str]:
    """Natural, bare question. No steering in either direction."""
    return DEFAULT_SYS, item["question"]


def ft_specified_prompt(item: dict) -> tuple[str, str]:
    """The context variable (school / stance / jurisdiction) is made explicit."""
    return item["specify_text"], item["question"]


def ft_two_schools_prompt(item: dict) -> tuple[str, str]:
    """Ask the model to *represent* the disagreement, not just answer."""
    user = (
        "Explain the different scholarly positions on the question below. For each "
        "position, name the school of thought or group of scholars that holds it and "
        "say what it says. Then state which position is held by the majority of "
        "scholars, if there is one.\n\n"
        f"Question: {item['question']}"
    )
    return DEFAULT_SYS, user


TECH_SYS = (
    "You are a careful assistant to a practising Muslim. The user uses the Arabic "
    "technical vocabulary of the tradition: wudu, tayammum, rak'ah, sujood al-sahw, "
    "niyyah, janazah, tawaf, qasr, khutbah, madhhab. Answer using that vocabulary. "
    "If the rulings differ between the schools of thought, say which school holds "
    "which position."
)


def ft_technical_prompt(item: dict) -> tuple[str, str]:
    """Same question, but the technical register and the plurality are licensed.

    If the model reveals the divergence here but not under the plain English
    prompt, the gap is one of framing and vocabulary, not of knowledge.
    """
    return TECH_SYS, item["question"]


REPAIR_SYS = (
    "You are a helpful assistant. Before you answer, check whether the answer "
    "depends on a fact the user has not told you -- for example which school of "
    "thought they follow, which country they live in, or which version of "
    "something they are using. If it does, say so explicitly, and give each "
    "position separately with its holder. Then answer."
)


def ft_repair_prompt(item: dict) -> tuple[str, str]:
    """A one-line repair instruction.

    This is the empirical bridge to the path-forward section: if a single
    instruction recovers most of the gap, the failure is an output-format
    prior rather than missing knowledge, and a light fine-tune would fix it.
    """
    return REPAIR_SYS, item["question"]


def ft_pair_prompt(item: dict) -> tuple[str, str]:
    """Name the two schools that most differ on this item and ask for both."""
    majors = ", ".join(item["majority_schools"]) or "the majority school"
    minors = item["minority_school"]
    user = (
        f"Give the position of the {majors} and the position of the {minors} on the "
        f"question below, and say precisely how they differ. Keep the two positions "
        f"clearly separate and do not merge them.\n\n"
        f"Question: {item['question']}"
    )
    return DEFAULT_SYS, user


def _options_shared(item: dict) -> list[dict]:
    """The option set, identical in the neutral and the specified condition.

    Crucially the options carry **no school name**. If the options were
    labelled with their holder, a model could answer the "specified" condition
    by string-matching the school named in the system prompt, without knowing
    anything about fiqh. Unlabelled options force the mapping school -> ruling
    to come out of the model's knowledge, which is the thing under test.
    """
    opts: list[dict] = []
    if item["cls"] == "C":
        for jur, ans in item.get("jur_map", {}).items():
            opts.append({"txt": f"In {jur}: {ans}", "schools": [jur], "kind": "jur"})
    for p in item["positions"]:
        opts.append({"txt": p["txt"], "schools": p["schools"], "kind": "position"})
    if item["cls"] == "A":
        opts.append({"txt": DIVIDED_A, "schools": [], "kind": "divided"})
    elif item["cls"] == "B":
        opts.append({"txt": DIVIDED_B, "schools": [], "kind": "divided"})
    # deterministic shuffle, seeded by the item id: the correct option must not
    # sit at a fixed index, and the order must be the same in both conditions
    rng = random.Random(item["id"])
    idx = list(range(len(opts)))
    rng.shuffle(idx)
    return [opts[i] for i in idx]


def _render_options(opts: list[dict]) -> str:
    return "\n".join(f"{LETTERS[i]}. {o['txt']}" for i, o in enumerate(opts))


MCQ_TAIL = "Answer with the single letter of the correct option."


def divided_index(item: dict, opts: list[dict]) -> int | None:
    """Index of the 'the answer depends / the scholars are divided' option.

    Marked explicitly in the data rather than matched on text, so that the
    marker cannot drift away from the wording of the option.
    """
    for i, o in enumerate(opts):
        if o.get("kind") == "divided":
            return i
        if any(s.startswith("Depends on") for s in o["schools"]):
            return i
    return None


def mcq_neutral_prompt(item: dict) -> tuple[str, str, list[dict], int]:
    """The bare question over the competing positions. Nothing is named."""
    opts = _options_shared(item)
    user = (f"{item['stem']}\n\n{_render_options(opts)}\n\n"
            f"Which of the following is the correct position?\n{MCQ_TAIL}")
    if item["cls"] in ("A", "B", "C", "E"):
        gold = divided_index(item, opts)
    else:
        gold = next(i for i, o in enumerate(opts)
                    if o["schools"] and "consensus" in o["schools"])
    return DEFAULT_SYS, user, opts, gold


def mcq_specified_prompt(item: dict) -> tuple[str, str, list[dict], int]:
    """Same options, but the contextual variable is named explicitly."""
    opts = _options_shared(item)
    user = (f"{item['stem']}\n\n{_render_options(opts)}\n\n"
            f"{item['specify_question']}\n{MCQ_TAIL}")
    if item["cls"] == "C":
        gold = next(i for i, o in enumerate(opts)
                    if o["kind"] == "jur" and item["minority_school"] in o["schools"])
    else:
        gold = next(i for i, o in enumerate(opts)
                    if item["minority_school"] in o["schools"])
    return item["specify_text"], user, opts, gold


# ----------------------------------------------------------------- metrics --
def entropy(probs: list[float]) -> float:
    import math
    return -sum(p * math.log(max(p, 1e-12)) for p in probs)