app.py

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