code/synth_common.py

"""
Shared pieces for the synthetic Arabic math-reasoning corpus.

Holds everything both the generator and the dataset builder need, so the accept/reject decision
is made by exactly one implementation:

  * AXES        — the variation grid every prompt is drawn from
  * build_task  — task_id -> (axes draw, prompt). Deterministic: task 7 always has the same draw,
                  which is what makes the run resumable and reproducible.
  * parse_items — split a raw completion into {instruction, reasoning, answer} records
  * audit       — verify the arithmetic *inside* the reasoning and reject rows that don't check out

Why the audit matters: a 12B model writes fluent Arabic derivations that are wrong roughly one
time in ten. The reasoning states its own equations ("35 × 4 = 140"), so they can be re-evaluated
exactly. This is the same discipline as the numeral audit in build_dataset.py, applied to
generated arithmetic rather than translated arithmetic.
"""
import random
import re
import unicodedata

# ---------------------------------------------------------------- variation grid

DOMAINS = [
    "التسوق وشراء البقالة", "الرواتب والأجور اليومية", "المسافات والسفر بالسيارة",
    "وصفات الطبخ والمقادير", "المدرسة والطلاب والدرجات", "المزرعة والمحاصيل والحيوانات",
    "الرياضة والمباريات والنقاط", "الادخار والميزانية الشهرية", "البناء والدهان والبلاط",
    "فواتير الهاتف والإنترنت", "الوقود واستهلاك السيارة", "السوق وبيع الخضار والفواكه",
    "المكتبة وشراء الكتب", "ورشة الخياطة والأقمشة", "توصيل الطلبات والمندوبين",
    "المستشفى والجرعات والمواعيد", "الحديقة والأشجار والري", "المخبز وإنتاج الخبز",
    "محل الحلويات والطلبيات", "النقل المدرسي والحافلات", "تربية الدواجن والبيض",
    "مصنع صغير وخط الإنتاج", "الرحلات المدرسية والتذاكر", "تأجير المعدات باليوم",
    "صيد السمك وبيعه", "محطة تحلية ومياه الشرب", "متجر إلكتروني والشحن",
    "معرض السيارات والأقساط", "الزكاة وتوزيع الصدقات", "تنظيم حفل وعدد الضيوف",
]

OPERATIONS = [
    ("الجمع والطرح", "خطوتان من الجمع والطرح على أعداد صحيحة"),
    ("الضرب والقسمة", "الضرب ثم القسمة للوصول إلى نصيب الواحد"),
    ("النسبة المئوية", "حساب نسبة مئوية من مبلغ ثم الزيادة أو النقصان"),
    ("الكسور", "أخذ كسر من كمية ثم التعامل مع الباقي"),
    ("النسبة والتناسب", "توزيع كمية بنسبة معلومة بين طرفين أو ثلاثة"),
    ("سعر الوحدة", "استنتاج سعر الوحدة ثم حساب تكلفة كمية مختلفة"),
    ("المتوسط الحسابي", "حساب متوسط مجموعة قيم ثم استنتاج قيمة ناقصة"),
    ("السرعة والزمن والمسافة", "علاقة السرعة بالزمن والمسافة على مرحلتين"),
    ("المساحة والمحيط", "مساحة أو محيط شكل مستطيل ثم تكلفة تغطيته"),
    ("الربح والخسارة", "شراء بسعر وبيع بسعر آخر وحساب الربح أو نسبته"),
    ("الخصم والضريبة", "خصم على السعر ثم إضافة ضريبة القيمة المضافة"),
    ("العمل والزمن", "عاملان أو ثلاثة ينجزون عملاً في أزمنة مختلفة"),
    ("تحويل الوحدات", "تحويل بين الوحدات (كجم/جرام، ساعة/دقيقة، متر/سم) ضمن الحل"),
    ("المتتاليات البسيطة", "نمط يزيد بمقدار ثابت كل يوم أو كل أسبوع"),
    ("المعادلة ذات المجهول الواحد", "صياغة معادلة بسيطة بمجهول واحد وحلها"),
]

# --- relational comparison pool -------------------------------------------------------------
# A SEPARATE pool, not appended to OPERATIONS, so build_task(N) stays byte-identical for the task
# ids that produced the original 100,323-row corpus — that corpus must remain reproducible.
#
# Why this exists: v5 scores pass@8 = 0 on prompts like "عمر أخيه ضعف عمره" because the grid never
# generated the relation. Measured coverage in the 100k corpus: ضعف/أضعاف 1.34%, أكثر بـ/أقل بـ
# 0.26%. RL cannot repair that class (no correct sample in the group -> no gradient), so it has to
# come from data.
#
# The defining trait: one quantity is stated ONLY in terms of another, so the solver must compute
# the second quantity before it can do anything else. That first step is exactly what v5 skips.
RELATIONAL_OPS = [
    ("المقارنة بالمضاعفة", "كمية ثانية تساوي ضعف الأولى، ثم يُطلب المجموع أو الفرق"),
    ("المقارنة بالتضعيف المتعدد", "كمية ثانية تساوي ثلاثة أو أربعة أضعاف الأولى، ثم المجموع"),
    ("المقارنة بالتنصيف", "كمية ثانية تساوي نصف الأولى، ثم المجموع أو الباقي"),
    ("المقارنة بالكسر", "كمية ثانية تساوي ثلث أو ربع الأولى، ثم المجموع"),
    ("المقارنة بالزيادة المطلقة", "كمية ثانية تزيد عن الأولى بمقدار ثابت، ثم المجموع"),
    ("المقارنة بالنقصان المطلق", "كمية ثانية تقل عن الأولى بمقدار ثابت، ثم المجموع"),
    ("المقارنة بين ثلاث كميات", "ثلاث كميات كل واحدة معرّفة بدلالة سابقتها، ثم المجموع"),
    ("المقارنة العكسية", "تُعطى الكمية الثانية وعلاقتها بالأولى، والمطلوب الأولى ثم المجموع"),
]

# Relational problems are naturally short — a 5-step chain forces padding with filler steps.
RELATIONAL_STEPS = [(2, "خطوتين"), (2, "خطوتين"), (3, "ثلاث خطوات")]

STEPS = [(2, "خطوتين"), (3, "ثلاث خطوات"), (3, "ثلاث خطوات"), (4, "أربع خطوات")]

NUMBER_STYLE = [
    "أعداد صحيحة صغيرة (أقل من 100)",
    "أعداد صحيحة متوسطة (بين 100 و 5000)",
    "مبالغ مالية تحتوي على كسور عشرية بمنزلتين",
    "أعداد كبيرة نسبيًا (بالآلاف) تقبل القسمة بدون باقٍ",
]

REGIONS = [
    ("السعودية", "ريال"), ("مصر", "جنيه"), ("الإمارات", "درهم"), ("الأردن", "دينار"),
    ("المغرب", "درهم"), ("الكويت", "دينار"), ("تونس", "دينار"), ("العراق", "دينار"),
    ("سوريا", "ليرة"), ("عُمان", "ريال"), ("قطر", "ريال"), ("الجزائر", "دينار"),
]

NAMES = [
    "أحمد", "فاطمة", "محمد", "سارة", "خالد", "نور", "يوسف", "مريم", "عمر", "ليلى",
    "سلمان", "هند", "طارق", "رنا", "بلال", "أسماء", "زياد", "دعاء", "كريم", "شيماء",
    "ياسر", "بثينة", "حسن", "رغد", "إبراهيم", "جنى", "مصطفى", "لمى", "سعيد", "أروى",
]

TWISTS = [
    "أضف رقمًا واحدًا في نص المسألة لا يُستخدم في الحل (معلومة زائدة)",
    "اجعل المسألة تقارن بين شخصين أو بين يومين",
    "اجعل السؤال يطلب الباقي أو المتبقي وليس المجموع",
    "اجعل السؤال يطلب عدد المرات أو عدد المجموعات",
    "اجعل المسألة على مرحلتين زمنيتين (اليوم الأول ثم اليوم الثاني)",
    "اجعل الإجابة النهائية تحتاج تقريبًا لأقرب عدد صحيح مع توضيح سبب التقريب",
    "لا تضف أي تعقيد إضافي، اجعلها مباشرة وواضحة",
    "اجعل السؤال يطلب النسبة المئوية للنتيجة من الإجمالي",
]

ITEMS_PER_TASK = 4

# ---------------------------------------------------------------- prompt

Q, T, A, E = "### مسألة", "### تفكير", "### الإجابة", "### نهاية"

PROMPT = """أنت معلم رياضيات عربي تكتب مسائل تدريبية عالية الجودة باللغة العربية الفصحى.

اكتب {k} مسائل حسابية **مختلفة تمامًا عن بعضها** بالمواصفات التالية:

- المجال: {domain}
- نوع العملية: {op_name} — {op_hint}
- عدد خطوات الحل: {steps_word} تقريبًا (اكتب فقط الخطوات التي يحتاجها الحل فعلًا)
- طبيعة الأرقام: {numbers}
- السياق: {region}، والعملة {currency}
- استخدم أسماء مثل: {names}
- {twist}

اكتب كل مسألة بهذا الشكل بالضبط، ولا تكتب أي شيء آخر خارج هذه الوسوم:

{Q}
نص المسألة في جملة أو جملتين، وينتهي بسؤال واضح.
{T}
خطوات الحل مرقمة، وكل خطوة تحتوي على معادلة صريحة بالأرقام مثل: 45 × 3 = 135
{A}
إذن، جملة واحدة تذكر الإجابة النهائية بالرقم مع وحدتها.
{E}

قواعد إلزامية:
1. كل معادلة تكتبها يجب أن تكون **صحيحة حسابيًا**. تحقق من كل عملية قبل كتابتها.
2. الرقم في سطر الإجابة يجب أن يساوي ناتج آخر خطوة في التفكير.
3. اكتب الأرقام بالأرقام الإنجليزية (0-9) وليس بالحروف ولا بالأرقام الهندية.
4. اكتب بالعربية فقط، بدون أي كلمة إنجليزية.
5. لا تكرر مسألة سبق أن كتبتها في هذه الإجابة.
6. لا تكتب خطوات فارغة أو بلا فائدة مثل «63 + 0 = 63» أو «4500 = 4500» أو «24 × 1 = 24».
   كل خطوة يجب أن تُنتج قيمة جديدة لم تكن معروفة قبلها.
7. إذا لم تناسب العملية المطلوبة المجالَ المطلوب، غيّر تفاصيل المسألة لتناسبها من البداية،
   ولا تعلّق على ذلك ولا تعِد صياغة المسألة داخل خطوات الحل.
"""


def build_task(task_id: int, seed: int = 1234, pool: str = "default"):
    """task_id -> (axes, prompt). Deterministic, so a resumed run redraws identical prompts.

    pool="default"    the original 15-operation grid — draws are byte-identical to the run that
                      produced the 100,323-row corpus, so that corpus stays reproducible.
    pool="relational" the relational-comparison pool (RELATIONAL_OPS). Use a disjoint task-id
                      range for it, the same rule that keeps two generating nodes from colliding.
    """
    ops, steps_pool = (RELATIONAL_OPS, RELATIONAL_STEPS) if pool == "relational" else (OPERATIONS, STEPS)
    rng = random.Random(seed * 1_000_003 + task_id)
    domain = rng.choice(DOMAINS)
    op_name, op_hint = rng.choice(ops)
    _, steps_word = rng.choice(steps_pool)
    numbers = rng.choice(NUMBER_STYLE)
    region, currency = rng.choice(REGIONS)
    names = "، ".join(rng.sample(NAMES, 3))
    twist = rng.choice(TWISTS)
    axes = {"domain": domain, "op": op_name, "steps": steps_word, "numbers": numbers,
            "region": region, "twist": twist}
    prompt = PROMPT.format(k=ITEMS_PER_TASK, domain=domain, op_name=op_name, op_hint=op_hint,
                           steps_word=steps_word, numbers=numbers, region=region,
                           currency=currency, names=names, twist=twist, Q=Q, T=T, A=A, E=E)
    return axes, prompt


# ---------------------------------------------------------------- parsing

AR_DIGITS = str.maketrans("٠١٢٣٤٥٦٧٨٩٪", "0123456789%")
BLOCK_RE = re.compile(
    re.escape(Q) + r"(?P<q>.*?)" + re.escape(T) + r"(?P<t>.*?)" + re.escape(A) + r"(?P<a>.*?)"
    + r"(?:" + re.escape(E) + r"|$)", re.S)
LATIN_RE = re.compile(r"[A-Za-z]")
CONCLUSION = ("إذن", "لذلك", "بالتالي", "وبالتالي", "لذا", "في النهاية", "الخلاصة")
# The model sometimes breaks frame and rewrites the problem inside its own derivation
# ("*تصحيح للمسألة لتناسب العملية المطلوبة*") when a drawn operation type does not fit the drawn
# domain. The arithmetic in those rows often still checks out, so the audit alone won't catch them.
META_MARKERS = ("تصحيح للمسألة", "سنعيد صياغة", "لتناسب العملية", "المسألة تطلب",
                "إعادة صياغة", "لنفترض أن المسألة", "تصحيح السؤال", "بما أن المسألة")


def norm(text: str) -> str:
    text = unicodedata.normalize("NFC", text.replace("‏", "").replace("‎", ""))
    text = text.translate(AR_DIGITS)
    text = re.sub(r"[ \t]+", " ", text)
    text = re.sub(r"\n{3,}", "\n\n", text)
    return text.strip()


def parse_items(raw: str):
    """Raw completion -> list of {instruction, reasoning, answer}. Malformed blocks are skipped."""
    out = []
    for m in BLOCK_RE.finditer(raw):
        q, t, a = (norm(m.group(g)) for g in ("q", "t", "a"))
        if q and t and a:
            out.append({"instruction": q, "reasoning": t, "answer": a})
    return out


# ---------------------------------------------------------------- arithmetic audit

NUM = r"\d+(?:\.\d+)?"
# "x + 0 = x", "x × 1 = x", "x = x" — vacuous steps the model emits to hit a requested step count.
NOOP_RE = re.compile(r"(?<![\d.])(" + NUM + r")\s*(?:[+\-−–]\s*0+(?:\.0+)?|"
                     r"[×xX*]\s*1(?:\.0+)?|[÷/]\s*1(?:\.0+)?)?\s*=\s*\1(?![\d.])")

# One arithmetic segment: digits and operators only, at least one digit. Deliberately excludes
# newlines — a chain never spans two steps.
_SC = r"[\d.()+\-−–×xX*÷/ \t]"
SEG = _SC + r"*\d" + _SC + r"*"
# "1." / "2)" list numbering at the head of a step is not part of the arithmetic; left in, it
# turns "1. 75 + 15" into the un-evaluatable "1. 75 + 15".
LIST_NUM_RE = re.compile(r"^[ \t]*\d+[.)][ \t]*", re.M)
# A *chain*: "a × b = c + d = e". Models routinely show their working this way, and reading only
# the first "=" turns a correct chain into a false rejection — every segment must be compared.
CHAIN_RE = re.compile(r"(?<![\d.])(" + SEG + r"(?:=" + SEG + r")+)")


def _evaluate(expr: str):
    """Evaluate a pure-arithmetic expression with normal precedence. None if it isn't one."""
    expr = expr.replace("−", "-").replace("–", "-").replace("×", "*").replace("x", "*") \
               .replace("X", "*").replace("÷", "/")
    expr = expr.strip()
    if not expr or not re.fullmatch(r"[\d.+\-*/() ]+", expr) or not re.search(r"\d", expr):
        return None
    try:
        # The regex above admits only digits, operators, dots, parens and spaces — no names,
        # no calls, no attribute access — so this cannot execute anything else.
        value = eval(expr, {"__builtins__": {}}, {})            # noqa: S307
    except Exception:
        return None
    return value if isinstance(value, (int, float)) else None


def _decimals(text: str) -> int:
    text = text.strip()
    return len(text.split(".")[1]) if "." in text else 0


def _close(a, b, shown: str | None = None):
    """
    Is the stated value `b` an acceptable rendering of the true value `a`?

    Exact, or `b` is `a` rounded to the precision `b` is written with — "3200 / 60 = 53.33" and
    "9.44 × 60 = 566.4 → 566" are how people write arithmetic, not arithmetic errors. Only
    rounding at the displayed precision is forgiven, so a genuinely wrong number still fails.
    """
    if abs(a - b) <= max(1e-6, abs(b) * 1e-9):
        return True
    if shown is not None:
        d = _decimals(shown)
        if abs(round(a, d) - b) <= 1e-9:
            return True
        if d == 0 and b in (int(a), int(a) + 1 if a > 0 else int(a) - 1):
            return True            # a step that floors or ceils, e.g. "3 crates with a remainder"
    return False


def audit(reasoning: str, answer: str, min_equations: int = 1):
    """
    -> (ok, reason, n_checked)

    Every "a op b = c [= d ...]" chain the reasoning states is re-evaluated end to end. One
    segment that disagrees rejects the row. The answer's number must also match the last chain's
    final value — a derivation that is internally correct but ends on a different number is the
    failure mode that survives any fluency check.
    """
    checked = 0
    last = None
    body = LIST_NUM_RE.sub("", reasoning)
    for m in CHAIN_RE.finditer(body):
        parts = m.group(1).split("=")
        values = [_evaluate(part) for part in parts]
        if any(v is None for v in values) or len(values) < 2:
            continue                       # not pure arithmetic (units, words, %) — not our call
        checked += 1
        head = values[0]
        for part, v in zip(parts[1:], values[1:]):
            # Compare each stated value against the *first* segment, which is the one written in
            # full precision; `part` carries how many decimals the model chose to show.
            if not _close(head, v, part):
                return False, f"bad_equation:{m.group(1).strip()[:60]}", checked
        last = values[-1]
    if checked < min_equations:
        return False, "no_equations", checked

    ans_shown = re.findall(NUM, answer)
    if not ans_shown:
        return False, "answer_has_no_number", checked
    if last is not None and not any(_close(last, float(x), x) for x in ans_shown):
        return False, "answer_not_last_result", checked
    return True, "ok", checked


def validate(item, min_reasoning=40, max_reasoning=1200):
    """Full accept/reject for one parsed item -> (ok, reason)."""
    q, t, a = item["instruction"], item["reasoning"], item["answer"]
    if LATIN_RE.search(q) or LATIN_RE.search(t) or LATIN_RE.search(a):
        return False, "latin_residue"
    if not (20 <= len(q) <= 600):
        return False, "question_length"
    if not (min_reasoning <= len(t) <= max_reasoning):
        return False, "reasoning_length"
    if not (10 <= len(a) <= 300):
        return False, "answer_length"
    if not a.lstrip("*-# ").startswith(CONCLUSION):
        return False, "no_conclusion_marker"
    if any(mark in t for mark in (Q, A, E)):
        return False, "tag_leak"
    if any(mark in t or mark in q for mark in META_MARKERS):
        return False, "meta_commentary"
    if NOOP_RE.search(t):
        return False, "noop_step"
    ok, reason, _ = audit(t, a)
    return ok, reason


def dedup_key(instruction: str) -> str:
    """Numbers masked out, so the same template with different values collapses to one key."""
    return re.sub(r"\d+(?:\.\d+)?", "#", re.sub(r"\W+", "", instruction))