# 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'{text[start:end]}' ) last_end = end if last_end < len(text): parts.append(text[last_end:]) return ( '
' + "".join(parts) + "
" ) 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()