SHK4K's picture
Add @spaces.GPU inference function for ZeroGPU
2c71b9c verified
Raw History Blame Contribute Delete
12 kB
"""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)