rules_common.py

"""
Shared pieces for the Arabic rule-checking corpus.

Same discipline as synth_common, applied to a task whose ground truth is a set of boolean
verdicts instead of a number:

  * RULES       — a library of rules that are *programmatically decidable*. Each carries an Arabic
                  statement for the prompt and a checker over the raw text.
  * build_task  — task_id -> (axes draw, prompt). Deterministic, so the run stays resumable.
  * parse_items — split a raw completion into {text, analysis, verdicts, rule_ids}
  * validate    — recompute every verdict with the checkers and reject rows the model got wrong

Why the checkers matter: if the model both writes the text and judges it, the labels are only as
good as the model, and its errors correlate with the text it produced. Here the text is generated
and the labels are *computed*, so a row is either consistent with ground truth or it is dropped.
That is the same move audit() makes for arithmetic.

Balance comes from the prompt asking for a specific satisfy/violate pattern per item. The model
often misses the target - that is fine and expected. The checkers record what the text actually
does, so a missed target costs balance, never correctness.
"""
import os
import random
import re

ITEMS_PER_TASK = 3

Q, T, A, E = "### نص", "### تحليل", "### الحكم", "### نهاية"

# ---------------------------------------------------------------- checkers

AR_DIGITS = str.maketrans("٠١٢٣٤٥٦٧٨٩", "0123456789")
CURRENCY = "ريال|ريالا|ريالًا|جنيه|جنيها|جنيهًا|درهم|درهما|درهمًا|دينار|دينارا|دينارًا|ليرة|دولار"
# A short list is a correctness bug, not a shortcut: any city missing from it silently
# mislabels a row, and the label is what the whole corpus is for. The smoke test caught
# الدمام missing. Kept deliberately broad across every region in REGIONS.
CITIES = [
    "القاهرة", "الإسكندرية", "الجيزة", "المنصورة", "طنطا", "الزقازيق", "أسيوط", "المنيا",
    "بورسعيد", "السويس", "الإسماعيلية", "دمياط", "الأقصر", "أسوان", "الفيوم", "بني سويف",
    "الرياض", "جدة", "مكة", "المدينة", "الدمام", "الخبر", "الظهران", "الطائف", "تبوك",
    "أبها", "خميس مشيط", "بريدة", "عنيزة", "حائل", "نجران", "جازان", "الأحساء", "الهفوف",
    "ينبع", "الجبيل", "القطيف",
    "دبي", "أبوظبي", "الشارقة", "عجمان", "رأس الخيمة", "الفجيرة", "أم القيوين", "العين",
    "الدوحة", "الريان", "الوكرة", "الخور",
    "الكويت", "حولي", "الفروانية", "الجهراء", "الأحمدي",
    "المنامة", "المحرق", "الرفاع", "مدينة عيسى",
    "مسقط", "صلالة", "صحار", "نزوى", "صور",
    "بغداد", "البصرة", "الموصل", "أربيل", "النجف", "كربلاء", "السليمانية", "كركوك",
    "الناصرية", "الرمادي", "دهوك", "بعقوبة",
    "عمّان", "عمان", "إربد", "اربد", "الزرقاء", "العقبة", "السلط", "مادبا", "الكرك",
    "تونس", "صفاقس", "سوسة", "القيروان", "بنزرت", "قابس", "نابل", "المنستير",
    "الدار البيضاء", "الرباط", "مراكش", "فاس", "طنجة", "أكادير", "مكناس", "وجدة",
    "تطوان", "القنيطرة", "سلا",
    "الخرطوم", "بيروت", "دمشق", "حلب", "طرابلس", "بنغازي", "صنعاء", "عدن", "نواكشوط",
]
DAYS = ["السبت", "الأحد", "الاثنين", "الإثنين", "الثلاثاء", "الأربعاء", "الخميس", "الجمعة"]
MONTHS = ["يناير", "فبراير", "مارس", "أبريل", "مايو", "يونيو", "يوليو", "أغسطس",
          "سبتمبر", "أكتوبر", "نوفمبر", "ديسمبر"]


def _n(text):
    return text.translate(AR_DIGITS)


def _words(text):
    return [w for w in re.split(r"\s+", text.strip()) if w]


def _has_phone(text):
    """A run of 8-15 digits that is not a price.

    8 is the floor, not 9: Kuwait, Qatar, Bahrain, Oman and Tunisia all use 8-digit numbers, and
    the smoke test surfaced 98765432 and 55443322 as phone numbers the old 9-digit threshold
    silently missed. A run immediately followed by a currency word is a price, not a phone.
    """
    t = _n(text)
    for m in re.finditer(r"(?:\+|00)?\d[\d\s\-]{6,}\d", t):
        digits = re.sub(r"\D", "", m.group(0))
        if not (8 <= len(digits) <= 15):
            continue
        if re.match(r"\s*(?:" + CURRENCY + r")", t[m.end():m.end() + 14]):
            continue                       # "12000000 دينار" is a price
        return True
    return False


CHECKS = {
    "has_price":     lambda t: bool(re.search(r"\d+\s*(?:" + CURRENCY + r")", _n(t))),
    "has_phone":     _has_phone,
    "no_phone":      lambda t: not _has_phone(t),
    "no_latin":      lambda t: not re.search(r"[A-Za-z]", t),
    "has_city":      lambda t: any(c in t for c in CITIES),
    "no_url":        lambda t: not re.search(r"https?://|www\.|\.com|\.net|\.org", t, re.I),
    "no_email":      lambda t: not re.search(r"[^\s@]+@[^\s@]+\.[^\s@]+", t),
    "has_number":    lambda t: bool(re.search(r"\d", _n(t))),
    "no_excess_punct": lambda t: not re.search(r"[!؟]\s*[!؟]", t),
    "ends_question": lambda t: t.strip().rstrip("‏‎").endswith(("؟", "?")),
    "has_date":      lambda t: any(d in t for d in DAYS + MONTHS)
                               or bool(re.search(r"\d{1,2}\s*/\s*\d{1,2}", _n(t))),
}

# Arabic statement shown in the prompt, keyed the same way.
STATEMENTS = {
    "has_price":       "يجب أن يذكر النص سعرًا مقترنًا بعملة",
    "has_phone":       "يجب أن يذكر النص رقم هاتف للتواصل (من 8 إلى 15 خانة)",
    "no_phone":        "يجب ألا يحتوي النص على أي رقم هاتف (من 8 إلى 15 خانة)",
    "no_latin":        "يجب ألا يحتوي النص على أي حروف لاتينية",
    "has_city":        "يجب أن يذكر النص اسم مدينة",
    "no_url":          "يجب ألا يحتوي النص على رابط أو موقع إلكتروني",
    "no_email":        "يجب ألا يحتوي النص على بريد إلكتروني",
    "has_number":      "يجب أن يحتوي النص على رقم واحد على الأقل",
    "no_excess_punct": "يجب ألا يحتوي النص على علامات تعجب أو استفهام متتالية",
    "ends_question":   "يجب أن ينتهي النص بعلامة استفهام",
    "has_date":        "يجب أن يذكر النص تاريخًا أو يومًا أو شهرًا",
}

# Parameterised length rules are built per task, so their id carries the bound.
def _length_rule(kind, n):
    rid = f"{kind}_{n}"
    if kind == "min_words":
        return rid, f"يجب ألا يقل النص عن {n} كلمة", (lambda t, n=n: len(_words(t)) >= n)
    return rid, f"يجب ألا يزيد النص عن {n} كلمة", (lambda t, n=n: len(_words(t)) <= n)


def rule_check(rid, text):
    """-> bool. Works for both library rules and the parameterised length ones."""
    if rid in CHECKS:
        return CHECKS[rid](text)
    m = re.fullmatch(r"(min_words|max_words)_(\d+)", rid or "")
    if m:
        n = int(m.group(2))
        return len(_words(text)) >= n if m.group(1) == "min_words" else len(_words(text)) <= n
    return None                                    # unknown id - caller rejects the row


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

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

REGIONS = ["مصر", "السعودية", "الإمارات", "قطر", "الكويت", "العراق", "الأردن", "تونس", "المغرب"]

# "Is this a city?" has no exact answer - district, town and neighbourhood shade into each other,
# and the smoke test produced السالمية and الصليبخات (real Kuwaiti districts) alongside الأردن
# (a country the model called a city). A rule library that claims decidable labels cannot hold a
# predicate that needs world knowledge, so the rule now *names* acceptable cities per region and
# the checker's gazetteer is a superset of every list shown.
REGION_CITIES = {
    "مصر":       ["القاهرة", "الإسكندرية", "الجيزة", "المنصورة", "أسيوط"],
    "السعودية":  ["الرياض", "جدة", "الدمام", "الخبر", "الطائف"],
    "الإمارات":  ["دبي", "أبوظبي", "الشارقة", "عجمان", "العين"],
    "قطر":       ["الدوحة", "الريان", "الوكرة", "الخور"],
    "الكويت":    ["الكويت", "حولي", "الفروانية", "الجهراء", "الأحمدي"],
    "العراق":    ["بغداد", "البصرة", "الموصل", "أربيل", "النجف"],
    "الأردن":    ["عمّان", "إربد", "الزرقاء", "العقبة", "السلط"],
    "تونس":      ["تونس", "صفاقس", "سوسة", "القيروان", "بنزرت"],
    "المغرب":    ["الدار البيضاء", "الرباط", "مراكش", "فاس", "طنجة"],
}

# Rules that co-determine each other give a free label: a model can score them without reading
# the text. no_phone is the exact negation of has_phone, and both has_phone and has_price force
# has_number to pass because a phone number and a price both contain digits. Never draw a
# conflicting pair into the same task.
CONFLICTS = [
    ("has_phone", "no_phone"),
    ("has_phone", "has_number"),
    ("has_price", "has_number"),
]


def _conflicts(rid, chosen_ids):
    return any((rid == a and b in chosen_ids) or (rid == b and a in chosen_ids)
               for a, b in CONFLICTS)


PROMPT = """أنت تكتب بيانات تدريبية لنموذج يتحقق من مطابقة النصوص لقواعد محددة.

اكتب {k} أمثلة مستقلة. كل مثال عبارة عن نص عربي واقعي من نوع: {doc_type} ({doc_hint}) في {region}.

القواعد المطلوب فحصها:
{rules_block}

هدف الكتابة (للتنويع فقط، وليس هو الحكم):
{pattern_block}

⚠️ الأهم في هذه المهمة:
بعد أن تكتب النص، افحصه من جديد كأنك تراه لأول مرة واحكم على ما يحتويه **فعلًا**، وليس على ما
كنت تنوي كتابته. إن خالف النص هدف الكتابة أعلاه، فاكتب الحقيقة كما هي في قسم الحكم. الحكم وصف
لما في النص، وليس تكرارًا للتعليمات.

اكتب كل مثال بهذا الشكل بالضبط:

{Q}
(النص العربي هنا، من سطر إلى ثلاثة أسطر)
{T}
(سطر واحد لكل قاعدة: اقتبس الدليل الحرفي من النص بين علامتي تنصيص، ثم اذكر النتيجة)
{A}
{verdict_template}
{E}

قواعد مهمة:
- في سطر الحكم اكتب المعرّف الإنجليزي للقاعدة كما هو، ثم كلمة واحدة: مطابق أو مخالف.
- "مطابق" تعني أن النص يحقق ما تطلبه القاعدة. "مخالف" تعني أنه لا يحققه. لا تخلط بينها وبين
  هدف الكتابة.
- النص نفسه يجب أن يكون طبيعيًا وواقعيًا، وليس مصنوعًا ليبدو كاختبار.
- نوّع الأسماء والأرقام والتفاصيل بين الأمثلة.
- لا تكتب أي شيء خارج الوسوم."""


def build_task(task_id: int, seed: int = 1234):
    """task_id -> (axes, prompt). Deterministic, so a resumed run redraws identical prompts."""
    rng = random.Random(seed * 1_000_003 + task_id)
    doc_type, doc_hint = rng.choice(DOC_TYPES)
    region = rng.choice(REGIONS)

    n_rules = rng.choice([3, 3, 4])
    pool = [(rid, STATEMENTS[rid], None) for rid in CHECKS]
    if rng.random() < 0.55:                        # roughly half the tasks carry a length bound
        kind = rng.choice(["min_words", "max_words"])
        n = rng.choice([15, 20, 25, 30]) if kind == "min_words" else rng.choice([25, 30, 40, 50])
        pool.append(_length_rule(kind, n))
    # Draw one at a time so a conflicting rule can be skipped rather than poisoning the task.
    chosen = []
    for cand in rng.sample(pool, len(pool)):
        if len(chosen) >= n_rules:
            break
        if not _conflicts(cand[0], {c[0] for c in chosen}):
            chosen.append(cand)
    n_rules = len(chosen)
    examples = "، ".join(REGION_CITIES.get(region, CITIES[:5]))
    chosen = [(rid, (stmt + f" من هذه المدن: {examples}") if rid == "has_city" else stmt, chk)
              for rid, stmt, chk in chosen]

    # Force at least one satisfied and one violated so the labels do not collapse to all-pass.
    pattern = [True] * n_rules
    n_violate = rng.choice([1, 1, 2])
    for i in rng.sample(range(n_rules), min(n_violate, n_rules - 1)):
        pattern[i] = False

    rules_block = "\n".join(f"- {rid}: {stmt}" for rid, stmt, _ in chosen)
    # Deliberately different vocabulary from the verdict words (مطابق/مخالف): when the writing
    # goal and the verdict share wording, the model echoes the instruction instead of inspecting
    # the text it wrote. The smoke test showed exactly that - "مخالف because the text DOES
    # contain a phone number" on a rule requiring one.
    want = [rid for (rid, _, _), ok in zip(chosen, pattern) if ok]
    avoid = [rid for (rid, _, _), ok in zip(chosen, pattern) if not ok]
    parts = []
    if want:
        parts.append("- اجعل النص يستوفي: " + "، ".join(want))
    if avoid:
        parts.append("- واجعله لا يستوفي: " + "، ".join(avoid))
    pattern_block = "\n".join(parts)
    # Listed in a different order than the writing goal, so position carries no hint.
    order = sorted((rid for rid, _, _ in chosen))
    verdict_template = "\n".join(f"{rid}: ..." for rid in order)

    axes = {"doc_type": doc_type, "region": region,
            "rule_ids": [rid for rid, _, _ in chosen],
            "target_pattern": pattern}
    prompt = PROMPT.format(k=ITEMS_PER_TASK, doc_type=doc_type, doc_hint=doc_hint, region=region,
                           rules_block=rules_block, pattern_block=pattern_block,
                           verdict_template=verdict_template, Q=Q, T=T, A=A, E=E)
    return axes, prompt


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

BLOCK_RE = re.compile(
    re.escape(Q) + r"(?P<text>.*?)" + re.escape(T) + r"(?P<analysis>.*?)"
    + re.escape(A) + r"(?P<verdicts>.*?)" + re.escape(E), re.S)
VERDICT_RE = re.compile(r"^\s*([A-Za-z_][A-Za-z0-9_]*)\s*[:：]\s*(مطابق|مخالف)\s*$", re.M)


def parse_items(raw: str):
    """-> [{text, analysis, verdicts: {rule_id: bool}, rule_ids: [...]}]"""
    items = []
    for m in BLOCK_RE.finditer(raw or ""):
        verdicts, order = {}, []
        for v in VERDICT_RE.finditer(m.group("verdicts")):
            rid = v.group(1)
            if rid not in verdicts:
                order.append(rid)
            verdicts[rid] = (v.group(2) == "مطابق")
        items.append({"text": m.group("text").strip(),
                      "analysis": m.group("analysis").strip(),
                      "verdicts": verdicts, "rule_ids": order})
    return items


# ---------------------------------------------------------------- validation

def validate(item, min_text=20, max_text=900, min_analysis=20, max_analysis=1200):
    """
    Full accept/reject for one parsed item -> (ok, reason).

    The decisive check is the last one: every verdict the model stated is recomputed from the
    text by the rule's own checker. A row whose analysis reads fluently but whose verdicts do not
    match what the text actually does is exactly the failure mode a fluency check cannot catch.
    """
    text, analysis, verdicts = item["text"], item["analysis"], item["verdicts"]
    if not (min_text <= len(text) <= max_text):
        return False, "text_length"
    if not (min_analysis <= len(analysis) <= max_analysis):
        return False, "analysis_length"
    if not verdicts:
        return False, "no_verdicts"
    if len(verdicts) < 2:
        return False, "too_few_rules"
    if any(tag in text for tag in (Q, T, A, E)):
        return False, "tag_leak"
    for rid, stated in verdicts.items():
        truth = rule_check(rid, text)
        if truth is None:
            return False, f"unknown_rule:{rid}"
        if truth != stated:
            return False, f"verdict_mismatch:{rid}"
    return True, "ok"


def truth_verdicts(item):
    """Ground truth recomputed from the text - what a reward function should compare against."""
    return {rid: rule_check(rid, item["text"]) for rid in item["rule_ids"]}


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