#!/usr/bin/env python3
"""
normaere — Middle High German text normalizer.
Original design: vertical workflow.
1. Input textarea at top (wide, line-numbered)
2. Control bar with options + Normalize button
3. Side-by-side numbered comparison table as output
"""
import argparse
import base64
import html as html_mod
import json
import os
import sys
from pathlib import Path
from threading import Lock
from typing import Optional
# ZeroGPU: `spaces` MUST be imported before torch (which is loaded lazily via
# src.inference) so that ZeroGPU's CUDA patches are active before any CUDA call.
# Do not set SPACES=false or other env hacks — they disable ZeroGPU integration.
import spaces
import gradio as gr
# ---------------------------------------------------------------------------
# Lazy-loaded normalizer
# ---------------------------------------------------------------------------
_normalizer = None
_lock = Lock()
def _get_normalizer(model_path: str, config_path: Optional[str] = None):
global _normalizer
if _normalizer is not None:
return _normalizer
with _lock:
if _normalizer is not None:
return _normalizer
from src.inference import MHGNormalizer
_normalizer = MHGNormalizer(model_path, config_path, defer_gpu=True)
return _normalizer
# ---------------------------------------------------------------------------
# Normalization
# ---------------------------------------------------------------------------
@spaces.GPU(duration=120)
def run_normalize(
input_text: str,
preserve_punctuation: bool,
preserve_capitalization: bool,
attach_en_proclitic: bool,
lenition_t_after_n: bool,
lenition_t_after_l: bool,
niet_to_niht: bool,
common_apocopes: bool,
model_path: str,
config_path: str,
):
"""Normalize text. Returns (raw_input, raw_output) or (input, error_msg)."""
if not input_text or not input_text.strip():
return "", ""
# Enforce the word limit BEFORE touching the GPU so oversized inputs never
# consume ZeroGPU quota. Truncate to the first _MAX_WORDS words while
# preserving line breaks (the UI also auto-caps; this is the safety net).
wc = len(input_text.split())
if wc > _MAX_WORDS:
kept_lines = []
running = 0
for line in input_text.split("\n"):
line_words = line.split()
if running + len(line_words) > _MAX_WORDS:
# Partial last line: take only the words that fit
remaining = _MAX_WORDS - running
if remaining > 0:
kept_lines.append(" ".join(line_words[:remaining]))
break
kept_lines.append(line)
running += len(line_words)
input_text = "\n".join(kept_lines)
config = Path(config_path) if config_path else None
norm = _get_normalizer(model_path, str(config) if config else None)
norm._ensure_on_device()
try:
result = norm.normalize_multiline_text(
input_text,
preserve_punctuation=preserve_punctuation,
preserve_capitalization=preserve_capitalization,
attach_en_proclitic=attach_en_proclitic,
lenition_t_after_n=lenition_t_after_n,
lenition_t_after_l=lenition_t_after_l,
niet_to_niht=niet_to_niht,
common_apocopes=common_apocopes,
)
return input_text, result
except Exception as exc:
return input_text, f"⚠ Error: {exc}"
# ---------------------------------------------------------------------------
# Input size limit
# ---------------------------------------------------------------------------
# ZeroGPU grants a limited daily GPU quota. To prevent a single large request
# from consuming it (or timing out the @spaces.GPU call), we cap the input at
# _MAX_WORDS. This is enforced server-side in run_normalize and surfaced in
# the UI via a live word counter.
_MAX_WORDS = 4000
def load_file(file_obj):
if file_obj is None:
return ""
try:
return file_obj.decode("utf-8")
except UnicodeDecodeError:
return file_obj.decode("latin-1")
# ---------------------------------------------------------------------------
# Comparison table builder
# ---------------------------------------------------------------------------
def build_table(input_text: str, output_text: str) -> str:
"""Build a numbered 3-column HTML table."""
if not input_text and not output_text:
return ""
in_lines = input_text.split("\n")
out_lines = output_text.split("\n")
n = max(len(in_lines), len(out_lines))
rows = []
for i in range(n):
il = in_lines[i] if i < len(in_lines) else ""
ol = out_lines[i] if i < len(out_lines) else ""
if not il.strip() or not ol.strip():
continue
rows.append(
f""
f"| {i + 1} | "
f"{html_mod.escape(il)} | "
f"{html_mod.escape(ol)} | "
f"
"
)
if not rows:
return 'Nothing to show. Enter some text and press Normalize.
'
return (
""
"| # | Input | Normalized |
|---|
"
f"{''.join(rows)}
Remove Punctuation
Convert to Lowercase
Attach en-Proclitic
n-Lenition? süntære → sündære
l-Lenition? solte → solde
niet, niut → niht
Common Apocopes? vile → vil
vore → vor
wile → wil
wole → wol
Highlight Changes
''')
# Main content (input card → replaced by output card after normalize)
with gr.Column(scale=4):
with gr.Column(elem_classes=["input-wrap"]) as input_card:
with gr.Row():
with gr.Column(elem_classes=["input-col"]):
gr.HTML('0 words · 0 lines · 0 chars (limit: 4,000 words)
')
input_text = gr.Textbox(
lines=12,
max_lines=30,
show_label=False,
container=False,
elem_id="input-text",
placeholder="Paste or type Middle High German text here…",
)
with gr.Column(elem_classes=["output-wrap"], visible=False) as output_card:
comparison_output = gr.HTML(value="")
# ---- State ----
mp_state = gr.State(model_path)
cp_state = gr.State(config_path or "")
# ---- Normalize (no post-processing — handled client-side) ----
def _do(text, *args):
mp, cp = args[-2], args[-1]
# Only normalize — post-processing happens client-side.
# run_normalize is decorated with @spaces.GPU, so inference runs
# on an allocated ZeroGPU (CUDA) device. A server-side word limit
# (run_normalize) prevents oversized inputs from exhausting quota.
raw_in, raw_out = run_normalize(text, True, True, False, False, False, False, False, mp, cp)
if not raw_in and not raw_out:
return "", gr.update(visible=False), gr.update(visible=True)
if raw_out.startswith("⚠"):
warning_html = (
f'{html_mod.escape(raw_out)}
'
)
return warning_html, gr.update(visible=True), gr.update(visible=False)
out_lines = raw_out.split("\n")
b64 = base64.b64encode(json.dumps(out_lines).encode()).decode()
data_div = f''
table_html = build_table(raw_in, raw_out) + data_div
return table_html, gr.update(visible=True), gr.update(visible=False)
btn_normalize.click(
fn=_do,
inputs=[input_text, mp_state, cp_state],
outputs=[comparison_output, output_card, input_card],
)
# ---- Clear ----
def _clear():
return "", "", gr.update(visible=True), gr.update(visible=False)
btn_clear_all.click(fn=_clear, outputs=[comparison_output, input_text, input_card, output_card])
# ---- Revert to input (triggered when user clears output) ----
with gr.Column(elem_classes=["revert-col"]):
btn_revert = gr.Button(elem_id="revertBtn")
def _revert():
return "", "", gr.update(visible=False), gr.update(visible=True)
btn_revert.click(
fn=_revert,
outputs=[comparison_output, input_text, output_card, input_card],
)
# ---- Show input without clearing (for file upload) ----
with gr.Column(elem_classes=["revert-col"]):
btn_show_input = gr.Button(elem_id="showInputBtn")
def _show_input_only():
return gr.update(visible=False), gr.update(visible=True)
btn_show_input.click(
fn=_show_input_only,
outputs=[output_card, input_card],
)
# ---- Copy / Download (JS) ----
btn_copy.click(js="window._copyOut()")
btn_dl.click(js="window._dlOut()")
# ---- Footer ----
logo_path = Path(__file__).parent / "static" / "logo.png"
logo_src = ""
if logo_path.exists():
import base64 as _b64
logo_data = logo_path.read_bytes()
logo_src = f"data:image/png;base64,{_b64.b64encode(logo_data).decode()}"
gr.HTML(
f''
)
return demo
# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------
def _ensure_model_downloaded(model_path: str) -> str:
"""If model_path is a Hub repo ID, download it to a local cache.
On HF Spaces the model is not bundled, so we must download it at runtime.
For local paths, returns the path unchanged.
"""
p = Path(model_path)
if p.exists():
return str(p)
# Treat as Hub repo ID and download
from huggingface_hub import snapshot_download
print(f"Model not found locally — downloading '{model_path}' from Hub …")
cached = snapshot_download(repo_id=model_path, cache_dir=None)
print(f" Downloaded to: {cached}")
return cached
def main():
# Respect HF Spaces $PORT env var (defaults to 7860)
default_port = int(os.environ.get("PORT", 7860))
default_host = os.environ.get("HOST", "0.0.0.0")
parser = argparse.ArgumentParser(description="normaere")
parser.add_argument("--model_path", default=None)
parser.add_argument("--config_path", default="config.yaml")
parser.add_argument("--port", type=int, default=default_port)
parser.add_argument("--host", default=default_host)
parser.add_argument("--share", action="store_true")
args = parser.parse_args()
import yaml
config_path = Path(args.config_path) if args.config_path else None
model_path = args.model_path
if model_path is None and config_path and config_path.exists():
try:
with open(config_path, "r") as f:
cfg = yaml.safe_load(f)
model_path = cfg.get("inference", {}).get("model_path")
except Exception:
pass
if model_path is None:
model_path = "JonasHermann/normaere-model"
model_path = _ensure_model_downloaded(model_path)
cp = str(config_path) if config_path else None
# Preload the model at startup (loads weights to CPU RAM as BF16).
# The slow disk I/O + deserialization happens here, OUTSIDE the
# ZeroGPU duration budget. The CPU→GPU transfer is deferred to
# run_normalize (inside @spaces.GPU) via _ensure_on_device().
print("Preloading model to CPU …")
_get_normalizer(model_path, cp)
print("Model preloaded.\n")
print(f"Starting normaere on http://{args.host}:{args.port}")
print(f"Model: {model_path}\n")
demo = build_app(model_path, cp)
demo.launch(
server_name=args.host,
server_port=args.port,
share=args.share,
css=_CSS,
head=_HEAD_JS,
)
if __name__ == "__main__":
main()