app.py

import streamlit as st
import torch
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 / PAGE CONFIG ---
st.set_page_config(
    initial_sidebar_state="collapsed",
    layout="wide",
    page_title="LinkBERT 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
# ----------------------------------
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# ----------------------------------
# Config
# ----------------------------------

@dataclass
class AppConfig:
    """Configuration for the LinkBERT application"""
    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}...")
    logger.info(f"Using device: {config.device}")

    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

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]

        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 for special tokens
        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 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)
            }
            outputs = model(**inputs)
            probabilities = torch.softmax(outputs.logits, dim=-1).squeeze(0)
            link_probs = probabilities[:, 1].cpu().numpy()

            # --- ROBUSTNESS FIX ---
            # This loop is modified to prevent the "cannot unpack non-iterable int object" error.
            for i, offset in enumerate(window['offset_mapping']):
                # Check if the offset is a valid tuple before unpacking
                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]
                        )

    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 and i < len(offsets):
            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 start_char < len(char_link_probabilities) 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"title='Link Probability: {word_prob:.1%}'>"
                              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
# ----------------------------------
st.title("LinkBERT")

DEFAULT_THRESHOLD = 70.0
THRESHOLD_STEP = 10.0
THRESHOLD_BOUNDARY_PERCENT = 10.0  # Top/Bottom 10% for finer control

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("Analyzing text..."):
            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", use_container_width=True):
    run_classification(slider_threshold)

if st.session_state.output_html:
    st.markdown("---")
    st.markdown(st.session_state.output_html, unsafe_allow_html=True)
    st.markdown("---")

    st.markdown(
        f"<p style='text-align: center;'>Confidence Threshold: {st.session_state.current_threshold:.1f}%</p>",
        unsafe_allow_html=True
    )

    col1, col2, col3 = st.columns(3)

    with col1:
        if st.button(
            "Less",
            icon=":material/playlist_remove:",
            use_container_width=True,
            help="Show fewer, more probable links"
        ):
            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=":material/notes:",
            use_container_width=True,
            help="Reset to default threshold (70%)"
        ):
            run_classification(DEFAULT_THRESHOLD)

    with col3:
        if st.button(
            "More",
            icon=":material/docs_add_on:",
            use_container_width=True,
            help="Show more potential links"
        ):
            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))

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.
""")