Delete src

src/.streamlit/config.toml

@@ -1,24 +0,0 @@
-[server]
-enableStaticServing = true
-toolbarMode = "viewer"
-
-[[theme.fontFaces]]
-family="montserrat-sans"
-url="app/static/Montserrat-Italic-VariableFont_wght.ttf"
-style="italic"
-weight=500
-[[theme.fontFaces]]
-family="montserrat-sans"
-url="app/static/Montserrat-VariableFont_wght.ttf"
-style="normal"
-weight=500
-[[theme.fontFaces]]
-family="noto-mono"
-url="app/static/NotoSansMono-VariableFont_wdth,wght.ttf"
-
-[theme]
-font="montserrat-sans, noto-sans, sans-serif"
-codeFont="noto-mono, monospace"
-baseFontSize=16
-primaryColor="#28a745"
-backgroundColor="#FFFFFF"

src/app.py

@@ -1,298 +0,0 @@
-import streamlit as st
-import torch
-# Use AutoModel and AutoTokenizer for easier loading from Hub
-from transformers import AutoModelForTokenClassification, AutoTokenizer
-import numpy as np
-import logging
-from dataclasses import dataclass
-from typing import Optional, Dict, List, Tuple
-
-# --- HIDE STREAMLIT MENU ---
-st.set_page_config(
-    initial_sidebar_state="collapsed"
-)
-
-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/",
-)
-
-# ----------------------------------
-# Logging
-# ----------------------------------
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-# ----------------------------------
-# Config
-# ----------------------------------
-
-@dataclass
-class AppConfig:
-    """Configuration for the LinkBERT application"""
-    # <<< CHANGE 1: Point to the Hugging Face Hub repository >>>
-    model_name: str = "dejanseo/link-prediction"
-    max_length: int = 512
-    doc_stride: int = 128
-    device: str = "cuda" if torch.cuda.is_available() else "cpu"
-
-# ----------------------------------
-# Load model/tokenizer from Hugging Face Hub
-# ----------------------------------
-
-@st.cache_resource
-def load_model_from_hub():
-    """Loads the fine-tuned model and tokenizer from the Hugging Face Hub."""
-    config = AppConfig()
-
-    logger.info(f"Loading model and tokenizer from Hugging Face Hub: {config.model_name}...")
-
-    # <<< CHANGE 2: Use Auto* classes for direct loading from the Hub >>>
-    model = AutoModelForTokenClassification.from_pretrained(config.model_name)
-    tokenizer = AutoTokenizer.from_pretrained(config.model_name)
-
-    model.to(config.device)
-    model.eval()
-
-    logger.info("Model and tokenizer loaded successfully.")
-    return model, tokenizer, config.device, config.max_length, config.doc_stride
-
-# <<< CHANGE 3: Call the new loading function >>>
-model, tokenizer, device, MAX_LENGTH, DOC_STRIDE = load_model_from_hub()
-
-
-# ----------------------------------
-# 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]
-
-        # Properly slice word_ids based on character spans
-        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.build_inputs_with_special_tokens(ids_slice)
-        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)
-
-        # Pad offset mapping correctly
-        window_offset_mapping = tokenizer.build_inputs_with_special_tokens([]) # Get special tokens offsets
-        window_offset_mapping = window_offset_mapping[:-1] + offsets_slice + window_offset_mapping[-1:]
-        window_offset_mapping += [(0, 0)] * padding_length
-
-        window_word_ids = [None] + word_ids_slice + [None] * (padding_length + 1)
-
-        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 classify_text(
-    text: str, prediction_threshold_percent: float
-) -> Tuple[str, Optional[str]]:
-    """Classify link tokens with windowing. Returns (html, warning)."""
-    if not text.strip(): return "", "Input text is empty."
-
-    windows = windowize_inference(text, tokenizer, MAX_LENGTH, DOC_STRIDE)
-    if not windows: return "", "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)
-            }
-            # <<< CHANGE 4: The output object from AutoModel has a 'logits' attribute directly >>>
-            outputs = model(**inputs)
-            probabilities = torch.softmax(outputs.logits, dim=-1).squeeze(0)
-            link_probs = probabilities[:, 1].cpu().numpy()
-
-            for i, (start, end) in enumerate(window['offset_mapping']):
-                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]
-                    )
-
-    final_threshold = prediction_threshold_percent / 100.0
-
-    full_encoding = tokenizer(text, return_offsets_mapping=True, truncation=False)
-    word_ids = full_encoding.word_ids(batch_index=0)
-    offsets = full_encoding['offset_mapping']
-
-    word_max_prob_map: Dict[int, float] = {}
-    word_char_spans: Dict[int, List[int]] = {}
-
-    for i, word_id in enumerate(word_ids):
-        if word_id is not None:
-            start_char, end_char = offsets[i]
-            if start_char < end_char:
-                current_token_max_prob = np.max(char_link_probabilities[start_char:end_char]) if np.any(char_link_probabilities[start_char:end_char]) else 0.0
-
-                if word_id not in word_max_prob_map:
-                    word_max_prob_map[word_id] = current_token_max_prob
-                    word_char_spans[word_id] = [start_char, end_char]
-                else:
-                    word_max_prob_map[word_id] = max(word_max_prob_map[word_id], current_token_max_prob)
-                    word_char_spans[word_id][1] = end_char
-
-    highlight_candidates: Dict[int, float] = {}
-    for word_id, max_prob in word_max_prob_map.items():
-        if max_prob >= final_threshold:
-            highlight_candidates[word_id] = max_prob
-
-    max_highlight_prob = 0.0
-    if highlight_candidates:
-        max_highlight_prob = max(highlight_candidates.values())
-
-    html_parts, current_char = [], 0
-    sorted_word_ids = sorted(word_char_spans.keys(), key=lambda k: word_char_spans[k][0])
-
-    for word_id in sorted_word_ids:
-        start_char, end_char = word_char_spans[word_id]
-
-        if start_char > current_char:
-            html_parts.append(text[current_char:start_char])
-
-        word_text = text[start_char:end_char]
-
-        if word_id in highlight_candidates:
-            word_prob = highlight_candidates[word_id]
-            normalized_opacity = 1.0
-            if max_highlight_prob > 0:
-                normalized_opacity = (word_prob / max_highlight_prob) * 0.9 + 0.1
-
-            base_bg_color = "#D4EDDA"
-            base_text_color = "#155724"
-
-            html_parts.append(f"<span style='background-color: {base_bg_color}; color: {base_text_color}; "
-                              f"padding: 0.1em 0.2em; border-radius: 0.2em; opacity: {normalized_opacity:.2f};'>"
-                              f"{word_text}</span>")
-        else:
-            html_parts.append(word_text)
-        current_char = end_char
-
-    if current_char < len(text):
-        html_parts.append(text[current_char:])
-
-    return "".join(html_parts), None
-
-# ----------------------------------
-# Streamlit UI (No changes needed from here down)
-# ----------------------------------
-st.set_page_config(layout="wide", page_title="LinkBERT by DEJAN AI")
-st.title("LinkBERT")
-
-DEFAULT_THRESHOLD = 70.0
-THRESHOLD_STEP = 10.0
-THRESHOLD_BOUNDARY_PERCENT = 10.0
-
-if 'current_threshold' not in st.session_state:
-    st.session_state.current_threshold = DEFAULT_THRESHOLD
-if 'output_html' not in st.session_state:
-    st.session_state.output_html = ""
-if 'user_input' not in st.session_state:
-    st.session_state.user_input = "DEJAN AI is the world's leading AI SEO agency. This tool showcases the capability of our latest link prediction model called LinkBERT. This model is trained on the highest quality organic link data and can predict natural link placement in plain text."
-
-user_input = st.text_area(
-    "Paste your text here:",
-    st.session_state.user_input,
-    height=200,
-    key="text_area"
-)
-
-with st.expander('Settings'):
-    slider_threshold = st.slider(
-        "Link Probability Threshold (%)",
-        min_value=0, max_value=100, value=int(st.session_state.current_threshold), step=1,
-        help="The minimum probability for a word to be considered a link candidate."
-    )
-
-def run_classification(new_threshold: float):
-    st.session_state.current_threshold = float(new_threshold)
-    st.session_state.user_input = user_input
-    if not st.session_state.user_input.strip():
-        st.warning("Please enter some text to classify.")
-        st.session_state.output_html = ""
-    else:
-        with st.spinner("Processing..."):
-            html, warning = classify_text(st.session_state.user_input, st.session_state.current_threshold)
-            if warning: st.warning(warning)
-            st.session_state.output_html = html
-    st.rerun()
-
-if st.button("Classify Text", type="primary"):
-    run_classification(slider_threshold)
-
-if st.session_state.output_html:
-    st.markdown("---")
-    st.subheader(f"Results (Threshold: {st.session_state.current_threshold:.1f}%)")
-    st.markdown(st.session_state.output_html, unsafe_allow_html=True)
-
-    col1, col2, col3 = st.columns(3)
-
-    with col1:
-        if st.button("Less", icon="➖", use_container_width=True, disabled=not st.session_state.output_html):
-            current_thr = st.session_state.current_threshold
-            if current_thr >= (100.0 - THRESHOLD_BOUNDARY_PERCENT):
-                new_threshold = current_thr + (100.0 - current_thr) / 2.0
-            else:
-                new_threshold = current_thr + THRESHOLD_STEP
-            run_classification(min(100.0, new_threshold))
-
-    with col2:
-        if st.button("Default", icon="🔄", use_container_width=True, disabled=not st.session_state.output_html):
-            run_classification(DEFAULT_THRESHOLD)
-
-    with col3:
-        if st.button("More", icon="➕", use_container_width=True, disabled=not st.session_state.output_html):
-            current_thr = st.session_state.current_threshold
-            if current_thr <= THRESHOLD_BOUNDARY_PERCENT:
-                new_threshold = current_thr / 2.0
-            else:
-                new_threshold = current_thr - THRESHOLD_STEP
-            run_classification(max(0.0, new_threshold))

src/packages.txt

@@ -1,3 +0,0 @@
-build-essential
-curl
-git

src/static/Montserrat-Italic-VariableFont_wght.ttf

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:acaff344a059669be7699d869c923bd5bb194973dc23748074f3f21deb1452dd
-size 701156

src/static/Montserrat-VariableFont_wght.ttf

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:e12288c03e4fa3721aca7ca984f25c042089dc3590e207c43a57199d7b4a5cdb
-size 688600

src/static/NotoSans-Italic-VariableFont_wdth,wght.ttf

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:4bf7b366af79c434984d67eae3967e9cd7a2f51c196101c43f21a7e21e608844
-size 2300468

src/static/NotoSans-VariableFont_wdth,wght.ttf

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:7b8cac46a1c86d2533a616b1fcf4e1926b8e39bda69034508b0df96791f56d97
-size 2044548

src/static/NotoSansMono-VariableFont_wdth,wght.ttf

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:47d856141683450ee297592b27be447eb5141c68516e5e0e748c66b6e0a54afe
-size 1707908

src/streamlit_app.py

@@ -1,328 +0,0 @@
-import streamlit as st
-import torch
-import torch.nn as nn
-# Use AutoModel and AutoTokenizer for easier loading from Hub
-from transformers import AutoModelForTokenClassification, AutoTokenizer
-from pathlib import Path
-import numpy as np
-import logging
-from dataclasses import dataclass
-from typing import Optional, Dict, List, Tuple
-
-# --- HIDE STREAMLIT MENU ---
-st.set_page_config(
-    initial_sidebar_state="collapsed"
-)
-
-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/",
-)
-
-# ----------------------------------
-# Logging
-# ----------------------------------
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-# ----------------------------------
-# Config
-# ----------------------------------
-
-@dataclass
-class AppConfig:
-    """Configuration for the LinkBERT application"""
-    # <<< CHANGE 1: Point to the Hugging Face Hub repository >>>
-    model_name: str = "dejanseo/link-prediction"
-    max_length: int = 512
-    doc_stride: int = 128
-    device: str = "cuda" if torch.cuda.is_available() else "cpu"
-
-# ----------------------------------
-# Load model/tokenizer from Hugging Face Hub
-# ----------------------------------
-
-@st.cache_resource
-def load_model_from_hub():
-    """Loads the fine-tuned model and tokenizer from the Hugging Face Hub."""
-    config = AppConfig()
-
-    logger.info(f"Loading model and tokenizer from Hugging Face Hub: {config.model_name}...")
-
-    # <<< CHANGE 2: Use Auto* classes for direct loading >>>
-    # The `AutoModelForTokenClassification` class will automatically find the correct
-    # model architecture (DeBERTaV2) and load the pre-trained weights, including the
-    # classification head.
-    model = AutoModelForTokenClassification.from_pretrained(config.model_name)
-    tokenizer = AutoTokenizer.from_pretrained(config.model_name)
-
-    model.to(config.device)
-    model.eval()
-
-    logger.info("Model and tokenizer loaded successfully.")
-    return model, tokenizer, config.device, config.max_length, config.doc_stride
-
-# <<< CHANGE 3: Call the new loading function >>>
-model, tokenizer, device, MAX_LENGTH, DOC_STRIDE = load_model_from_hub()
-
-
-# ----------------------------------
-# Inference helpers (No changes needed here, but forward pass output is slightly different)
-# ----------------------------------
-
-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
-    )
-    # The tokenizer from the Hub might not have a `word_ids` method attached by default
-    # so we create a temporary tokenization just for that.
-    temp_tokenization = tokenizer(plain_text, truncation=False)
-    full_word_ids = temp_tokenization.word_ids(batch_index=0)
-
-
-    windows_data = []
-    # Use max() to prevent step from being 0 if doc_stride is larger than cap
-    step = max(cap - doc_stride, 1)
-    start_token_idx = 0
-    total_tokens = len(full_encoding["input_ids"])
-
-    # Ensure there is at least one window
-    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]
-        # Adjust word_ids slicing to match token slicing
-        word_ids_slice = [full_word_ids[i] for i in range(len(temp_tokenization.input_ids)) if temp_tokenization.token_to_chars(i) is not None][start_token_idx:end_token_idx]
-
-        input_ids = tokenizer.build_inputs_with_special_tokens(ids_slice)
-        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 = tokenizer.build_inputs_with_special_tokens(
-            offsets_slice
-        )
-        # Pad the offset mapping to match the input_ids length
-        window_offset_mapping += [(0, 0)] * padding_length
-
-        window_word_ids = [None] + word_ids_slice + [None] * (padding_length + 1)
-
-
-        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 classify_text(
-    text: str, prediction_threshold_percent: float
-) -> Tuple[str, Optional[str]]:
-    """Classify link tokens with windowing. Returns (html, warning)."""
-    if not text.strip(): return "", "Input text is empty."
-
-    windows = windowize_inference(text, tokenizer, MAX_LENGTH, DOC_STRIDE)
-    if not windows: return "", "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)
-            }
-            # <<< CHANGE 4: The output object from AutoModel has a 'logits' attribute directly >>>
-            outputs = model(**inputs)
-            probabilities = torch.softmax(outputs.logits, dim=-1).squeeze(0)
-            # The model predicts two labels: 0 (not a link), 1 (is a link). We need the prob of label 1.
-            link_probs = probabilities[:, 1].cpu().numpy()
-
-            for i, (start, end) in enumerate(window['offset_mapping']):
-                if window['word_ids'][i] is not None and start < end:
-                    # Aggregate probabilities for characters using maximum
-                    char_link_probabilities[start:end] = np.maximum(
-                        char_link_probabilities[start:end], link_probs[i]
-                    )
-
-    final_threshold = prediction_threshold_percent / 100.0
-
-    # Tokenize once to get word_ids and offsets for the full text
-    full_encoding = tokenizer(text, return_offsets_mapping=True, truncation=False)
-    word_ids = full_encoding.word_ids(batch_index=0)
-    offsets = full_encoding['offset_mapping']
-
-    word_max_prob_map: Dict[int, float] = {}
-    word_char_spans: Dict[int, List[int]] = {}
-
-    for i, word_id in enumerate(word_ids):
-        if word_id is not None:
-            start_char, end_char = offsets[i]
-            if start_char < end_char:
-                # Get the max probability for the characters spanned by this token
-                current_token_max_prob = np.max(char_link_probabilities[start_char:end_char])
-
-                # Aggregate max probabilities for each word
-                if word_id not in word_max_prob_map:
-                    word_max_prob_map[word_id] = current_token_max_prob
-                    word_char_spans[word_id] = [start_char, end_char]
-                else:
-                    word_max_prob_map[word_id] = max(word_max_prob_map[word_id], current_token_max_prob)
-                    word_char_spans[word_id][1] = end_char # Extend end char for the word
-
-    # Determine words that meet the threshold
-    highlight_candidates: Dict[int, float] = {}
-    for word_id, max_prob in word_max_prob_map.items():
-        if max_prob >= final_threshold:
-            highlight_candidates[word_id] = max_prob
-
-    # Calculate max probability among highlighted words for normalization
-    max_highlight_prob = 0.0
-    if highlight_candidates:
-        max_highlight_prob = max(highlight_candidates.values())
-
-    # Reconstruct HTML with dynamic opacity
-    html_parts, current_char = [], 0
-    # Sort word IDs by their starting character position to reconstruct the text in order
-    sorted_word_ids = sorted(word_char_spans.keys(), key=lambda k: word_char_spans[k][0])
-
-    for word_id in sorted_word_ids:
-        start_char, end_char = word_char_spans[word_id]
-
-        # Append any text that is between the last word and this one
-        if start_char > current_char:
-            html_parts.append(text[current_char:start_char])
-
-        word_text = text[start_char:end_char]
-
-        if word_id in highlight_candidates:
-            word_prob = highlight_candidates[word_id]
-            # Normalize probability to 0.1-1.0 range for opacity
-            normalized_opacity = 1.0 # Default if no candidates or max_prob is 0
-            if max_highlight_prob > 0:
-                normalized_opacity = (word_prob / max_highlight_prob) * 0.9 + 0.1 # Scale to 0.1-1.0
-
-            base_bg_color = "#D4EDDA" # Light green
-            base_text_color = "#155724" # Dark green
-
-            html_parts.append(f"<span style='background-color: {base_bg_color}; color: {base_text_color}; "
-                              f"padding: 0.1em 0.2em; border-radius: 0.2em; opacity: {normalized_opacity:.2f};'>"
-                              f"{word_text}</span>")
-        else:
-            html_parts.append(word_text)
-        current_char = end_char
-
-    # Append any remaining text at the end
-    if current_char < len(text):
-        html_parts.append(text[current_char:])
-
-    return "".join(html_parts), None
-
-
-# ----------------------------------
-# Streamlit UI (No changes needed from here down)
-# ----------------------------------
-st.set_page_config(layout="wide", page_title="LinkBERT by DEJAN AI")
-st.title("LinkBERT")
-
-DEFAULT_THRESHOLD = 70.0
-THRESHOLD_STEP = 10.0
-THRESHOLD_BOUNDARY_PERCENT = 10.0 # Top/Bottom 10% for half-way logic
-
-# Initialize session state for threshold and output
-if 'current_threshold' not in st.session_state:
-    st.session_state.current_threshold = DEFAULT_THRESHOLD
-if 'output_html' not in st.session_state:
-    st.session_state.output_html = ""
-if 'user_input' not in st.session_state:
-    st.session_state.user_input = "DEJAN AI is the world's leading AI SEO agency. This tool showcases the capability of our latest link prediction model called LinkBERT. This model is trained on the highest quality organic link data and can predict natural link placement in plain text."
-
-# --- UI Controls ---
-user_input = st.text_area(
-    "Paste your text here:",
-    st.session_state.user_input,
-    height=200,
-    key="text_area"
-)
-
-with st.expander('Settings'):
-    slider_threshold = st.slider(
-        "Link Probability Threshold (%)",
-        min_value=0, max_value=100, value=int(st.session_state.current_threshold), step=1,
-        help="The minimum probability for a word to be considered a link candidate."
-    )
-
-# --- Classification Function (re-run logic) ---
-def run_classification(new_threshold: float):
-    st.session_state.current_threshold = float(new_threshold)
-    st.session_state.user_input = user_input # Ensure latest input is used
-    if not st.session_state.user_input.strip():
-        st.warning("Please enter some text to classify.")
-        st.session_state.output_html = ""
-    else:
-        with st.spinner("Processing..."):
-            html, warning = classify_text(st.session_state.user_input, st.session_state.current_threshold)
-            if warning: st.warning(warning)
-            st.session_state.output_html = html
-    st.rerun() # Rerun to update the display immediately
-
-
-# --- Main Classify Button ---
-if st.button("Classify Text", type="primary"):
-    run_classification(slider_threshold)
-
-# --- Display Output ---
-if st.session_state.output_html:
-    st.markdown("---")
-    st.subheader(f"Results (Threshold: {st.session_state.current_threshold:.1f}%)")
-    st.markdown(st.session_state.output_html, unsafe_allow_html=True)
-
-    # --- Adjustment Buttons ---
-    col1, col2, col3 = st.columns(3)
-
-    with col1:
-        if st.button("Less", icon=":material/playlist_remove:", use_container_width=True, disabled=not st.session_state.output_html):
-            current_thr = st.session_state.current_threshold
-            if current_thr >= (100.0 - THRESHOLD_BOUNDARY_PERCENT): # In top 10% (90-100)
-                new_threshold = current_thr + (100.0 - current_thr) / 2.0
-                new_threshold = min(100.0, new_threshold) # Ensure not more than 100
-            else:
-                new_threshold = current_thr + THRESHOLD_STEP
-                new_threshold = min(100.0 - THRESHOLD_BOUNDARY_PERCENT, new_threshold) # Don't step into deep half-way zone too soon
-            run_classification(new_threshold)
-
-
-    with col2:
-        if st.button("Default", icon=":material/notes:", use_container_width=True, disabled=not st.session_state.output_html):
-            run_classification(DEFAULT_THRESHOLD)
-
-    with col3:
-        if st.button("More", icon=":material/docs_add_on:", use_container_width=True, disabled=not st.session_state.output_html):
-            current_thr = st.session_state.current_threshold
-            if current_thr <= THRESHOLD_BOUNDARY_PERCENT: # In bottom 10% (0-10)
-                new_threshold = current_thr / 2.0
-                new_threshold = max(0.0, new_threshold) # Ensure not less than 0
-            else:
-                new_threshold = current_thr - THRESHOLD_STEP
-                new_threshold = max(THRESHOLD_BOUNDARY_PERCENT, new_threshold) # Don't step into deep half-way zone too soon
-            run_classification(new_threshold)