src/streamlit_app.py

# streamlit_app.py — Streamlit HF Space comparing three LinkjeBERT2 checkpoints
# side by side. Binary token classification (O / LINK) fine-tuned from
# mdeberta-v3-base. Each checkpoint is a different training step of the same run;
# they are loaded from subfolders of dejanseo/LinkjeBERT2. Long inputs are handled
# with overlapping sliding windows; per-token LINK probabilities are averaged
# across overlaps, then promoted to whole words.
import streamlit as st
import torch
import torch.nn.functional as F
from transformers import AutoTokenizer, AutoModelForTokenClassification

MODEL_ID = "dejanseo/LinkjeBERT2"

# (subfolder, short label) — step 42732 is the eval-loss minimum (best-calibrated
# probabilities), 64098 is near the eval-f1 peak (highest recall), 53415 is in
# between.
CHECKPOINTS = [
    ("42732", "42732 · loss min / best-calibrated"),
    ("53415", "53415 · mid"),
    ("64098", "64098 · f1 peak / highest recall"),
]

st.set_page_config(page_title="LinkjeBERT2 — checkpoint compare", layout="wide")


@st.cache_resource
def load_tokenizer():
    # Tokenizer is identical across checkpoints; load once from the repo root.
    return AutoTokenizer.from_pretrained(MODEL_ID, use_fast=True)


@st.cache_resource
def load_model(subfolder):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = AutoModelForTokenClassification.from_pretrained(MODEL_ID, subfolder=subfolder)
    model = model.to(device)
    model.eval()
    return model, device


def group_tokens_into_words(tokens, offset_mapping, link_probs):
    """Group subword tokens into whole words (DeBERTa marks word starts with the
    SentencePiece boundary marker) and carry each token's LINK probability."""
    words = []
    cur_tokens, cur_offsets, cur_probs = [], [], []

    def flush():
        if cur_tokens:
            words.append({'tokens': list(cur_tokens),
                          'offsets': list(cur_offsets),
                          'probs': list(cur_probs)})

    for i, (token, offsets, prob) in enumerate(zip(tokens, offset_mapping, link_probs)):
        if offsets == [0, 0]:  # special token
            flush()
            cur_tokens.clear(); cur_offsets.clear(); cur_probs.clear()
            continue

        is_new_word = False
        if token.startswith("▁"):  # DeBERTa/SentencePiece word boundary
            is_new_word = True
        elif not token.startswith("##"):
            if i == 0 or offset_mapping[i - 1] == [0, 0]:
                is_new_word = True
            elif cur_offsets and offsets[0] > cur_offsets[-1][1]:
                is_new_word = True

        if is_new_word and cur_tokens:
            flush()
            cur_tokens.clear(); cur_offsets.clear(); cur_probs.clear()

        cur_tokens.append(token)
        cur_offsets.append(offsets)
        cur_probs.append(prob)

    flush()
    return words


def predict_links(text, tokenizer, model, device, threshold=0.5,
                  max_length=512, doc_stride=128):
    """Return (link_spans, link_details) for LINK words above `threshold`."""
    if not text.strip():
        return [], []

    full_enc = tokenizer(text, add_special_tokens=False, truncation=False,
                         return_offsets_mapping=True)
    all_ids = full_enc["input_ids"]
    all_offsets = full_enc["offset_mapping"]
    n_tokens = len(all_ids)
    if n_tokens == 0:
        return [], []

    prob_sums = [0.0] * n_tokens
    prob_counts = [0] * n_tokens

    cls_id = tokenizer.cls_token_id
    sep_id = tokenizer.sep_token_id
    cap = max_length - 2  # room for CLS + SEP
    step = max(cap - doc_stride, 1)

    start = 0
    while start < n_tokens:
        end = min(start + cap, n_tokens)
        window_ids = all_ids[start:end]

        input_ids = torch.tensor(
            [[cls_id] + window_ids + [sep_id]],
            device=device,
        )
        attention_mask = torch.ones_like(input_ids)

        with torch.no_grad():
            logits = model(input_ids=input_ids, attention_mask=attention_mask).logits
            probs = F.softmax(logits, dim=-1)[0].cpu()
            content_probs = probs[1:-1, 1].tolist()  # LINK prob, minus CLS/SEP

        for i, p in enumerate(content_probs):
            orig_idx = start + i
            if orig_idx < n_tokens:
                prob_sums[orig_idx] += p
                prob_counts[orig_idx] += 1

        if end == n_tokens:
            break
        start += step

    link_probs = [prob_sums[i] / prob_counts[i] if prob_counts[i] > 0 else 0.0
                  for i in range(n_tokens)]

    tokens = tokenizer.convert_ids_to_tokens(all_ids)
    offset_mapping = [list(o) for o in all_offsets]
    words = group_tokens_into_words(tokens, offset_mapping, link_probs)

    link_spans, link_details = [], []
    for w in words:
        if any(p >= threshold for p in w['probs']):
            s = w['offsets'][0][0]
            e = w['offsets'][-1][1]
            link_spans.append((s, e))
            link_details.append({
                "text": text[s:e],
                "start": s,
                "end": e,
                "max_confidence": round(max(w['probs']), 4),
                "avg_confidence": round(sum(w['probs']) / len(w['probs']), 4),
            })

    # Merge adjacent/touching highlighted words into contiguous spans.
    merged = []
    for s, e in sorted(link_spans):
        if merged and s <= merged[-1][1] + 1:
            merged[-1] = (merged[-1][0], max(merged[-1][1], e))
        else:
            merged.append((s, e))

    return merged, link_details


def render_highlighted_text(text, link_spans):
    if not text:
        return ""
    link_spans = sorted(link_spans, key=lambda x: x[0])
    parts = []
    last_end = 0
    for start, end in link_spans:
        if start > last_end:
            parts.append(text[last_end:start])
        parts.append(
            f'<span style="background-color:#90EE90;padding:2px 4px;'
            f'border-radius:3px;font-weight:500;">{text[start:end]}</span>'
        )
        last_end = end
    if last_end < len(text):
        parts.append(text[last_end:])

    return (
        '<div style="padding:20px;background-color:#f8f9fa;border-radius:8px;'
        'line-height:1.8;font-size:16px;white-space:pre-wrap;word-wrap:break-word;'
        'color:#111;font-family:-apple-system,BlinkMacSystemFont,Segoe UI,Roboto,sans-serif;">'
        + "".join(parts) + "</div>"
    )


SAMPLE = (
    "Het Nederlandse kabinet presenteerde vandaag een nieuw plan voor duurzame "
    "energie. Volgens de minister moeten windmolens op zee de komende jaren fors "
    "worden uitgebreid. Onderzoekers van de Technische Universiteit Delft "
    "waarschuwen echter dat de netcapaciteit tekortschiet."
)


def main():
    st.subheader("LinkjeBERT — Dutch Natural Link Prediction Model")
    st.write("Model trained by [DEJAN AI](https://dejan.ai/)")
    st.caption(f"Model: {MODEL_ID} · checkpoints " + ", ".join(c[0] for c in CHECKPOINTS))

    try:
        tokenizer = load_tokenizer()
        models = {sub: load_model(sub) for sub, _ in CHECKPOINTS}
        device = next(iter(models.values()))[1]
        st.success(f"Loaded {len(models)} checkpoints on {device}")
    except Exception as e:
        st.error(f"Failed to load models: {e}")
        return

    threshold = st.slider(
        "Confidence threshold (%)",
        min_value=0, max_value=100, value=50, step=1,
        help="A word is highlighted if any of its tokens reaches this LINK probability. "
             "Same threshold is applied to all three checkpoints.",
    ) / 100.0

    text = st.text_area("Input text (Dutch):", value=SAMPLE, height=220)

    if st.button("Detect links"):
        if not text.strip():
            st.warning("Please enter text.")
            return

        columns = st.columns(len(CHECKPOINTS))
        for (sub, label), col in zip(CHECKPOINTS, columns):
            model, device = models[sub]
            link_spans, link_details = predict_links(
                text, tokenizer, model, device, threshold
            )
            with col:
                st.subheader(label)
                st.markdown(render_highlighted_text(text, link_spans),
                            unsafe_allow_html=True)
                st.info(f"{len(link_details)} anchor word(s) @ {threshold:.0%}")
                if link_details:
                    with st.expander("Details"):
                        st.json(link_details)


if __name__ == "__main__":
    main()