Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,25 +1,39 @@
 import streamlit as st
 import torch
 import torch.nn.functional as F
-from torch.nn.functional import softmax
 from transformers import AutoTokenizer, AutoModelForTokenClassification
 import pandas as pd
 import trafilatura
 
-# Set Streamlit configuration
+# Streamlit config
 st.set_page_config(layout="wide", page_title="LinkBERT")
 
-# Load model and tokenizer (correct for XLM-RoBERTa Large)
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-tokenizer = AutoTokenizer.from_pretrained("dejanseo/LinkBERT-XL")
-model = AutoModelForTokenClassification.from_pretrained("dejanseo/LinkBERT-XL").to(device)
+# Load tokenizer & model (avoid meta-tensor .to() issue)
+MODEL_ID = "dejanseo/LinkBERT-XL"
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, use_fast=True)
+
+load_kwargs = {}
+if torch.cuda.is_available():
+    # Load directly onto GPU(s); do NOT call .to(...) afterward
+    load_kwargs.update(dict(device_map="auto", torch_dtype=torch.float16))
+else:
+    # CPU load without meta tensors
+    load_kwargs.update(dict(device_map=None))
+
+model = AutoModelForTokenClassification.from_pretrained(MODEL_ID, **load_kwargs)
 model.eval()
 
 # Functions
-
 def tokenize_with_indices(text: str):
-    encoded = tokenizer.encode_plus(text, return_offsets_mapping=True, add_special_tokens=True)
-    return encoded['input_ids'], encoded['offset_mapping']
+    encoded = tokenizer.encode_plus(
+        text,
+        return_offsets_mapping=True,
+        add_special_tokens=True,
+        truncation=True,
+        max_length=512
+    )
+    return encoded["input_ids"], encoded["offset_mapping"]
 
 def fetch_and_extract_content(url: str):
     downloaded = trafilatura.fetch_url(url)
@@ -29,109 +43,102 @@ def fetch_and_extract_content(url: str):
     return None
 
 def process_text(inputs: str, confidence_threshold: float):
-    max_chunk_length = 512 - 2
+    max_chunk_length = 512 - 2  # safe window for special tokens
     words = inputs.split()
     chunk_texts = []
-    current_chunk = []
-    current_length = 0
+    current_chunk, current_length = [], 0
     for word in words:
-        if len(tokenizer.tokenize(word)) + current_length > max_chunk_length:
+        tok_len = len(tokenizer.tokenize(word))
+        if tok_len + current_length > max_chunk_length:
             chunk_texts.append(" ".join(current_chunk))
             current_chunk = [word]
-            current_length = len(tokenizer.tokenize(word))
-
+            current_length = tok_len
         else:
             current_chunk.append(word)
-            current_length += len(tokenizer.tokenize(word))
-    chunk_texts.append(" ".join(current_chunk))
-
-    df_data = {
-        'Word': [],
-        'Prediction': [],
-        'Confidence': [],
-        'Start': [],
-        'End': []
-    }
+            current_length += tok_len
+    if current_chunk:
+        chunk_texts.append(" ".join(current_chunk))
+
+    df_data = {"Word": [], "Prediction": [], "Confidence": [], "Start": [], "End": []}
     reconstructed_text = ""
     original_position_offset = 0
 
-    for chunk in chunk_texts:
-        input_ids, token_offsets = tokenize_with_indices(chunk)
-        input_ids_tensor = torch.tensor(input_ids).unsqueeze(0).to(device)
-        with torch.no_grad():
-            outputs = model(input_ids_tensor)
-            logits = outputs.logits
-            predictions = torch.argmax(logits, dim=-1).squeeze().tolist()
-            softmax_scores = F.softmax(logits, dim=-1).squeeze().tolist()
-
-        word_info = {}
-
-        for idx, (start, end) in enumerate(token_offsets):
-            if idx == 0 or idx == len(token_offsets) - 1:
-                continue
-
-            word_start = start
-            while word_start > 0 and chunk[word_start-1] != ' ':
-                word_start -= 1
-
-            if word_start not in word_info:
-                word_info[word_start] = {'prediction': 0, 'confidence': 0.0, 'subtokens': []}
-
-            confidence_percentage = softmax_scores[idx][predictions[idx]] * 100
-
-            if predictions[idx] == 1 and confidence_percentage >= confidence_threshold:
-                word_info[word_start]['prediction'] = 1
+    with torch.no_grad():
+        for chunk in chunk_texts:
+            input_ids, token_offsets = tokenize_with_indices(chunk)
+            input_ids_tensor = torch.tensor(input_ids).unsqueeze(0).to(model.device)
 
-            word_info[word_start]['confidence'] = max(word_info[word_start]['confidence'], confidence_percentage)
-            word_info[word_start]['subtokens'].append((start, end, chunk[start:end]))
-
-        last_end = 0
-        for word_start in sorted(word_info.keys()):
-            word_data = word_info[word_start]
-            for subtoken_start, subtoken_end, subtoken_text in word_data['subtokens']:
-                escaped_subtoken_text = subtoken_text.replace('$', '\\$')
-                if last_end < subtoken_start:
-                    reconstructed_text += chunk[last_end:subtoken_start]
-                if word_data['prediction'] == 1:
-                    reconstructed_text += f"<span style='background-color: rgba(0, 255, 0); display: inline;'>{escaped_subtoken_text}</span>"
-                else:
-                    reconstructed_text += escaped_subtoken_text
-                last_end = subtoken_end
-
-                df_data['Word'].append(escaped_subtoken_text)
-                df_data['Prediction'].append(word_data['prediction'])
-                df_data['Confidence'].append(word_info[word_start]['confidence'])
-                df_data['Start'].append(subtoken_start + original_position_offset)
-                df_data['End'].append(subtoken_end + original_position_offset)
-
-            original_position_offset += len(chunk) + 1
-
-        reconstructed_text += chunk[last_end:].replace('$', '\\$')
+            outputs = model(input_ids_tensor)
+            logits = outputs.logits  # [1, seq_len, num_labels]
+            predictions = torch.argmax(logits, dim=-1).squeeze(0).tolist()
+            softmax_scores = F.softmax(logits, dim=-1).squeeze(0).tolist()
+
+            word_info = {}
+            for idx, (start, end) in enumerate(token_offsets):
+                if idx == 0 or idx == len(token_offsets) - 1:
+                    continue  # skip specials
+
+                word_start = start
+                while word_start > 0 and chunk[word_start - 1] != ' ':
+                    word_start -= 1
+
+                if word_start not in word_info:
+                    word_info[word_start] = {"prediction": 0, "confidence": 0.0, "subtokens": []}
+
+                conf_pct = softmax_scores[idx][predictions[idx]] * 100.0
+                if predictions[idx] == 1 and conf_pct >= confidence_threshold:
+                    word_info[word_start]["prediction"] = 1
+                word_info[word_start]["confidence"] = max(word_info[word_start]["confidence"], conf_pct)
+                word_info[word_start]["subtokens"].append((start, end, chunk[start:end]))
+
+            last_end = 0
+            for word_start in sorted(word_info.keys()):
+                word_data = word_info[word_start]
+                for subtoken_start, subtoken_end, subtoken_text in word_data["subtokens"]:
+                    escaped = subtoken_text.replace("$", "\\$")
+                    if last_end < subtoken_start:
+                        reconstructed_text += chunk[last_end:subtoken_start]
+                    if word_data["prediction"] == 1:
+                        reconstructed_text += (
+                            f"<span style='background-color: rgba(0, 255, 0); display: inline;'>{escaped}</span>"
+                        )
+                    else:
+                        reconstructed_text += escaped
+                    last_end = subtoken_end
+
+                    df_data["Word"].append(escaped)
+                    df_data["Prediction"].append(word_data["prediction"])
+                    df_data["Confidence"].append(word_info[word_start]["confidence"])
+                    df_data["Start"].append(subtoken_start + original_position_offset)
+                    df_data["End"].append(subtoken_end + original_position_offset)
+
+                original_position_offset += len(chunk) + 1
+
+            reconstructed_text += chunk[last_end:].replace("$", "\\$")
 
     df_tokens = pd.DataFrame(df_data)
     return reconstructed_text, df_tokens
 
-# Streamlit Interface
-
-st.title('LinkBERT')
+# UI
+st.title("LinkBERT")
 st.markdown("""
-LinkBERT is a model developed by [Dejan Marketing](https://dejanmarketing.com/) designed to predict natural link placement within web content. You can either enter plain text or the URL for automated plain text extraction. To reduce the number of link predictions increase the threshold slider value.
+LinkBERT predicts natural link placement within web content. Enter text or a URL for extraction. Increase the threshold to reduce link predictions.
 """)
 
-confidence_threshold = st.slider('Confidence Threshold', 50, 100, 50)
+confidence_threshold = st.slider("Confidence Threshold", 50, 100, 50)
 
 tab1, tab2 = st.tabs(["Text Input", "URL Input"])
 
 with tab1:
     user_input = st.text_area("Enter text to process:")
-    if st.button('Process Text'):
+    if st.button("Process Text"):
         highlighted_text, df_tokens = process_text(user_input, confidence_threshold)
         st.markdown(highlighted_text, unsafe_allow_html=True)
         st.dataframe(df_tokens)
 
 with tab2:
     url_input = st.text_input("Enter URL to process:")
-    if st.button('Fetch and Process'):
+    if st.button("Fetch and Process"):
         content = fetch_and_extract_content(url_input)
         if content:
             highlighted_text, df_tokens = process_text(content, confidence_threshold)
@@ -140,28 +147,22 @@ with tab2:
         else:
             st.error("Could not fetch content from the URL. Please check the URL and try again.")
 
-# Additional information at the end
 st.divider()
 st.markdown("""
-
 ## Applications of LinkBERT
-
-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.
+- **Anchor Text Suggestion**
+- **Evaluation of Existing Links**
+- **Link Placement Guide**
+- **Anchor Text Idea Generator**
+- **Spam and Inorganic SEO Detection**
 
 ## Training and Performance
-
 LinkBERT was fine-tuned on a dataset of organic web content and editorial links.
 
 [Watch the video](https://www.youtube.com/watch?v=A0ZulyVqjZo)
-            
+
 # Engage Our Team
 Interested in using this in an automated pipeline for bulk link prediction?
 
-Please [book an appointment](https://dejanmarketing.com/conference/) to discuss your needs.
-""")
+Please [book an appointment](https://dejanmarketing.com/conference/).
+""")