Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| import torch | |
| from transformers import AutoModelForTokenClassification, AutoTokenizer | |
| import numpy as np | |
| import re | |
| import markdown | |
| import requests | |
| from markdownify import markdownify | |
| from bs4 import BeautifulSoup | |
| from html import escape as html_escape | |
| import logging | |
| from typing import Optional, Dict, List, Tuple | |
| # --- HIDE STREAMLIT MENU / PAGE CONFIG --- | |
| st.set_page_config( | |
| initial_sidebar_state="collapsed", | |
| layout="wide", | |
| page_title="LinkjeBERT by DEJAN AI" | |
| ) | |
| hide_streamlit_style = """ | |
| <style> | |
| #MainMenu {visibility: hidden;} | |
| </style> | |
| """ | |
| st.markdown(hide_streamlit_style, unsafe_allow_html=True) | |
| st.logo( | |
| image="https://dejan.ai/wp-content/uploads/2024/02/dejan-300x103.png", | |
| link="https://dejan.ai/", | |
| size="large" | |
| ) | |
| logging.basicConfig(level=logging.INFO) | |
| logger = logging.getLogger(__name__) | |
| # ---------------------------------- | |
| # Config | |
| # ---------------------------------- | |
| MODEL_ID = "dejanseo/LinkjeBERT3" # main model at the repo root (no subfolder) | |
| MAX_LENGTH = 512 | |
| DOC_STRIDE = 128 | |
| DEVICE = "cuda" if torch.cuda.is_available() else "cpu" | |
| # Highlight opacity scales linearly from the document's most link-likely word | |
| # (peak = full opacity) down to 0. Words whose relative opacity falls below this | |
| # floor render as plain text — a display detail to avoid tinting the whole doc, | |
| # not a probability threshold. | |
| MIN_REL_OPACITY = 0.05 | |
| # --- URL scraping --- | |
| SCRAPE_UA = ( | |
| "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 " | |
| "(KHTML, like Gecko) Chrome/125.0.0.0 Safari/537.36" | |
| ) | |
| SCRAPE_TIMEOUT = 20 | |
| # Tags stripped before content detection (boilerplate / non-content). | |
| STRIP_TAGS = ['script', 'style', 'noscript', 'nav', 'header', 'footer', | |
| 'aside', 'form', 'iframe', 'svg', 'button'] | |
| # Content-area detection: semantic elements first, then common CSS id/class | |
| # signals. First strong match (or the largest candidate) wins. | |
| SEMANTIC_SELECTORS = ['article', 'main', '[role="main"]'] | |
| CSS_SELECTORS = ( | |
| '#content, #main, #article, [itemprop="articleBody"], ' | |
| '.post-content, .entry-content, .article-content, .article-body, ' | |
| '.post-body, .story-body, .content__article-body, .article__body, ' | |
| '.rich-text, .prose, .content' | |
| ) | |
| MIN_CONTENT_CHARS = 200 # a container must hold at least this much text to count | |
| # ---------------------------------- | |
| # Load tokenizer / model from the Hugging Face Hub | |
| # ---------------------------------- | |
| def load_tokenizer(): | |
| return AutoTokenizer.from_pretrained(MODEL_ID, use_fast=True) | |
| def load_model(): | |
| model = AutoModelForTokenClassification.from_pretrained(MODEL_ID) | |
| model.to(DEVICE) | |
| model.eval() | |
| return model | |
| # ---------------------------------- | |
| # Inference helpers | |
| # ---------------------------------- | |
| def windowize_inference( | |
| plain_text: str, tokenizer: AutoTokenizer, max_length: int, doc_stride: int | |
| ) -> List[Dict]: | |
| """Slice long text into overlapping windows for inference.""" | |
| specials = tokenizer.num_special_tokens_to_add(pair=False) | |
| cap = max_length - specials | |
| full_encoding = tokenizer( | |
| plain_text, add_special_tokens=False, return_offsets_mapping=True, truncation=False | |
| ) | |
| temp_tokenization = tokenizer(plain_text, truncation=False) | |
| full_word_ids = temp_tokenization.word_ids(batch_index=0) | |
| windows_data = [] | |
| step = max(cap - doc_stride, 1) | |
| start_token_idx = 0 | |
| total_tokens = len(full_encoding["input_ids"]) | |
| if total_tokens == 0 and len(plain_text) > 0: | |
| logger.warning("Tokenizer produced 0 tokens for a non-empty string.") | |
| return [] | |
| while start_token_idx < total_tokens: | |
| end_token_idx = min(start_token_idx + cap, total_tokens) | |
| ids_slice = full_encoding["input_ids"][start_token_idx:end_token_idx] | |
| offsets_slice = full_encoding["offset_mapping"][start_token_idx:end_token_idx] | |
| word_ids_slice = [] | |
| current_token = 0 | |
| for i, wid in enumerate(full_word_ids): | |
| if temp_tokenization.token_to_chars(i) is not None: | |
| if current_token >= start_token_idx and current_token < end_token_idx: | |
| word_ids_slice.append(wid) | |
| current_token += 1 | |
| input_ids = [tokenizer.cls_token_id] + ids_slice + [tokenizer.sep_token_id] | |
| attention_mask = [1] * len(input_ids) | |
| padding_length = max_length - len(input_ids) | |
| input_ids.extend([tokenizer.pad_token_id] * padding_length) | |
| attention_mask.extend([0] * padding_length) | |
| window_offset_mapping = [(0, 0)] + offsets_slice + [(0, 0)] | |
| window_offset_mapping += [(0, 0)] * padding_length | |
| window_word_ids = [None] + word_ids_slice + [None] | |
| window_word_ids += [None] * padding_length | |
| windows_data.append({ | |
| "input_ids": torch.tensor(input_ids, dtype=torch.long), | |
| "attention_mask": torch.tensor(attention_mask, dtype=torch.long), | |
| "word_ids": window_word_ids[:max_length], | |
| "offset_mapping": window_offset_mapping[:max_length], | |
| }) | |
| if end_token_idx >= total_tokens: | |
| break | |
| start_token_idx += step | |
| return windows_data | |
| def build_highlighted_html(raw_text: str, char_probs: np.ndarray, peak: float) -> str: | |
| """Render raw_text as Markdown -> HTML, then inject confidence-heatmap spans | |
| into the rendered TEXT nodes so headings / bold / lists survive. Each rendered | |
| word is matched back to its position in the raw Markdown (order is preserved, | |
| markers like #, **, - are simply skipped over) to read its LINK probability. | |
| Word opacity is linear and relative to the document peak.""" | |
| soup = BeautifulSoup(markdown.markdown(raw_text), "html.parser") | |
| cursor = 0 | |
| for node in list(soup.find_all(string=True)): | |
| s = str(node) | |
| if not s.strip(): | |
| continue | |
| out = [] | |
| for token in re.findall(r"\S+|\s+", s): | |
| if not token.strip(): | |
| out.append(html_escape(token)) | |
| continue | |
| idx = raw_text.find(token, cursor) | |
| if idx == -1: | |
| out.append(html_escape(token)) | |
| continue | |
| start, end = idx, idx + len(token) | |
| cursor = end | |
| wprob = float(char_probs[start:end].max()) if start < len(char_probs) else 0.0 | |
| rel = (wprob / peak) if peak > 0 else 0.0 | |
| if rel >= MIN_REL_OPACITY: | |
| out.append( | |
| f"<span style='background-color: rgba(100, 200, 100, {rel:.3f}); " | |
| f"padding: 0.1em 0.2em; border-radius: 0.2em;' " | |
| f"title='Link Probability: {wprob:.1%}'>{html_escape(token)}</span>" | |
| ) | |
| else: | |
| out.append(html_escape(token)) | |
| node.replace_with(BeautifulSoup("".join(out), "html.parser")) | |
| return str(soup) | |
| def classify_text(text: str, model, tokenizer) -> Tuple[str, float, Optional[str]]: | |
| """Run inference ONCE on the raw Markdown, then render a Markdown-formatted | |
| confidence heatmap. Returns (html, peak_prob, warning).""" | |
| if not text.strip(): | |
| return "", 0.0, "Input text is empty." | |
| windows = windowize_inference(text, tokenizer, MAX_LENGTH, DOC_STRIDE) | |
| if not windows: | |
| return "", 0.0, "Could not generate any windows for processing." | |
| char_link_probabilities = np.zeros(len(text), dtype=np.float32) | |
| with torch.no_grad(): | |
| for window in windows: | |
| inputs = { | |
| 'input_ids': window['input_ids'].unsqueeze(0).to(DEVICE), | |
| 'attention_mask': window['attention_mask'].unsqueeze(0).to(DEVICE) | |
| } | |
| outputs = model(**inputs) | |
| probabilities = torch.softmax(outputs.logits, dim=-1).squeeze(0) | |
| link_probs = probabilities[:, 1].cpu().numpy() | |
| for i, offset in enumerate(window['offset_mapping']): | |
| if isinstance(offset, (list, tuple)) and len(offset) == 2: | |
| start, end = offset | |
| if window['word_ids'][i] is not None and start < end: | |
| char_link_probabilities[start:end] = np.maximum( | |
| char_link_probabilities[start:end], link_probs[i] | |
| ) | |
| # Peak = the highest per-character LINK probability in the document (i.e. the | |
| # single most link-likely word). Every word's opacity is expressed relative to | |
| # it, so the brightest highlight is always the model's top pick. | |
| peak_prob = float(char_link_probabilities.max()) if len(char_link_probabilities) else 0.0 | |
| html = build_highlighted_html(text, char_link_probabilities, peak_prob) | |
| return html, peak_prob, None | |
| # ---------------------------------- | |
| # URL scraping -> Markdown | |
| # ---------------------------------- | |
| def detect_content_element(soup: BeautifulSoup): | |
| """Find the main content container using semantic HTML then CSS signals. | |
| Returns the element, or None if nothing substantial is found.""" | |
| for sel in SEMANTIC_SELECTORS: | |
| el = soup.select_one(sel) | |
| if el and len(el.get_text(strip=True)) >= MIN_CONTENT_CHARS: | |
| return el | |
| best, best_len = None, 0 | |
| for el in soup.select(CSS_SELECTORS): | |
| n = len(el.get_text(strip=True)) | |
| if n > best_len: | |
| best, best_len = el, n | |
| return best if best_len >= MIN_CONTENT_CHARS else None | |
| def html_to_markdown(el) -> str: | |
| """Convert an HTML element to Markdown. Anchor tags are stripped to their | |
| plain text (strip=['a']) so the input matches the link-stripped Markdown the | |
| model was trained on and existing links don't leak the answer.""" | |
| return markdownify(str(el), strip=['a'], heading_style='ATX').strip() | |
| def fetch_and_extract(url: str) -> Tuple[Optional[str], Optional[str]]: | |
| """Fetch a URL and return (markdown, error). On a detection miss a best-effort | |
| body extraction is returned alongside a non-fatal error message so the caller | |
| can offer it as editable manual-fallback text.""" | |
| url = url.strip() | |
| if not url: | |
| return None, "Please enter a URL." | |
| if not url.startswith(('http://', 'https://')): | |
| url = 'https://' + url | |
| try: | |
| resp = requests.get(url, headers={'User-Agent': SCRAPE_UA}, timeout=SCRAPE_TIMEOUT) | |
| resp.raise_for_status() | |
| except Exception as e: | |
| return None, f"Could not fetch the URL ({e})." | |
| soup = BeautifulSoup(resp.text, 'html.parser') | |
| for t in soup(STRIP_TAGS): | |
| t.decompose() | |
| container = detect_content_element(soup) | |
| if container is not None: | |
| md = html_to_markdown(container) | |
| if md.strip(): | |
| return md, None | |
| # Fallback: no content area detected -> best-effort whole-body extraction. | |
| body = soup.body or soup | |
| return html_to_markdown(body), "Couldn't detect a main content area — showing a best-effort extraction. Edit it below or use the Text tab." | |
| # ---------------------------------- | |
| # Streamlit UI | |
| # ---------------------------------- | |
| st.title("LinkjeBERT3") | |
| st.caption(f"Model: {MODEL_ID} · device {DEVICE}") | |
| if 'result' not in st.session_state: | |
| st.session_state.result = None | |
| if 'user_input' not in st.session_state: | |
| st.session_state.user_input = """# Nieuw kabinetsplan voor duurzame energie op zee | |
| Het Nederlandse kabinet presenteerde vandaag een **ambitieus plan** om de opwekking van **windenergie op zee** de komende tien jaar fors uit te breiden. Volgens de minister van Klimaat en Groene Groei is de uitbreiding nodig om de klimaatdoelen van 2035 te halen. | |
| ## Wat er verandert | |
| Onderzoekers van de *Technische Universiteit Delft* waarschuwen echter dat de netcapaciteit op land tekortschiet. Zonder forse investeringen in het hoogspanningsnet dreigt een deel van de opgewekte stroom verloren te gaan. | |
| De belangrijkste maatregelen op een rij: | |
| - Aanleg van drie nieuwe windparken voor de kust van **Noord-Holland** | |
| - Een verdubbeling van de subsidie voor waterstofproductie | |
| - Strengere eisen aan netbeheerders zoals TenneT | |
| > "We kunnen niet blijven bouwen aan windmolens als het net ze niet aankan", aldus een woordvoerder van de netbeheerder. | |
| ## Reacties uit de sector | |
| Brancheorganisatie NVDE noemde het plan een stap in de goede richting, maar hoorde naar eigen zeggen te weinig over de financiering. Volgens *Bloomberg* zou de totale investering kunnen oplopen tot 40 miljard euro. | |
| Meer details over de plannen zijn te vinden in de Kamerbrief die vandaag is gepubliceerd.""" | |
| try: | |
| tokenizer = load_tokenizer() | |
| model = load_model() | |
| except Exception as e: | |
| st.error(f"Failed to load model: {e}") | |
| st.stop() | |
| def do_classify(text: str): | |
| if not text.strip(): | |
| st.warning("Please enter some text to classify.") | |
| st.session_state.result = None | |
| return | |
| with st.spinner("Analyzing text..."): | |
| html, peak_prob, warning = classify_text(text, model, tokenizer) | |
| st.session_state.result = {"html": html, "peak": peak_prob, "warning": warning} | |
| tab_text, tab_url = st.tabs(["Text", "URL"]) | |
| with tab_text: | |
| user_input = st.text_area( | |
| "Paste your text here (Dutch, Markdown-formatted):", | |
| st.session_state.user_input, | |
| height=200, | |
| key="text_area" | |
| ) | |
| if st.button("Classify Text", type="primary", use_container_width=True): | |
| st.session_state.user_input = user_input | |
| do_classify(user_input) | |
| with tab_url: | |
| url = st.text_input("Article URL", placeholder="https://www.example.nl/artikel", key="url_input") | |
| if st.button("Fetch & Classify", type="primary", use_container_width=True): | |
| with st.spinner("Fetching and extracting content..."): | |
| md, err = fetch_and_extract(url) | |
| st.session_state.url_error = err | |
| st.session_state.url_fallback = md or "" | |
| if md and not err: | |
| # Clean content detected -> convert to MD and classify straight away. | |
| st.session_state.user_input = md | |
| do_classify(md) | |
| else: | |
| # Fetch failed or no content area -> hold for manual fallback below. | |
| st.session_state.result = None | |
| if st.session_state.get("url_error"): | |
| st.warning(st.session_state.url_error) | |
| fallback_text = st.text_area( | |
| "Extracted text (edit if needed, then classify):", | |
| st.session_state.get("url_fallback", ""), | |
| height=200, | |
| key="url_fallback_area", | |
| ) | |
| if st.button("Classify this text", use_container_width=True, key="url_fallback_btn"): | |
| st.session_state.user_input = fallback_text | |
| do_classify(fallback_text) | |
| if st.session_state.result: | |
| res = st.session_state.result | |
| if res["warning"]: | |
| st.warning(res["warning"]) | |
| else: | |
| st.caption( | |
| f"Peak link probability: {res['peak']:.1%} (brightest highlight). " | |
| "Opacity scales linearly down to that word — hover any word for its exact probability." | |
| ) | |
| st.markdown("---") | |
| st.markdown(res["html"], unsafe_allow_html=True) | |
| st.divider() | |
| st.markdown(""" | |
| ## SEO Use Cases | |
| LinkBERT's applications are vast and diverse, tailored to enhance both the efficiency and quality of web content creation and analysis: | |
| - **Anchor Text Suggestion:** Acts as a mechanism during internal link optimization, suggesting potential anchor texts to web authors. | |
| - **Evaluation of Existing Links:** Assesses the naturalness of link placements within existing content, aiding in the refinement of web pages. | |
| - **Link Placement Guide:** Offers guidance to link builders by suggesting optimal placement for links within content. | |
| - **Anchor Text Idea Generator:** Provides creative anchor text suggestions to enrich content and improve SEO strategies. | |
| - **Spam and Inorganic SEO Detection:** Helps identify unnatural link patterns, contributing to the detection of spam and inorganic SEO tactics. | |
| ## Engage Our Team | |
| - Interested in using this in an automated pipeline for bulk link prediction? | |
| - Please [book an appointment](https://dejan.ai/call/) to discuss your needs. | |
| """) | |