Blame for app.py

"""Hadith splitter: sanad / matn segmentation for SHK4K/hadith-segmentation-bert."""

import html
import os
import re
import sys

import gradio as gr
import spaces
import torch
from transformers import AutoModelForTokenClassification, AutoTokenizer

MODEL_ID = os.environ.get("MODEL_ID", "SHK4K/hadith-segmentation-bert")
# The Hub config still ships LABEL_0..3, so model.config.id2label is useless here.
SANAD_IDS, MATN_IDS = (0, 1), (2, 3)
MAX_LEN, STRIDE, BATCH = 512, 128, 8
# ponytail: hard cap on very long pastes. 64 windows ~= 25k subword tokens, far
# beyond any single hadith; drop the cap if the demo ever faces books.
MAX_WINDOWS = 64
# The model was trained on undiacritized text; feeding it harakat produces
# garbage, so they are stripped for inference and mapped back for the output.
_HARAKAT = re.compile("[\u064b-\u0652\u0670\u0640]")

EXAMPLES = [
    ['حَدَّثَنَا عَمْرُو بْنُ خَالِدٍ، قَالَ حَدَّثَنَا اللَّيْثُ، عَنْ يَزِيدَ، عَنْ أَبِي الْخَيْرِ، عَنْ عَبْدِ اللَّهِ بْنِ عَمْرٍو ـ رضى الله عنهما ـ أَنَّ رَجُلاً، سَأَلَ النَّبِيَّ صلى الله عليه وسلم أَىُّ الإِسْلاَمِ خَيْرٌ قَالَ \u200f "\u200f تُطْعِمُ الطَّعَامَ، وَتَقْرَأُ السَّلاَمَ عَلَى مَنْ عَرَفْتَ وَمَنْ لَمْ تَعْرِفْ \u200f"\u200f\u200f.\u200f'],
    ['حَدَّثَنَا أَبُو الْيَمَانِ، قَالَ أَخْبَرَنَا شُعَيْبٌ، قَالَ حَدَّثَنَا أَبُو الزِّنَادِ، عَنِ الأَعْرَجِ، عَنْ أَبِي هُرَيْرَةَ ـ رضى الله عنه ـ أَنَّ رَسُولَ اللَّهِ صلى الله عليه وسلم قَالَ \u200f "\u200f فَوَالَّذِي نَفْسِي بِيَدِهِ لاَ يُؤْمِنُ أَحَدُكُمْ حَتَّى أَكُونَ أَحَبَّ إِلَيْهِ مِنْ وَالِدِهِ وَوَلَدِهِ \u200f"\u200f\u200f.\u200f'],
    ['حَدَّثَنَا أَبُو الْوَلِيدِ، قَالَ حَدَّثَنَا شُعْبَةُ، قَالَ أَخْبَرَنِي عَبْدُ اللَّهِ بْنُ عَبْدِ اللَّهِ بْنِ جَبْرٍ، قَالَ سَمِعْتُ أَنَسًا، عَنِ النَّبِيِّ صلى الله عليه وسلم قَالَ \u200f "\u200f آيَةُ الإِيمَانِ حُبُّ الأَنْصَارِ، وَآيَةُ النِّفَاقِ بُغْضُ الأَنْصَارِ \u200f"\u200f\u200f.\u200f'],
    ['حَدَّثَنَا الْحُمَيْدِيُّ عَبْدُ اللَّهِ بْنُ الزُّبَيْرِ ، قَالَ : حَدَّثَنَا سُفْيَانُ ، قَالَ : حَدَّثَنَا يَحْيَى بْنُ سَعِيدٍ الْأَنْصَارِيُّ ، قَالَ : أَخْبَرَنِي مُحَمَّدُ بْنُ إِبْرَاهِيمَ التَّيْمِيُّ ، أَنَّهُ سَمِعَ عَلْقَمَةَ بْنَ وَقَّاصٍ اللَّيْثِيَّ ، يَقُولُ : سَمِعْتُ عُمَرَ بْنَ الْخَطَّابِ رَضِيَ اللَّهُ عَنْهُ عَلَى الْمِنْبَرِ، قَالَ : سَمِعْتُ رَسُولَ اللَّهِ صَلَّى اللَّهُ عَلَيْهِ وَسَلَّمَ، يَقُولُ : " إِنَّمَا الْأَعْمَالُ بِالنِّيَّاتِ، وَإِنَّمَا لِكُلِّ امْرِئٍ مَا نَوَى، فَمَنْ كَانَتْ هِجْرَتُهُ إِلَى دُنْيَا يُصِيبُهَا أَوْ إِلَى امْرَأَةٍ يَنْكِحُهَا، فَهِجْرَتُهُ إِلَى مَا هَاجَرَ إِلَيْهِ'],
]

_model = _tokenizer = None


def load():
    global _model, _tokenizer
    if _model is None:
        _tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
        _model = AutoModelForTokenClassification.from_pretrained(MODEL_ID).eval()
    # Called from inside an @GPU scope, where the CUDA device is attached.
    if torch.cuda.is_available() and _model.device.type == "cpu":
        _model.to("cuda")
    return _tokenizer, _model


def segment(text):
    """Return {'sanad', 'matn', 'html', 'note'} for raw hadith text."""
    if not text or not text.strip():
        raise ValueError("empty input")

    # Strip harakat for inference, remembering where each character came from,
    # so the cut can be mapped back onto the original diacritized text.
    norm_chars, origin = [], []
    for i, ch in enumerate(text):
        if not _HARAKAT.match(ch):
            norm_chars.append(ch)
            origin.append(i)
    norm = "".join(norm_chars)

    tok, model = load()
    enc = tok(
        norm,
        truncation=True,
        max_length=MAX_LEN,
        stride=STRIDE,
        return_overflowing_tokens=True,
        return_offsets_mapping=True,
        return_tensors="pt",
        padding=True,
    )
    offsets = enc.pop("offset_mapping")
    enc.pop("overflow_to_sample_mapping", None)
    windows = min(enc["input_ids"].shape[0], MAX_WINDOWS)

    preds = []
    with torch.inference_mode():
        for i in range(0, windows, BATCH):
            batch = {k: v[i : i + BATCH].to(model.device) for k, v in enc.items()}
            preds.append(model(**batch).logits.argmax(-1).cpu())
    preds = torch.cat(preds)

    # Majority vote per word. A word straddling two windows is keyed by its
    # character start, which both windows report identically.
    votes = {}
    for w in range(windows):
        for pos, word_id in enumerate(enc.word_ids(w)):
            if word_id is None:
                continue
            start, end = (int(x) for x in offsets[w][pos])
            if start == end:
                continue
            label = int(preds[w][pos])
            votes.setdefault(start, {}).setdefault(label, 0)
            votes[start][label] += 1

    # The labels describe a single cut, so take the first matn word as the
    # boundary: sanad + matn always cover the whole input, in order.
    cut = next(
        (s for s in sorted(votes) if max(votes[s], key=votes[s].get) in MATN_IDS),
        len(norm),
    )
    cut = origin[cut] if cut < len(origin) else len(text)

    note = f"{len(text.split())} words · {windows} window(s)"
    if windows < enc["input_ids"].shape[0]:
        note += f" · only the first {MAX_WINDOWS} windows were processed"
    if not text[cut:].strip():
        note += " · model found no matn boundary"

    return {
        "sanad": text[:cut].strip(),
        "matn": text[cut:].strip(),
        "html": _highlight(text, cut),
        "note": note,
    }


def _highlight(text, cut):
    escape = html.escape  # before wrapping: the text comes from a public textbox
    return (
        '<div class="hadith" dir="rtl">'
        f'<span class="sanad">{escape(text[:cut])}</span>'
        f'<span class="matn">{escape(text[cut:])}</span>'
        "</div>"
    )


@spaces.GPU(duration=60)
def run(text):
    try:
        r = segment(text)
    except ValueError:
        raise gr.Error("أدخل نص الحديث أولاً / Please paste a hadith first.")
    return r["html"], r["sanad"], r["matn"], r["note"]


load()  # warm the Hub cache at startup, outside any GPU allocation


CSS = """
#hadith-input textarea { direction: rtl; font-size: 1.05rem; line-height: 1.9; }
.hadith { font-size: 1.15rem; line-height: 2.1; padding: 1rem; border-radius: .5rem;
          background: #fffdf7; color: #1a1a1a; min-height: 3rem; }
.hadith .sanad { background: #fde68a; }
.hadith .matn { background: #bbf7d0; }
@media (prefers-color-scheme: dark) {
  .hadith { background: #1f2937; color: #f3f4f6; }
  .hadith .sanad { background: #92400e; }
  .hadith .matn { background: #166534; }
}
"""

with gr.Blocks(title="Hadith Splitter") as demo:
    gr.Markdown(
        "# فاصل الحديث — Hadith Splitter\n"
        "يقسم الحديث الخام إلى **سند** و**متن**. · Splits a raw hadith into its "
        "**sanad** (chain of narrators) and **matn** (the saying itself)."
    )
    inp = gr.Textbox(
        label="الحديث / Hadith",
        lines=8,
        rtl=True,
        elem_id="hadith-input",
        placeholder="حدثنا ... قال رسول الله صلى الله عليه وسلم ...",
    )
    with gr.Row():
        btn = gr.Button("قسّم / Segment", variant="primary")
        clear = gr.Button("مسح / Clear")

    out_html = gr.HTML()
    note = gr.Markdown()
    with gr.Row():
        sanad = gr.Textbox(
            label="السند / Sanad", lines=6, buttons=["copy"], rtl=True
        )
        matn = gr.Textbox(
            label="المتن / Matn", lines=6, buttons=["copy"], rtl=True
        )

    gr.Examples(examples=EXAMPLES, inputs=inp, label="أمثلة / Examples")

    gr.Markdown(
        "> **ملاحظات / Notes:** يُجرَّد التشكيل تلقائيًا قبل التحليل، ويُعاد النص الأصلي "
        "بتشكيله في المخرجات · diacritics are stripped for inference, your original "
        "text is returned unchanged. "
        "النص الأطول من 512 رمزًا يُعالج بنوافذ متداخلة · longer text is processed in "
        "overlapping windows, never truncated. "
        "النموذج: [SHK4K/hadith-segmentation-bert]"
        "(https://huggingface.co/SHK4K/hadith-segmentation-bert)."
    )

    outputs = [out_html, sanad, matn, note]
    btn.click(run, inp, outputs)
    inp.submit(run, inp, outputs)
    clear.click(lambda: (None, "", "", "", ""), None, [inp] + outputs)


def _selftest():
    sample = (
        "حدثنا الحميدي عبد الله بن الزبير ، قال : حدثنا سفيان ، قال : حدثنا يحيى بن سعيد "
        "الأنصاري ، قال : أخبرني محمد بن إبراهيم التيمي ، أنه سمع علقمة بن وقاص الليثي ، "
        'يقول : سمعت عمر بن الخطاب رضي الله عنه على المنبر، قال : سمعت رسول الله صلى الله '
        'عليه وسلم، يقول : " إنما الأعمال بالنيات، وإنما لكل امرئ ما نوى، فمن كانت هجرته '
        "إلى دنيا يصيبها أو إلى امرأة ينكحها، فهجرته إلى ما هاجر إليه"
    )
    r = segment(sample)
    print("SANAD:", r["sanad"][:100])
    print("MATN:", r["matn"][:100])
    assert r["sanad"].startswith("حدثنا الحميدي"), r["sanad"][:80]
    assert "إنما الأعمال بالنيات" in r["matn"]
    assert sample.strip().startswith(r["sanad"]) and sample.strip().endswith(r["matn"])

    long_text = " ".join([sample] * 6)
    lr = segment(long_text)
    assert lr["sanad"] and lr["matn"], "long input lost its boundary"
    assert sample.strip()[-40:] in lr["matn"], "tail of a long input was dropped"

    xss = '<script>alert(1)</script> ' + sample
    assert "<script>" not in segment(xss)["html"], "input was not escaped"

    # Diacritized input must segment too, and keep its diacritics in the output.
    vocalized = (
        "حَدَّثَنَا أَبُو الْوَلِيدِ، قَالَ حَدَّثَنَا شُعْبَةُ، عَنْ أَنَسٍ، قَالَ: "
        "قَالَ رَسُولُ اللَّهِ صَلَّى اللَّهُ عَلَيْهِ وَسَلَّمَ: «الدِّينُ النَّصِيحَةُ»"
    )
    vr = segment(vocalized)
    print("VOCALIZED SANAD:", vr["sanad"][:80])
    print("VOCALIZED MATN:", vr["matn"][:80])
    assert vr["sanad"].startswith("حَدَّثَنَا"), "diacritics broke the sanad"
    assert vr["matn"], "diacritized input lost its matn"
    print("selftest ok")


if __name__ == "__main__":
    if "--selftest" in sys.argv:
        _selftest()
    else:
        demo.launch(css=CSS)