"""Gradio demo for CTFAlign — document-level word alignment.

Pick a model, paste a source/target document pair, and the app shows the
predicted word alignments with hover-to-highlight (the same viewer as the repo's
viewer.html, embedded via an iframe).
"""
import gc
import html
import json
import os
import shutil
import tempfile

import gradio as gr
from transformers import AutoConfig

from ctfalign import align_from_similarity
from ctfalign.defaults import resolve_k, resolve_layer
from ctfalign.encoders import HFEncoder
from ctfalign.similarity import get_sim_matrix

# Display name -> HuggingFace id. Models from the paper's experiments that have a
# document-level tuned layer (plus LaBSE).
MODELS = {
    "mmBERT-base": "jhu-clsp/mmBERT-base",
    "Qwen3-Embedding-4B": "Qwen/Qwen3-Embedding-4B",
    "LaBSE": "sentence-transformers/LaBSE",
}
LANG_PAIRS = ["(auto)", "en-fr", "en-ro", "en-ja", "en-zh", "la-gr", "en-cz"]
DEFAULT_MODEL = "mmBERT-base"

# Dropdown label -> internal method name (order matters; CTFAlign first).
METHOD_CHOICES = {
    "ctfalign": "ctfalign",
    "mdpalign": "mdpalign-fuzzy",
}

EXAMPLE_A = 'Animals are multicellular, eukaryotic organisms belonging to the biological kingdom Animalia (/ˌænɪˈmeɪliə/[4]). With few exceptions, animals consume organic material, breathe oxygen, have myocytes and are able to move, can reproduce sexually, and grow from a hollow sphere of cells, the blastula, during embryonic development. Animals form a clade, meaning that they arose from a single common ancestor. Over 1.5 million living animal species have been described, of which around 1.05 million are insects, over 85,000 are molluscs, and around 65,000 are vertebrates. It has been estimated there are as many as 7.77 million animal species on Earth. Animal body lengths range from 8.5 μm (0.00033 in) to 33.6 m (110 ft). They have complex ecologies and interactions with each other and their environments, forming intricate food webs. The scientific study of animals is known as zoology, and the study of animal behaviour is known as ethology.\nThe animal kingdom is divided into five major clades, namely Porifera, Ctenophora, Placozoa, Cnidaria and Bilateria. Most living animal species belong to the clade Bilateria, a highly proliferative clade whose members have a bilaterally symmetric and significantly cephalised body plan, and the vast majority of bilaterians belong to two large clades: the protostomes, which include organisms such as arthropods, molluscs, flatworms, annelids and nematodes; and the deuterostomes, which include echinoderms, hemichordates and chordates, the last of which contains the vertebrates. The much smaller basal phylum Xenacoelomorpha have an uncertain position within Bilateria.'
EXAMPLE_B = "Les animaux sont des organismes eucaryotes pluricellulaires appartenant au règne biologique *Animalia*. À quelques exceptions près, les animaux consomment de la matière organique, respirent de l'oxygène, possèdent des myocytes et sont capables de se déplacer ainsi que de se reproduire par voie sexuée ; ils se développent à partir d'une sphère creuse de cellules, la blastula, au cours de l'embryogenèse. Les animaux forment un clade, ce qui signifie qu'ils sont issus d'un ancêtre commun unique. Plus de 1,5 million d'espèces animales vivantes ont été décrites ; parmi elles, environ 1,05 million sont des insectes, plus de 85 000 sont des mollusques et environ 65 000 sont des vertébrés. On estime qu'il pourrait exister jusqu'à 7,77 millions d'espèces animales sur Terre. La taille des animaux varie de 8,5 μm à 33,6 m. Ils entretiennent des relations écologiques complexes entre eux et avec leur environnement, formant des réseaux trophiques élaborés. L'étude scientifique des animaux est appelée zoologie, tandis que l'étude de leur comportement est nommée éthologie.\nLe règne animal se divise en cinq clades majeurs : les Porifera (éponges), les Ctenophora (cténaires), les Placozoa (placozoaires), les Cnidaria (cnidaires) et les Bilateria (bilatériens). La plupart des espèces animales vivantes appartiennent au clade des Bilateria, un groupe très diversifié dont les membres présentent un plan d'organisation à symétrie bilatérale et une céphalisation marquée. La grande majorité des bilatériens se répartit en deux grands clades : les protostomiens, qui incluent des organismes tels que les arthropodes, les mollusques, les plathelminthes (vers plats), les annélides et les nématodes ; et les deutérostomiens, qui comprennent les échinodermes, les hémichordés et les chordés (groupe incluant les vertébrés). Le phylum basal Xenacoelomorpha, beaucoup plus restreint, occupe une position incertaine au sein des Bilateria."

# Small models kept resident in RAM (never evicted) so their first Align is
# instant. The large default model (Qwen) instead uses a single rotating slot to
# bound peak memory: at most one *unpinned* model is held at a time.
PINNED_MODELS = ["mmBERT-base", "LaBSE"]

# hf_id -> HFEncoder. Keyed by model only (not layer): the layer just selects
# which hidden-states to read and is set per request, so changing the layer /
# language pair never reloads the model.
_enc_cache = {}


def _resolve_layer(hf_id, lang_pair):
    short = hf_id.split("/")[-1]
    layer = resolve_layer(short, lang_pair, "documents")
    if layer is None:                       # model not in the tuned tables
        cfg = AutoConfig.from_pretrained(hf_id)
        n = getattr(cfg, "num_hidden_layers", None) or 12
        layer = round((2 / 3) * n)
    return int(layer)


def _get_encoder(hf_id, layer):
    layer = int(layer)
    enc = _enc_cache.get(hf_id)
    if enc is None:
        pinned_ids = {MODELS[m] for m in PINNED_MODELS}
        if hf_id not in pinned_ids:
            # Loading a large unpinned model: free the previous unpinned one
            # first (pinned small models stay resident).
            for k in [k for k in _enc_cache if k not in pinned_ids]:
                del _enc_cache[k]
            gc.collect()
        enc = HFEncoder(hf_id, layer)
        _enc_cache[hf_id] = enc
    enc.layer = layer                       # cheap: just changes which layer is read
    return enc


def _default_layer(model_display, lp_choice):
    """The tuned layer that would actually be used for this model + lang pair."""
    hf_id = MODELS[model_display]
    lp = None if lp_choice == "(auto)" else lp_choice
    return _resolve_layer(hf_id, lp)


def _default_k(method_display):
    """The default constraint hyperparameter (w for CTFAlign, k for MDPAlign)."""
    return resolve_k(METHOD_CHOICES.get(method_display, method_display), None)


def _is_cjk(ch):
    """True for Chinese/Japanese/Korean characters (no whitespace word boundaries)."""
    o = ord(ch)
    return (0x4E00 <= o <= 0x9FFF      # CJK Unified Ideographs
            or 0x3400 <= o <= 0x4DBF   # CJK Extension A
            or 0x3040 <= o <= 0x30FF   # Hiragana + Katakana
            or 0xAC00 <= o <= 0xD7A3   # Hangul syllables
            or 0xF900 <= o <= 0xFAFF)  # CJK Compatibility Ideographs


def _needs_subword(text):
    """Heuristic: text is not whitespace-segmented (mostly CJK), so word
    projection would be meaningless — align at the subword level instead."""
    chars = [c for c in text if not c.isspace()]
    if not chars:
        return False
    return sum(_is_cjk(c) for c in chars) / len(chars) > 0.2


def _units_and_ids(text, encoded, seg):
    """Per-side display units and the b2w map fed to ``align_from_similarity``.

    ``seg`` is 'auto' | 'words' | 'subwords'. In subword mode each non-special
    token is its own unit and the b2w map is the identity, so the token alignment
    is returned as-is (no extrapolation to whitespace words).

    Returns ``(units, b2w_ids, level)`` where ``level`` is 'word' or 'subword' —
    the resolved segmentation, which the exported payload has to record because
    the alignment indices are meaningless without it.
    """
    if seg == "subwords" or (seg == "auto" and _needs_subword(text)):
        # EncodedText.tokens comes out of the same pass as the embeddings, so it
        # is always row-aligned with them (this is the package's units="tokens").
        units = list(encoded.tokens)
        return units, list(range(len(units))), "subword"
    return text.split(), encoded.word_ids, "word"


def _viewer_html(units_a, units_b, pairs, sim_pairs):
    data = {
        "a": units_a,
        "b": units_b,
        "pairs": [[int(i), int(j)] for i, j in pairs],
        "sim_pairs": [[int(i), int(j)] for i, j in sim_pairs],
    }
    doc = _VIEWER_TEMPLATE.replace("__DATA__", json.dumps(data))
    srcdoc = html.escape(doc, quote=True)
    return (f'<iframe srcdoc="{srcdoc}" '
            f'style="width:100%;height:560px;border:1px solid #e9ebf0;'
            f'border-radius:14px;background:#fff;"></iframe>')


def _pharaoh(pairs):
    """Standard ``i-j`` alignment string (the GIZA/AER interchange format)."""
    return " ".join(f"{int(i)}-{int(j)}" for i, j in pairs)


def align_payload(model_display, method_display, mode, lp_choice, layer_in, k_in,
                  text_a, text_b, seg="auto"):
    """Run one alignment and return the canonical result dict.

    Single source of truth for both the UI (everything shown is rendered from
    this) and the ``/align`` API endpoint. Failures raise ``gr.Error`` so an API
    caller gets a failed request instead of a successful response whose body is
    an error message.
    """
    if not (text_a or "").strip() or not (text_b or "").strip():
        raise gr.Error("Please enter both a source and a target text.")
    if model_display not in MODELS:
        raise gr.Error(f"Unknown model {model_display!r}; choose one of {list(MODELS)}.")
    if method_display not in METHOD_CHOICES and method_display not in METHOD_CHOICES.values():
        raise gr.Error(f"Unknown method {method_display!r}; choose one of {list(METHOD_CHOICES)}.")
    hf_id = MODELS[model_display]
    method = METHOD_CHOICES.get(method_display, method_display)
    lang_pair = None if lp_choice == "(auto)" else lp_choice
    try:
        layer = int(layer_in) if layer_in not in (None, "") else _resolve_layer(hf_id, lang_pair)
        enc = _get_encoder(hf_id, layer)
        k = float(k_in) if k_in not in (None, "") else None
        chunk_size = max(enc.max_length - 2, 1)
        n_tok_a = len(enc.tokenizer.encode(text_a, add_special_tokens=False, verbose=False))
        n_tok_b = len(enc.tokenizer.encode(text_b, add_special_tokens=False, verbose=False))
        # Encode once and share the similarity matrix between the selected method
        # and the unconstrained SimAlign baseline (the latter is just an argmax).
        a = enc.encode(text_a)
        b = enc.encode(text_b)
        # Pick units per side: whitespace words, or subwords for non-segmented
        # (zh/ja) text — in which case the alignment is left at the token level.
        units_a, ids_a, level_a = _units_and_ids(text_a, a, seg)
        units_b, ids_b, level_b = _units_and_ids(text_b, b, seg)
        sim = get_sim_matrix(a.embeddings, b.embeddings)
        pairs = align_from_similarity(sim, ids_a, ids_b,
                                      method=method, mode=mode, k=k)
        sim_pairs = align_from_similarity(sim, ids_a, ids_b,
                                          method="simalign", mode=mode)
    except gr.Error:
        raise
    except Exception as e:                  # surface model/config errors to the caller
        raise gr.Error(f"Error: {e}") from e

    warnings = []
    if n_tok_a > chunk_size or n_tok_b > chunk_size:
        model_name = model_display.split(" — ")[0]
        warnings.append(
            f"The text exceeds {model_name}'s context window of {enc.max_length} tokens "
            f"(source {n_tok_a}, target {n_tok_b} tokens). It was split into chunks, which "
            f"degrades alignment quality — switch to a model with a larger context window, "
            f"e.g. mmBERT-base."
        )

    pairs = [[int(i), int(j)] for i, j in pairs]
    sim_pairs = [[int(i), int(j)] for i, j in sim_pairs]

    def _unit(units, i):
        return units[i] if 0 <= i < len(units) else ""

    return {
        "config": {
            "model": hf_id,
            "method": method,
            "mode": mode,
            "layer": int(layer),
            "k": resolve_k(method, k),
            "segmentation": seg,
            "lang_pair": lang_pair,
        },
        "source": {"level": level_a, "units": units_a},
        "target": {"level": level_b, "units": units_b},
        # Indices are what scorers need; the surface strings make the same rows
        # readable (and checkable) without re-running the tokeniser.
        "alignment": [{"s": i, "t": j,
                       "source": _unit(units_a, i), "target": _unit(units_b, j)}
                      for i, j in pairs],
        "pharaoh": _pharaoh(pairs),
        "baseline_simalign": {"pairs": sim_pairs, "pharaoh": _pharaoh(sim_pairs)},
        "warnings": warnings,
    }


_EXPORT_ROOT = os.path.join(tempfile.gettempdir(), "ctfalign-exports")


def _download_files(payload):
    """Write the payload to disk and return (json_path, talp_path) for download.

    One directory per call (so a concurrent request can't overwrite a file that
    is still being served), pruned to the most recent few so a long-lived Space
    doesn't accumulate them.
    """
    os.makedirs(_EXPORT_ROOT, exist_ok=True)
    stale = sorted(os.scandir(_EXPORT_ROOT), key=lambda e: e.stat().st_mtime)[:-20]
    for e in stale:
        shutil.rmtree(e.path, ignore_errors=True)
    d = tempfile.mkdtemp(prefix="run-", dir=_EXPORT_ROOT)
    json_path = os.path.join(d, "alignment.json")
    talp_path = os.path.join(d, "alignment.talp")
    with open(json_path, "w", encoding="utf-8") as f:
        json.dump(payload, f, ensure_ascii=False, indent=2)
    with open(talp_path, "w", encoding="utf-8") as f:
        f.write(payload["pharaoh"] + "\n")
    return json_path, talp_path


def run(model_display, method, mode, lp_choice, layer_in, k_in, text_a, text_b, seg="auto"):
    # The align button is re-enabled here (it was greyed out + relabelled while processing).
    align_ready = gr.update(interactive=True, value="Align")
    empty = ("", [], "", None, None)        # cleared export widgets on the error paths

    if not (text_a or "").strip() or not (text_b or "").strip():
        return ("<p style='font-family:Inter,sans-serif;color:#b91c1c'>"
                "Please enter both a source and a target text.</p>", align_ready) + empty
    try:
        payload = align_payload(model_display, method, mode, lp_choice,
                                layer_in, k_in, text_a, text_b, seg)
    except gr.Error as e:
        # gr.Error's str() is the repr'd message, so take .message when present.
        msg = getattr(e, "message", None) or str(e)
        return (f"<pre style='color:#b91c1c;white-space:pre-wrap'>"
                f"{html.escape(msg)}</pre>", align_ready) + empty

    units_a = payload["source"]["units"]
    units_b = payload["target"]["units"]
    pairs = [[p["s"], p["t"]] for p in payload["alignment"]]
    sim_pairs = payload["baseline_simalign"]["pairs"]

    warning = ""
    for w in payload["warnings"]:
        warning += (
            f"<p style='font-family:Inter,sans-serif;font-size:13px;color:#92400e;"
            f"background:#fef3c7;border:1px solid #fde68a;border-radius:10px;padding:8px 12px;margin:0 0 10px'>"
            f"⚠️ {html.escape(w)}</p>"
        )
    header = (f"<p style='font-family:Inter,sans-serif;font-size:12.5px;color:#9aa1b1;"
              f"margin:0 0 8px'>layer {payload['config']['layer']} · "
              f"<b style='color:#6366f1'>{len(pairs)}</b> {payload['config']['method']} links · "
              f"<b style='color:#6b7280'>{len(sim_pairs)}</b> SimAlign links · "
              f"hover a word to compare</p>")

    table = [[p["s"], p["source"], p["t"], p["target"]] for p in payload["alignment"]]
    json_path, talp_path = _download_files(payload)
    return (warning + header + _viewer_html(units_a, units_b, pairs, sim_pairs),
            align_ready,
            payload["pharaoh"],
            table,
            json.dumps(payload, ensure_ascii=False, indent=2),
            gr.update(value=json_path),
            gr.update(value=talp_path))


THEME = gr.themes.Soft(
    primary_hue="indigo",
    secondary_hue="indigo",
    neutral_hue="slate",
    font=[gr.themes.GoogleFont("Inter"), "system-ui", "sans-serif"],
).set(
    body_background_fill="#f6f7f9",
    block_radius="14px",
    button_large_radius="999px",
    button_small_radius="999px",
)

CUSTOM_CSS = """
.gradio-container{max-width:1120px !important;margin:0 auto !important;}
#title h1{font-weight:700;letter-spacing:-.02em;}
#title{margin-bottom:2px;}
footer{display:none !important;}
"""

with gr.Blocks(title="CTFAlign Demo", theme=THEME, css=CUSTOM_CSS) as demo:
    gr.Markdown(
        "# CTFAlign · Document-level Word Alignment\n"
        "Choose a model, paste a document pair, and see the "
        "predicted word alignments. Hover a word to highlight its correspondence(s) "
        "in the other text.",
        elem_id="title",
    )
    with gr.Row():
        model = gr.Dropdown(list(MODELS), value=DEFAULT_MODEL, label="Model")
        method = gr.Dropdown(list(METHOD_CHOICES), value="ctfalign", label="Method")
        mode = gr.Dropdown(["argmax", "itermax"], value="argmax", label="Mode")
        seg = gr.Dropdown(["auto", "words", "subwords"], value="auto",
                          label="Segmentation (subwords for zh/ja)")
    with gr.Row():
        lp = gr.Dropdown(LANG_PAIRS, value="(auto)", label="Language pair (for tuned layer)")
        layer = gr.Number(value=_default_layer(DEFAULT_MODEL, "(auto)"), precision=0,
                          label="Layer (auto-set from model + lang pair; editable)")
        k = gr.Number(value=_default_k("ctfalign"),
                      label="k / w (auto-set from method; editable)")
    with gr.Row():
        ta = gr.Textbox(lines=8, label="Text A", value=EXAMPLE_A)
        tb = gr.Textbox(lines=8, label="Text B", value=EXAMPLE_B)
    btn = gr.Button("Align", variant="primary")
    out = gr.HTML()

    # Everything below is rendered from the same payload the /align endpoint
    # returns, so what a human copies and what an agent fetches cannot drift.
    with gr.Accordion("Export alignment", open=False):
        pharaoh_box = gr.Textbox(
            label="Pharaoh format (i-j, 0-indexed over the units below)",
            lines=3, show_copy_button=True, interactive=False,
        )
        pair_table = gr.Dataframe(
            headers=["source #", "source", "target #", "target"],
            datatype=["number", "str", "number", "str"],
            label="Aligned pairs", wrap=True,
        )
        with gr.Row():
            dl_json = gr.DownloadButton("Download alignment.json", size="sm")
            dl_talp = gr.DownloadButton("Download alignment.talp", size="sm")
        json_box = gr.Code(label="Full result (JSON)", language="json", interactive=False)

    exports = [pharaoh_box, pair_table, json_box, dl_json, dl_talp]
    btn.click(
        lambda: gr.update(interactive=False, value="Aligning…"), None, btn,
        api_name=False,
    ).then(
        run, [model, method, mode, lp, layer, k, ta, tb, seg], [out, btn] + exports,
        api_name=False,
    )

    # Machine-readable endpoint: `Client(...).predict(..., api_name="/align")`
    # returns the payload dict directly (and raises on failure).
    api_btn = gr.Button(visible=False)
    api_out = gr.JSON(visible=False)
    api_btn.click(
        align_payload, [model, method, mode, lp, layer, k, ta, tb, seg], api_out,
        api_name="align",
    )

    # Keep the hyperparameter fields in sync with the choices.
    model.change(_default_layer, [model, lp], layer, api_name=False)
    lp.change(_default_layer, [model, lp], layer, api_name=False)
    method.change(_default_k, method, k, api_name=False)
    gr.Markdown(
        "_The use of Qwen3-Embedding-4B can take some time due to model loading and the CPU backend. "
        "Segmentation auto aligns whitespace words, but switches to subwords for "
        "non-segmented text (e.g. zh/ja) so the alignment is shown token-by-token instead._"
    )


# Single-sample version of the repo's viewer.html, with data injected at __DATA__.
_VIEWER_TEMPLATE = """<!DOCTYPE html>
<html><head><meta charset='utf-8'>
<link href='https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700&display=swap' rel='stylesheet'>
<style>
:root{--hl:#6366f1;--border:#e9ebf0;--muted:#9aa1b1;}
body{font-family:'Inter',system-ui,-apple-system,BlinkMacSystemFont,'Segoe UI',Roboto,sans-serif;margin:0;color:#2a2f3a;background:transparent;-webkit-font-smoothing:antialiased;}
.hint{font-size:12px;color:var(--muted);padding:4px 2px 12px;}
.panels{display:flex;gap:14px;align-items:flex-start;}
.panel{flex:1;min-width:0;border:1px solid var(--border);border-radius:14px;padding:16px 18px;max-height:420px;overflow:auto;line-height:2.3;background:#fff;box-shadow:0 1px 2px rgba(16,24,40,.04),0 4px 16px rgba(16,24,40,.05);}
.panel h2{font-size:10.5px;text-transform:uppercase;letter-spacing:.08em;font-weight:600;color:var(--muted);margin:0 0 12px;}
.word{padding:2px 4px;border-radius:6px;cursor:default;transition:background .12s,color .12s,box-shadow .12s;}
.word:hover{background:rgba(99,102,241,.10);}
/* SimAlign (unconstrained) target = grey */
.word.aligned-sim{background:#d1d5db;color:#1f2937;}
/* selected method target = blue */
.word.active,.word.aligned{background:var(--hl);color:#fff;}
/* both methods agree = blue fill + grey frame */
.word.aligned.aligned-sim{background:var(--hl);color:#fff;box-shadow:0 0 0 2px #fff,0 0 0 3px #9ca3af;}
.word.active{box-shadow:0 0 0 2px #fff,0 0 0 4px var(--hl);}
.word[data-un]{color:#b8bdc9;}
#info{font-size:12px;color:#4b5160;padding:12px 2px 2px;min-height:16px;}
.legend{font-size:11.5px;color:var(--muted);}
.legend b{display:inline-block;width:10px;height:10px;border-radius:3px;vertical-align:-1px;margin:0 3px 0 8px;}
</style></head><body>
<div class='hint'>Hover a word to highlight its counterpart(s).
<span class='legend'><b style='background:#6366f1'></b>selected method<b style='background:#d1d5db'></b>SimAlign<b style='background:#6366f1;box-shadow:0 0 0 2px #fff,0 0 0 3px #9ca3af'></b>both agree</span></div>
<div class='panels'>
  <div class='panel' id='pa'><h2>Text A</h2></div>
  <div class='panel' id='pb'><h2>Text B</h2></div>
</div>
<div id='info'></div>
<script>
const DATA = __DATA__;
const a2b={},b2a={},sa2b={},sb2a={};  // selected-method and SimAlign maps
for(const p of DATA.pairs){const i=p[0],j=p[1];(a2b[i]=a2b[i]||[]).push(j);(b2a[j]=b2a[j]||[]).push(i);}
for(const p of (DATA.sim_pairs||[])){const i=p[0],j=p[1];(sa2b[i]=sa2b[i]||[]).push(j);(sb2a[j]=sb2a[j]||[]).push(i);}
const info=document.getElementById('info');
function el(side,idx){return document.querySelector('.word[data-side="'+side+'"][data-idx="'+idx+'"]');}
function clear(){document.querySelectorAll('.word.active,.word.aligned,.word.aligned-sim').forEach(e=>e.classList.remove('active','aligned','aligned-sim'));info.textContent='';}
function fmt(list,ow){return list.map(function(j){return j+' \\u201c'+(ow[j]||'?')+'\\u201d';}).join(', ');}
function hi(side,idx){clear();const self=el(side,idx);if(self)self.classList.add('active');
  const other=side==='a'?'b':'a';const ow=side==='a'?DATA.b:DATA.a;
  const sw=(side==='a'?DATA.a:DATA.b);
  const mtg=(side==='a'?a2b:b2a)[idx]||[];
  const stg=(side==='a'?sa2b:sb2a)[idx]||[];
  // Grey (SimAlign) first, then blue (method) so words both agree on get both classes.
  stg.forEach(function(j){const e=el(other,j);if(e)e.classList.add('aligned-sim');});
  mtg.forEach(function(j){const e=el(other,j);if(e)e.classList.add('aligned');});
  // Jump only to the primary (method) correspondence; never follow a possibly
  // distant SimAlign-only match (it gets a directional hint in the info bar instead).
  if(mtg.length){const f=el(other,mtg[0]);if(f)f.scrollIntoView({block:'nearest'});}
  let txt=(side==='a'?'Text A':'Text B')+' word '+idx+' \\u201c'+sw[idx]+'\\u201d';
  txt += mtg.length ? ' \\u2194 '+fmt(mtg,ow) : ' \\u2014 no method alignment';
  const extra=stg.filter(function(j){return mtg.indexOf(j)<0;});  // SimAlign-only targets
  if(extra.length){
    const srcRel=idx/(sw.length||1), tgtRel=extra[0]/(ow.length||1);
    const arrow = tgtRel>srcRel ? ' \\u2193' : ' \\u2191';   // scroll down / up to find it
    txt += '  \\u00b7  SimAlign: '+fmt(extra,ow)+arrow;
  }
  info.textContent=txt;}
function build(side,words,mmap,smap){const p=document.getElementById(side==='a'?'pa':'pb');
  words.forEach(function(w,idx){const s=document.createElement('span');s.className='word';s.textContent=w;s.dataset.side=side;s.dataset.idx=idx;
    if(!(idx in mmap)&&!(idx in smap))s.dataset.un='1';   // dim only if neither method aligns it
    s.addEventListener('mouseenter',function(){hi(side,idx);});
    s.addEventListener('mouseleave',clear);
    p.appendChild(s);p.appendChild(document.createTextNode(' '));});}
build('a',DATA.a,a2b,sa2b);build('b',DATA.b,b2a,sb2a);
</script></body></html>"""


def _preload():
    """Warm encoders on Space startup so the first alignment is fast.

    Loads the pinned small models (kept resident) plus the default model (in the
    rotating slot). Failures are non-fatal (the model still loads lazily on first
    use) so a download hiccup doesn't take the whole app down.
    """
    for display in PINNED_MODELS + [DEFAULT_MODEL]:
        try:
            hf_id = MODELS[display]
            _get_encoder(hf_id, _resolve_layer(hf_id, None))
        except Exception as e:
            print(f"Preload of {display} failed (will load on first use): {e}")


_preload()


if __name__ == "__main__":
    demo.queue().launch()