Upload 8 files

app.py

@@ -0,0 +1,406 @@
+import streamlit as st
+import torch
+import torch.nn as nn
+from transformers import DebertaV2Model, DebertaV2TokenizerFast, DebertaV2Config, AutoTokenizer
+from pathlib import Path
+import numpy as np
+import json
+import logging
+from dataclasses import dataclass
+from typing import Optional, Dict, List, Tuple
+from tqdm import tqdm
+from skimage.filters import threshold_otsu
+
+# ----------------------------------
+# Logging
+# ----------------------------------
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# ----------------------------------
+# Config / Model
+# ----------------------------------
+
+@dataclass
+class TrainingConfig:
+    """Training configuration for link token classification"""
+    model_name: str = "microsoft/deberta-v3-large"
+    num_labels: int = 2  # 0: not link, 1: link token
+
+    # Inference windowing
+    max_length: int = 512
+    doc_stride: int = 128  # match _prep.py for consistent windowing
+
+    # Train-only placeholders
+    train_file: str = ""
+    val_file: str = ""
+    batch_size: int = 1
+    gradient_accumulation_steps: int = 1
+    num_epochs: int = 1
+    learning_rate: float = 1e-5
+    warmup_ratio: float = 0.1
+    weight_decay: float = 0.01
+    max_grad_norm: float = 1.0
+    label_smoothing: float = 0.0
+
+    device: str = "cuda" if torch.cuda.is_available() else "cpu"
+    num_workers: int = 0
+    bf16: bool = False
+    seed: int = 42
+
+    logging_steps: int = 1
+    eval_steps: int = 100
+    save_steps: int = 100
+    output_dir: str = "./deberta_link_output"  # model is loaded from here
+
+    wandb_project: str = ""
+    wandb_name: str = ""
+
+    patience: int = 2
+    min_delta: float = 0.0001
+
+
+class DeBERTaForTokenClassification(nn.Module):
+    """DeBERTa model for token classification"""
+
+    def __init__(self, model_name: str, num_labels: int, dropout_rate: float = 0.1):
+        super().__init__()
+        self.config = DebertaV2Config.from_pretrained(model_name)
+        self.deberta = DebertaV2Model.from_pretrained(model_name)
+        self.dropout = nn.Dropout(dropout_rate)
+        self.classifier = nn.Linear(self.config.hidden_size, num_labels)
+        nn.init.xavier_uniform_(self.classifier.weight)
+        nn.init.zeros_(self.classifier.bias)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor,
+        labels: Optional[torch.Tensor] = None
+    ) -> Dict[str, torch.Tensor]:
+        outputs = self.deberta(input_ids=input_ids, attention_mask=attention_mask)
+        sequence_output = self.dropout(outputs.last_hidden_state)
+        logits = self.classifier(sequence_output)
+        return {'loss': None, 'logits': logits}
+
+# ----------------------------------
+# Load model/tokenizer (robust)
+# ----------------------------------
+
+@st.cache_resource
+def load_model():
+    """Loads pre-trained model and tokenizer. Handles raw state_dict and wrapped checkpoints."""
+    config = TrainingConfig()
+    final_dir = Path(config.output_dir) / "final_model"
+    model_path = final_dir / "pytorch_model.bin"
+
+    if not model_path.exists():
+        st.error(f"Model checkpoint not found at {model_path}.")
+        st.stop()
+
+    logger.info(f"Loading model from {model_path}...")
+    model = DeBERTaForTokenClassification(config.model_name, config.num_labels)
+
+    # Load checkpoint robustly
+    try:
+        checkpoint = torch.load(model_path, map_location=torch.device('cpu'), weights_only=False)
+    except TypeError:
+        checkpoint = torch.load(model_path, map_location=torch.device('cpu'))
+
+    # Determine state_dict
+    state_dict = None
+    if isinstance(checkpoint, dict):
+        # Case A: raw state_dict (keys -> tensors)
+        if checkpoint and all(isinstance(v, torch.Tensor) for v in checkpoint.values()):
+            state_dict = checkpoint
+            logger.info("Detected raw state_dict checkpoint.")
+        # Case B: wrapped dicts
+        elif 'model_state_dict' in checkpoint and isinstance(checkpoint['model_state_dict'], dict):
+            state_dict = checkpoint['model_state_dict']
+            logger.info("Detected 'model_state_dict' in checkpoint.")
+        elif 'state_dict' in checkpoint and isinstance(checkpoint['state_dict'], dict):
+            state_dict = checkpoint['state_dict']
+            logger.info("Detected 'state_dict' in checkpoint.")
+        else:
+            raise KeyError(f"Unrecognized checkpoint format keys: {list(checkpoint.keys())}")
+    else:
+        raise TypeError(f"Unexpected checkpoint type: {type(checkpoint)}")
+
+    missing, unexpected = model.load_state_dict(state_dict, strict=False)
+    if missing:
+        logger.warning(f"Missing keys: {missing}")
+    if unexpected:
+        logger.warning(f"Unexpected keys: {unexpected}")
+
+    model.to(config.device)
+    model.eval()
+
+    logger.info(f"Loading tokenizer {config.model_name}...")
+    tokenizer = DebertaV2TokenizerFast.from_pretrained(config.model_name)
+    logger.info("Tokenizer loaded.")
+
+    return model, tokenizer, config.device, config.max_length, config.doc_stride
+
+model, tokenizer, device, MAX_LENGTH, DOC_STRIDE = load_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
+    if cap <= 0:
+        raise ValueError(f"max_length too small; specials={specials}")
+
+    full_encoding = tokenizer(
+        plain_text,
+        add_special_tokens=False,
+        return_offsets_mapping=True,
+        return_attention_mask=False,
+        return_token_type_ids=False,
+        truncation=False,
+    )
+    input_ids_no_special = full_encoding["input_ids"]
+    offsets_no_special = full_encoding["offset_mapping"]
+
+    temp_encoding_for_word_ids = tokenizer(
+        plain_text, return_offsets_mapping=True, truncation=False, padding=False
+    )
+    full_word_ids = temp_encoding_for_word_ids.word_ids(batch_index=0)
+
+    windows_data = []
+    step = max(cap - doc_stride, 1)
+    start_token_idx = 0
+    total_tokens_no_special = len(input_ids_no_special)
+
+    while start_token_idx < total_tokens_no_special:
+        end_token_idx = min(start_token_idx + cap, total_tokens_no_special)
+
+        ids_slice_no_special = input_ids_no_special[start_token_idx:end_token_idx]
+        offsets_slice_no_special = offsets_no_special[start_token_idx:end_token_idx]
+        word_ids_slice = full_word_ids[start_token_idx:end_token_idx]
+
+        input_ids_with_special = tokenizer.build_inputs_with_special_tokens(ids_slice_no_special)
+        attention_mask_with_special = [1] * len(input_ids_with_special)
+
+        padding_length = max_length - len(input_ids_with_special)
+        if padding_length > 0:
+            input_ids_with_special.extend([tokenizer.pad_token_id] * padding_length)
+            attention_mask_with_special.extend([0] * padding_length)
+
+        window_offset_mapping = offsets_slice_no_special[:]
+        window_word_ids = word_ids_slice[:]
+
+        if tokenizer.cls_token_id is not None:
+            window_offset_mapping.insert(0, (0, 0))
+            window_word_ids.insert(0, None)
+        if tokenizer.sep_token_id is not None and len(window_offset_mapping) < max_length:
+            window_offset_mapping.append((0, 0))
+            window_word_ids.append(None)
+
+        while len(window_offset_mapping) < max_length:
+            window_offset_mapping.append((0, 0))
+            window_word_ids.append(None)
+
+        windows_data.append({
+            "input_ids": torch.tensor(input_ids_with_special, dtype=torch.long),
+            "attention_mask": torch.tensor(attention_mask_with_special, dtype=torch.long),
+            "word_ids": window_word_ids,
+            "offset_mapping": window_offset_mapping,
+        })
+
+        if end_token_idx == total_tokens_no_special:
+            break
+        start_token_idx += step
+
+    return windows_data
+
+
+def classify_text(
+    text: str,
+    otsu_mode: str,
+    prediction_threshold_override: Optional[float] = None
+) -> Tuple[str, Optional[str], Optional[float]]:
+    """Classify link tokens with windowing. Returns (html, warning, threshold%)."""
+    if not text.strip():
+        return "", None, None
+
+    windows = windowize_inference(text, tokenizer, MAX_LENGTH, DOC_STRIDE)
+    if not windows:
+        return "", "Could not generate any windows for processing.", None
+
+    char_link_probabilities = np.zeros(len(text), dtype=np.float32)
+    char_covered = np.zeros(len(text), dtype=bool)
+    all_content_token_probs = []
+
+    with torch.no_grad():
+        for window in tqdm(windows, desc="Processing windows"):
+            inputs = {
+                'input_ids': window['input_ids'].unsqueeze(0).to(device),
+                'attention_mask': window['attention_mask'].unsqueeze(0).to(device)
+            }
+            outputs = model(**inputs)
+            logits = outputs['logits'].squeeze(0)
+            probabilities = torch.softmax(logits, dim=-1)
+            link_probs_for_window_tokens = probabilities[:, 1].cpu().numpy()
+
+            for i, (offset_start, offset_end) in enumerate(window['offset_mapping']):
+                if window['word_ids'][i] is not None and offset_start < offset_end:
+                    char_link_probabilities[offset_start:offset_end] = np.maximum(
+                        char_link_probabilities[offset_start:offset_end],
+                        link_probs_for_window_tokens[i]
+                    )
+                    char_covered[offset_start:offset_end] = True
+                    all_content_token_probs.append(link_probs_for_window_tokens[i])
+
+    # Threshold selection (Otsu or manual)
+    determined_threshold_float = None
+    determined_threshold_for_display = None  # 0-100%
+
+    if prediction_threshold_override is not None:
+        determined_threshold_float = prediction_threshold_override / 100.0
+        determined_threshold_for_display = prediction_threshold_override
+    else:
+        if len(all_content_token_probs) > 1:
+            try:
+                otsu_base_threshold = threshold_otsu(np.array(all_content_token_probs))
+                conservative_delta = 0.1  # stricter
+                generous_delta = 0.1      # more lenient
+                if otsu_mode == 'conservative':
+                    determined_threshold_float = otsu_base_threshold + conservative_delta
+                elif otsu_mode == 'generous':
+                    determined_threshold_float = otsu_base_threshold - generous_delta
+                else:
+                    determined_threshold_float = otsu_base_threshold
+                determined_threshold_float = max(0.0, min(1.0, determined_threshold_float))
+                determined_threshold_for_display = determined_threshold_float * 100
+            except ValueError:
+                logger.warning("Otsu failed; defaulting to 0.5.")
+                determined_threshold_float = 0.5
+                determined_threshold_for_display = 50.0
+        else:
+            logger.warning("Insufficient tokens for Otsu; defaulting to 0.5.")
+            determined_threshold_float = 0.5
+            determined_threshold_for_display = 50.0
+
+    final_threshold = determined_threshold_float
+
+    # Word-level aggregation
+    full_text_encoding = tokenizer(text, return_offsets_mapping=True, truncation=False, padding=False)
+    full_word_ids = full_text_encoding.word_ids(batch_index=0)
+    full_offset_mapping = full_text_encoding['offset_mapping']
+
+    word_prob_map: Dict[int, List[float]] = {}
+    word_char_spans: Dict[int, List[int]] = {}
+
+    for i, word_id in enumerate(full_word_ids):
+        if word_id is not None:
+            start_char, end_char = full_offset_mapping[i]
+            if start_char < end_char and np.any(char_covered[start_char:end_char]):
+                if word_id not in word_prob_map:
+                    word_prob_map[word_id] = []
+                    word_char_spans[word_id] = [start_char, end_char]
+                else:
+                    word_char_spans[word_id][0] = min(word_char_spans[word_id][0], start_char)
+                    word_char_spans[word_id][1] = max(word_char_spans[word_id][1], end_char)
+
+                token_span_probs = char_link_probabilities[start_char:end_char]
+                word_prob_map[word_id].append(np.max(token_span_probs) if token_span_probs.size > 0 else 0.0)
+            elif word_id not in word_prob_map:
+                word_prob_map[word_id] = [0.0]
+                word_char_spans[word_id] = list(full_offset_mapping[i])
+
+    words_to_highlight_status: Dict[int, bool] = {}
+    for word_id, probs in word_prob_map.items():
+        max_word_prob = np.max(probs) if probs else 0.0
+        words_to_highlight_status[word_id] = (max_word_prob >= final_threshold)
+
+    # Reconstruct HTML with highlights
+    html_output_parts: List[str] = []
+    current_char_idx = 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_idx:
+            html_output_parts.append(text[current_char_idx:start_char])
+
+        word_text = text[start_char:end_char]
+        if words_to_highlight_status.get(word_id, False):
+            html_output_parts.append(
+                "<span style='background-color: #D4EDDA; color: #155724; padding: 0.1em 0.2em; border-radius: 0.2em;'>"
+                + word_text +
+                "</span>"
+            )
+        else:
+            html_output_parts.append(word_text)
+        current_char_idx = end_char
+
+    if current_char_idx < len(text):
+        html_output_parts.append(text[current_char_idx:])
+
+    return "".join(html_output_parts), None, determined_threshold_for_display
+
+# ----------------------------------
+# Streamlit UI
+# ----------------------------------
+
+st.set_page_config(layout="wide", page_title="LinkBERT by DEJAN AI")
+st.title("LinkBERT")
+
+user_input = st.text_area(
+    "Paste your text here:",
+    "DEJAN AI is the world's leading AI SEO agency.",
+    height=200
+)
+
+with st.expander('Settings'):
+    auto_threshold_enabled = st.checkbox(
+        "Automagic",
+        value=True,
+        help="Uncheck to set manual threshold value for link prediction."
+    )
+
+    otsu_mode_options = ['Conservative', 'Standard', 'Generous']
+    selected_otsu_mode = 'Standard'
+    if auto_threshold_enabled:
+        selected_otsu_mode = st.radio(
+            "Generosity:",
+            otsu_mode_options,
+            index=1,
+            help="Generous suggests more links; conservative suggests fewer."
+        )
+
+    prediction_threshold_manual = 50.0
+    if not auto_threshold_enabled:
+        prediction_threshold_manual = st.slider(
+            "Manual Link Probability Threshold (%)",
+            min_value=0,
+            max_value=100,
+            value=50,
+            step=1,
+            help="Minimum probability to classify a token as a link when Automagic is off."
+        )
+
+if st.button("Classify Text"):
+    if not user_input.strip():
+        st.warning("Please enter some text to classify.")
+    else:
+        threshold_to_pass = None if auto_threshold_enabled else prediction_threshold_manual
+        highlighted_html, warning_message, determined_threshold_for_display = classify_text(
+            user_input,
+            selected_otsu_mode.lower(),
+            threshold_to_pass
+        )
+        if warning_message:
+            st.warning(warning_message)
+        if determined_threshold_for_display is not None and auto_threshold_enabled:
+            st.info(f"Auto threshold: {determined_threshold_for_display:.1f}% ({selected_otsu_mode})")
+        st.markdown(highlighted_html, unsafe_allow_html=True)

config.json

@@ -0,0 +1,22 @@
+{
+		"model_type": "deberta-v2",
+		"attention_probs_dropout_prob": 0.1,
+		"hidden_act": "gelu",
+		"hidden_dropout_prob": 0.1,
+		"hidden_size": 1024,
+		"initializer_range": 0.02,
+		"intermediate_size": 4096,
+		"max_position_embeddings": 512,
+		"relative_attention": true,
+		"position_buckets": 256,
+		"norm_rel_ebd": "layer_norm",
+		"share_att_key": true,
+		"pos_att_type": "p2c|c2p",
+		"layer_norm_eps": 1e-7,
+		"max_relative_positions": -1,
+		"position_biased_input": false,
+		"num_attention_heads": 16,
+		"num_hidden_layers": 24,
+		"type_vocab_size": 0,
+		"vocab_size": 128100
+}

generator_config.json

@@ -0,0 +1,22 @@
+{
+	"model_type": "deberta-v2",
+	"attention_probs_dropout_prob": 0.1,
+	"hidden_act": "gelu",
+	"hidden_dropout_prob": 0.1,
+	"hidden_size": 1024,
+	"initializer_range": 0.02,
+	"intermediate_size": 4096,
+	"max_position_embeddings": 512,
+	"relative_attention": true,
+	"position_buckets": 256,
+	"norm_rel_ebd": "layer_norm",
+	"share_att_key": true,
+	"pos_att_type": "p2c|c2p",
+	"layer_norm_eps": 1e-7,
+	"max_relative_positions": -1,
+	"position_biased_input": false,
+	"num_attention_heads": 16,
+	"num_hidden_layers": 12,
+	"type_vocab_size": 0,
+	"vocab_size": 128100
+}

pytorch_model.bin

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

spm.model

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

tokenizer_config.json

@@ -0,0 +1,4 @@
+{
+  "do_lower_case": false,
+  "vocab_type": "spm"
+}

train.py

@@ -0,0 +1,698 @@
+import json
+import shutil
+import torch
+import torch.nn as nn
+from torch.utils.data import Dataset, DataLoader
+from transformers import (
+    DebertaV2Model,
+    DebertaV2TokenizerFast,
+    DebertaV2Config,
+    get_linear_schedule_with_warmup,
+    set_seed
+)
+from torch.cuda.amp import autocast
+from tqdm import tqdm
+import numpy as np
+from pathlib import Path
+import logging
+from dataclasses import dataclass
+from typing import Optional, Dict, List, Tuple
+import wandb
+from sklearn.metrics import accuracy_score, f1_score, precision_recall_fscore_support
+import functools  # Import functools for partial
+import re
+
+# Setup logging
+logging.basicConfig(
+    format='%(asctime)s - %(levelname)s - %(name)s - %(message)s',
+    datefmt='%m/%d/%Y %H:%M:%S',
+    level=logging.INFO
+)
+logger = logging.getLogger(__name__)
+
+@dataclass
+class TrainingConfig:
+    """Training configuration for link token classification"""
+    # Model
+    model_name: str = "microsoft/deberta-v3-large"
+    num_labels: int = 2  # 0: not link, 1: link token
+
+    # Data
+    train_file: str = "train_windows.jsonl"
+    val_file: str = "val_windows.jsonl"
+    max_length: int = 512  # This is the crucial fixed length for padding
+
+    # Training
+    batch_size: int = 8
+    gradient_accumulation_steps: int = 8
+    num_epochs: int = 3
+    learning_rate: float = 1e-6
+    warmup_ratio: float = 0.1
+    weight_decay: float = 0.01
+    max_grad_norm: float = 1.0
+    label_smoothing: float = 0.0  # Not currently used in CrossEntropyLoss
+
+    # System
+    device: str = "cuda" if torch.cuda.is_available() else "cpu"
+    num_workers: int = 0  # Set to 0 for Windows to avoid multiprocessing issues
+    seed: int = 42
+    bf16: bool = True  # Using BF16 for RTX 4090
+
+    # Logging
+    logging_steps: int = 1  # Log every step to wandb
+    eval_steps: int = 5000
+    save_steps: int = 10000
+    output_dir: str = "./deberta_link_output"
+
+    # WandB
+    wandb_project: str = "deberta-link-classification"
+    wandb_name: str = "deberta-v3-large-link-tokens"
+
+    # Early stopping
+    patience: int = 2
+    min_delta: float = 0.0001
+
+    # Checkpoint retention (Scope A: count all subdirs except 'final_model')
+    max_checkpoints: int = 5
+    protect_latest_epoch_step: bool = True  # Always keep latest best_model_epoch_* and best_model_step_*
+
+
+class LinkTokenDataset(Dataset):
+    """Dataset for link token classification"""
+
+    def __init__(self, file_path: str, max_samples: Optional[int] = None):
+        self.data = []
+
+        logger.info(f"Loading data from {file_path}")
+        seq_lengths = []
+
+        with open(file_path, 'r') as f:
+            for i, line in enumerate(f):
+                if max_samples and i >= max_samples:
+                    break
+                sample = json.loads(line)
+
+                seq_len = len(sample['input_ids'])
+                seq_lengths.append(seq_len)
+
+                # Convert to tensors
+                sample['input_ids'] = torch.tensor(sample['input_ids'], dtype=torch.long)
+                sample['attention_mask'] = torch.tensor(sample['attention_mask'], dtype=torch.long)
+                sample['labels'] = torch.tensor(sample['labels'], dtype=torch.long)
+
+                self.data.append(sample)
+
+        logger.info(f"Loaded {len(self.data)} samples")
+        logger.info(f"Sequence lengths - Min: {min(seq_lengths)}, Max: {max(seq_lengths)}, Avg: {np.mean(seq_lengths):.1f}")
+
+        # Calculate class weights for imbalanced data (for logging info)
+        total_labels = []
+        for s in self.data:
+            # Only count non-padded positions (where labels are not -100)
+            valid_labels = s['labels'][s['labels'] != -100]
+            total_labels.append(valid_labels)
+
+        # Ensure total_labels is not empty before concatenating
+        if total_labels:
+            total_labels = torch.cat(total_labels)
+            num_link_tokens = (total_labels == 1).sum().item()
+            num_non_link = (total_labels == 0).sum().item()
+
+            logger.info(f"Label distribution - Non-link: {num_non_link}, Link: {num_link_tokens}")
+            if (num_link_tokens + num_non_link) > 0:
+                logger.info(f"Link token ratio: {num_link_tokens / (num_link_tokens + num_non_link):.4%}")
+        else:
+            logger.info("No valid labels found in the dataset.")
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        return self.data[idx]
+
+
+def collate_fn(batch: List[Dict], max_seq_length: int) -> Dict[str, torch.Tensor]:
+    """
+    Custom collate function for batching with padding to a fixed max_seq_length.
+
+    Args:
+        batch (List[Dict]): A list of samples from the dataset.
+        max_seq_length (int): The maximum sequence length to pad all samples to.
+
+    Returns:
+        Dict[str, torch.Tensor]: A dictionary containing stacked and padded tensors.
+    """
+
+    input_ids = []
+    attention_mask = []
+    labels = []
+
+    for x in batch:
+        seq_len = len(x['input_ids'])
+
+        # Truncate if sequence is longer than max_seq_length (shouldn't happen with preprocessed data)
+        if seq_len > max_seq_length:
+            x['input_ids'] = x['input_ids'][:max_seq_length]
+            x['attention_mask'] = x['attention_mask'][:max_seq_length]
+            x['labels'] = x['labels'][:max_seq_length]
+            seq_len = max_seq_length
+
+        # Pad sequences to the global max_seq_length
+        padding_len = max_seq_length - seq_len
+
+        # Pad input_ids with 0 (typically the pad token id)
+        padded_input = torch.cat([
+            x['input_ids'],
+            torch.zeros(padding_len, dtype=torch.long)
+        ])
+
+        # Pad attention_mask with 0 (ignore padded tokens)
+        padded_mask = torch.cat([
+            x['attention_mask'],
+            torch.zeros(padding_len, dtype=torch.long)
+        ])
+
+        # Pad labels with -100 (ignored in loss calculation)
+        padded_labels = torch.cat([
+            x['labels'],
+            torch.full((padding_len,), -100, dtype=torch.long)
+        ])
+
+        input_ids.append(padded_input)
+        attention_mask.append(padded_mask)
+        labels.append(padded_labels)
+
+    return {
+        'input_ids': torch.stack(input_ids),
+        'attention_mask': torch.stack(attention_mask),
+        'labels': torch.stack(labels)
+    }
+
+
+class DeBERTaForTokenClassification(nn.Module):
+    """DeBERTa model for token classification"""
+
+    def __init__(self, model_name: str, num_labels: int, dropout_rate: float = 0.1):
+        super().__init__()
+
+        self.config = DebertaV2Config.from_pretrained(model_name)
+        self.deberta = DebertaV2Model.from_pretrained(model_name)
+
+        self.dropout = nn.Dropout(dropout_rate)
+        self.classifier = nn.Linear(self.config.hidden_size, num_labels)
+
+        # Initialize classifier weights
+        nn.init.xavier_uniform_(self.classifier.weight)
+        nn.init.zeros_(self.classifier.bias)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor,
+        labels: Optional[torch.Tensor] = None
+    ) -> Dict[str, torch.Tensor]:
+
+        outputs = self.deberta(
+            input_ids=input_ids,
+            attention_mask=attention_mask
+        )
+
+        sequence_output = outputs.last_hidden_state
+        sequence_output = self.dropout(sequence_output)
+        logits = self.classifier(sequence_output)
+
+        loss = None
+        if labels is not None:
+            # Calculate class weights for imbalanced dataset
+            # Link tokens are ~3.88% of data, so weight them ~25x more
+            # Ensure weight tensor is on the correct device
+            weight = torch.tensor([1.0, 25.0]).to(logits.device)
+
+            loss_fct = nn.CrossEntropyLoss(weight=weight, ignore_index=-100)
+            # Reshape logits to (batch_size * sequence_length, num_labels)
+            # Reshape labels to (batch_size * sequence_length)
+            loss = loss_fct(logits.view(-1, self.config.num_labels), labels.view(-1))
+
+        return {
+            'loss': loss,
+            'logits': logits
+        }
+
+
+def compute_metrics(predictions: np.ndarray, labels: np.ndarray, mask: np.ndarray) -> Dict[str, float]:
+    """Compute metrics for token classification"""
+    # Flatten and remove padding
+    # Only consider positions where attention_mask is 1 AND labels are not -100
+    # The -100 in labels already implies an ignored position, so we can primarily filter by that.
+
+    # Flatten all predictions, labels, and masks
+    predictions_flat = predictions.flatten()
+    labels_flat = labels.flatten()
+    mask_flat = mask.flatten()
+
+    # Create a combined filter for valid tokens (not padding, not -100 label)
+    valid_indices = (labels_flat != -100) & (mask_flat == 1)
+
+    preds_filtered = predictions_flat[valid_indices]
+    labels_filtered = labels_flat[valid_indices]
+
+    # Handle cases where no valid tokens are present
+    if len(labels_filtered) == 0:
+        return {
+            'accuracy': 0.0,
+            'precision': 0.0,
+            'recall': 0.0,
+            'f1': 0.0,
+            'f1_non_link': 0.0,
+            'f1_link': 0.0,
+            'precision_link': 0.0,
+            'recall_link': 0.0,
+            'num_valid_tokens': 0
+        }
+
+    # Calculate metrics
+    accuracy = accuracy_score(labels_filtered, preds_filtered)
+
+    precision, recall, f1, support = precision_recall_fscore_support(
+        labels_filtered, preds_filtered, average='binary', pos_label=1, zero_division=0
+    )
+
+    # Per-class metrics
+    unique_labels_in_data = np.unique(labels_filtered)
+
+    precision_per_class = [0.0, 0.0]
+    recall_per_class = [0.0, 0.0]
+    f1_per_class = [0.0, 0.0]
+
+    # Class 0 (non-link)
+    if 0 in unique_labels_in_data:
+        p0, r0, f0, _ = precision_recall_fscore_support(
+            labels_filtered, preds_filtered, labels=[0], average='binary', pos_label=0, zero_division=0
+        )
+        precision_per_class[0] = p0
+        recall_per_class[0] = r0
+        f1_per_class[0] = f0
+
+    # Class 1 (link)
+    if 1 in unique_labels_in_data:
+        p1, r1, f1_1, _ = precision_recall_fscore_support(
+            labels_filtered, preds_filtered, labels=[1], average='binary', pos_label=1, zero_division=0
+        )
+        precision_per_class[1] = p1
+        recall_per_class[1] = r1
+        f1_per_class[1] = f1_1
+
+    return {
+        'accuracy': accuracy,
+        'precision': precision,
+        'recall': recall,
+        'f1': f1,
+        'f1_non_link': f1_per_class[0],
+        'f1_link': f1_per_class[1],
+        'precision_link': precision_per_class[1],
+        'recall_link': recall_per_class[1],
+        'num_valid_tokens': len(labels_filtered)
+    }
+
+
+class Trainer:
+    """Trainer class for DeBERTa token classification"""
+
+    def __init__(self, config: TrainingConfig):
+        self.config = config
+        set_seed(config.seed)
+
+        # Initialize wandb
+        wandb.init(
+            project=config.wandb_project,
+            name=config.wandb_name,
+            config=vars(config)
+        )
+
+        # Create output directory
+        Path(config.output_dir).mkdir(parents=True, exist_ok=True)
+
+        # Load datasets
+        self.train_dataset = LinkTokenDataset(config.train_file)
+        self.val_dataset = LinkTokenDataset(config.val_file)
+
+        # Create dataloaders
+        # Use functools.partial to pass the fixed max_length to collate_fn
+        self.train_loader = DataLoader(
+            self.train_dataset,
+            batch_size=config.batch_size,
+            shuffle=False,
+            num_workers=config.num_workers,
+            collate_fn=functools.partial(collate_fn, max_seq_length=config.max_length),
+            pin_memory=True
+        )
+
+        self.val_loader = DataLoader(
+            self.val_dataset,
+            batch_size=config.batch_size * 2,  # Often larger batch size for validation
+            shuffle=False,
+            num_workers=config.num_workers,
+            collate_fn=functools.partial(collate_fn, max_seq_length=config.max_length),
+            pin_memory=True
+        )
+
+        # Initialize model
+        self.model = DeBERTaForTokenClassification(
+            config.model_name,
+            config.num_labels
+        ).to(config.device)
+
+        # Count parameters
+        total_params = sum(p.numel() for p in self.model.parameters())
+        trainable_params = sum(p.numel() for p in self.model.parameters() if p.requires_grad)
+        logger.info(f"Total parameters: {total_params:,}")
+        logger.info(f"Trainable parameters: {trainable_params:,}")
+
+        # Initialize optimizer
+        no_decay = ['bias', 'LayerNorm.weight']
+        optimizer_grouped_parameters = [
+            {
+                'params': [p for n, p in self.model.named_parameters()
+                           if not any(nd in n for nd in no_decay)],
+                'weight_decay': config.weight_decay
+            },
+            {
+                'params': [p for n, p in self.model.named_parameters()
+                           if any(nd in n for nd in no_decay)],
+                'weight_decay': 0.0
+            }
+        ]
+
+        self.optimizer = torch.optim.AdamW(
+            optimizer_grouped_parameters,
+            lr=config.learning_rate,
+            eps=1e-6
+        )
+
+        # Initialize scheduler
+        total_steps = len(self.train_loader) * config.num_epochs // config.gradient_accumulation_steps
+        warmup_steps = int(total_steps * config.warmup_ratio)
+
+        self.scheduler = get_linear_schedule_with_warmup(
+            self.optimizer,
+            num_warmup_steps=warmup_steps,
+            num_training_steps=total_steps
+        )
+
+        # Tracking variables
+        self.global_step = 0
+        self.best_val_loss = float('inf')
+        self.patience_counter = 0
+
+    def train_epoch(self, epoch: int) -> float:
+        """Train for one epoch"""
+        self.model.train()
+        total_loss = 0
+        progress_bar = tqdm(self.train_loader, desc=f"Epoch {epoch}")
+
+        # Flag to indicate if early stopping was triggered mid-epoch
+        early_stop_triggered = False
+
+        for step, batch in enumerate(progress_bar):
+            # Move batch to device
+            batch = {k: v.to(self.config.device) for k, v in batch.items()}
+
+            # Forward pass with BF16 mixed precision
+            if self.config.bf16:
+                with torch.amp.autocast(device_type='cuda', dtype=torch.bfloat16):
+                    outputs = self.model(**batch)
+                    loss = outputs['loss'] / self.config.gradient_accumulation_steps
+            else:
+                outputs = self.model(**batch)
+                loss = outputs['loss'] / self.config.gradient_accumulation_steps
+
+            # Check if loss is NaN or inf, and skip if it is
+            if torch.isnan(loss) or torch.isinf(loss):
+                logger.warning(f"NaN or Inf loss encountered at step {self.global_step}. Skipping backward pass.")
+                self.optimizer.zero_grad()  # Clear gradients for current batch
+                continue  # Skip this step
+
+            loss.backward()
+            total_loss += loss.item()
+
+            # Gradient accumulation
+            if (step + 1) % self.config.gradient_accumulation_steps == 0:
+                torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.config.max_grad_norm)
+                self.optimizer.step()
+                self.scheduler.step()
+                self.optimizer.zero_grad()
+                self.global_step += 1
+
+                # Logging - every step to wandb
+                if self.global_step % self.config.logging_steps == 0:
+                    current_loss = loss.item() * self.config.gradient_accumulation_steps
+                    wandb.log({
+                        'train/loss': current_loss,
+                        'train/learning_rate': self.scheduler.get_last_lr()[0],
+                        'train/global_step': self.global_step,
+                        'train/epoch': epoch
+                    })
+                    progress_bar.set_postfix({'loss': f'{current_loss:.4f}'})
+
+                # Evaluation
+                if self.global_step % self.config.eval_steps == 0:
+                    eval_metrics = self.evaluate()
+                    logger.info(f"Step {self.global_step} - Eval metrics: {eval_metrics}")
+
+                    # Early stopping check based on validation loss
+                    current_val_loss = eval_metrics['loss']
+                    if current_val_loss < self.best_val_loss - self.config.min_delta:
+                        self.best_val_loss = current_val_loss
+                        self.patience_counter = 0
+                        self.save_model(f"best_model_step_{self.global_step}")
+                        logger.info(f"New best validation loss: {self.best_val_loss:.4f}")
+                    else:
+                        self.patience_counter += 1
+                        logger.info(f"No improvement in validation loss. Patience: {self.patience_counter}/{self.config.patience}")
+                        if self.patience_counter >= self.config.patience:
+                            logger.info("Early stopping triggered mid-epoch!")
+                            early_stop_triggered = True
+                            break  # Break from the inner loop (current epoch)
+
+            if early_stop_triggered:
+                break  # Break from the outer loop (current epoch)
+
+        return total_loss / len(self.train_loader) if len(self.train_loader) > 0 else 0.0  # Return 0 if loader is empty
+
+    def evaluate(self) -> Dict[str, float]:
+        """Evaluate on validation set"""
+        self.model.eval()
+
+        all_predictions = []
+        all_labels = []
+        all_masks = []
+        total_loss = 0
+        num_batches = 0
+
+        with torch.no_grad():
+            for batch in tqdm(self.val_loader, desc="Evaluating"):
+                batch = {k: v.to(self.config.device) for k, v in batch.items()}
+
+                # Use BF16 for evaluation too
+                if self.config.bf16:
+                    with torch.amp.autocast(device_type='cuda', dtype=torch.bfloat16):
+                        outputs = self.model(**batch)
+                else:
+                    outputs = self.model(**batch)
+
+                if outputs['loss'] is not None:
+                    total_loss += outputs['loss'].item()
+                    num_batches += 1
+
+                predictions = torch.argmax(outputs['logits'], dim=-1)
+
+                all_predictions.append(predictions.cpu().numpy())
+                all_labels.append(batch['labels'].cpu().numpy())
+                all_masks.append(batch['attention_mask'].cpu().numpy())
+
+        all_predictions = np.concatenate(all_predictions, axis=0)
+        all_labels = np.concatenate(all_labels, axis=0)
+        all_masks = np.concatenate(all_masks, axis=0)
+
+        # Compute metrics
+        metrics = compute_metrics(all_predictions, all_labels, all_masks)
+        metrics['loss'] = total_loss / num_batches if num_batches > 0 else 0.0
+
+        # Log to wandb
+        wandb.log({f'eval/{k}': v for k, v in metrics.items()}, step=self.global_step)
+
+        self.model.train()  # Set model back to train mode after evaluation
+        return metrics
+
+    def _enforce_checkpoint_limit(self):
+        """
+        Enforce checkpoint retention:
+        - Count all subdirectories in output_dir except 'final_model'
+        - Keep at most config.max_checkpoints
+        - Delete oldest by modification time
+        - Always protect:
+            * 'final_model'
+            * latest 'best_model_epoch_*'
+            * latest 'best_model_step_*'
+        """
+        output_dir = Path(self.config.output_dir)
+        if not output_dir.exists():
+            return
+
+        # List all subdirectories
+        subdirs = [p for p in output_dir.iterdir() if p.is_dir()]
+        if not subdirs:
+            return
+
+        # Identify protected directories
+        protected = set()
+
+        # Always protect 'final_model' if present
+        final_dir = output_dir / "final_model"
+        if final_dir.exists() and final_dir.is_dir():
+            protected.add(final_dir.resolve())
+
+        if self.config.protect_latest_epoch_step:
+            # Latest best_model_epoch_*
+            epoch_dirs = [d for d in subdirs if re.match(r"best_model_epoch_\d+$", d.name)]
+            if epoch_dirs:
+                latest_epoch = max(epoch_dirs, key=lambda d: d.stat().st_mtime)
+                protected.add(latest_epoch.resolve())
+
+            # Latest best_model_step_*
+            step_dirs = [d for d in subdirs if re.match(r"best_model_step_\d+$", d.name)]
+            if step_dirs:
+                latest_step = max(step_dirs, key=lambda d: d.stat().st_mtime)
+                protected.add(latest_step.resolve())
+
+        # Candidates counted toward limit: all except 'final_model'
+        counted = [d for d in subdirs if d.resolve() != final_dir.resolve()]
+
+        # Nothing to do if within limit
+        if len(counted) <= self.config.max_checkpoints:
+            return
+
+        # Sort by mtime (oldest first)
+        counted_sorted = sorted(counted, key=lambda d: d.stat().st_mtime)
+
+        # Iteratively delete oldest non-protected until within limit
+        to_delete = []
+        current = len(counted)
+        for d in counted_sorted:
+            if current <= self.config.max_checkpoints:
+                break
+            if d.resolve() in protected:
+                continue
+            to_delete.append(d)
+            current -= 1
+
+        # If still above limit because everything old was protected,
+        # continue deleting oldest even if protected EXCEPT final_model,
+        # but try to avoid removing the most recent protected items by re-check.
+        if current > self.config.max_checkpoints:
+            # Recompute deletable set excluding final_model only
+            extras = [d for d in counted_sorted if d.resolve() != final_dir.resolve() and d not in to_delete]
+            for d in extras:
+                if current <= self.config.max_checkpoints:
+                    break
+                # Do not delete the most recent protected epoch/step if possible
+                if d.resolve() in protected:
+                    continue
+                to_delete.append(d)
+                current -= 1
+
+        # Execute deletions
+        for d in to_delete:
+            try:
+                shutil.rmtree(d)
+                logger.info(f"Deleted old checkpoint: {d}")
+            except Exception as e:
+                logger.warning(f"Failed to delete {d}: {e}")
+
+    def save_model(self, name: str):
+        """Save model checkpoint"""
+        save_path = Path(self.config.output_dir) / name
+        save_path.mkdir(parents=True, exist_ok=True)
+
+        # Only save model state dict to keep file size manageable
+        torch.save(self.model.state_dict(), save_path / 'pytorch_model.bin')
+
+        # Save config separately
+        with open(save_path / 'training_config.json', 'w') as f:
+            json.dump(vars(self.config), f, indent=4)
+
+        logger.info(f"Model saved to {save_path}")
+
+        # Enforce retention after each save
+        self._enforce_checkpoint_limit()
+
+    def train(self):
+        """Main training loop"""
+        logger.info("Starting training...")
+        logger.info(f"Training samples: {len(self.train_dataset)}")
+        logger.info(f"Validation samples: {len(self.val_dataset)}")
+
+        # Calculate total optimization steps accurately
+        total_optimization_steps = (len(self.train_loader) + self.config.gradient_accumulation_steps - 1) // self.config.gradient_accumulation_steps * self.config.num_epochs
+        logger.info(f"Total optimization steps: {total_optimization_steps}")
+        logger.info(f"Early stopping: monitoring validation loss with patience={self.config.patience}")
+
+        for epoch in range(self.config.num_epochs):
+            logger.info(f"\n{'='*50}")
+            logger.info(f"Epoch {epoch + 1}/{self.config.num_epochs}")
+
+            # Train
+            avg_train_loss = self.train_epoch(epoch + 1)
+            logger.info(f"Average training loss: {avg_train_loss:.4f}")
+
+            # Check if early stopping was already triggered mid-epoch from train_epoch
+            if self.patience_counter >= self.config.patience:
+                logger.info("Training stopped due to early stopping during epoch.")
+                break
+
+            # Evaluate at end of epoch if not already stopped
+            eval_metrics = self.evaluate()
+            logger.info(f"Epoch {epoch + 1} - Eval metrics:")
+            for key, value in eval_metrics.items():
+                logger.info(f"  {key}: {value:.4f}")
+
+            # Check for early stopping at epoch level
+            current_val_loss = eval_metrics['loss']
+            if current_val_loss < self.best_val_loss - self.config.min_delta:
+                self.best_val_loss = current_val_loss
+                self.patience_counter = 0
+                self.save_model(f"best_model_epoch_{epoch + 1}")
+                logger.info(f"New best validation loss at epoch end: {self.best_val_loss:.4f}")
+            else:
+                self.patience_counter += 1
+                logger.info(f"No improvement in validation loss. Patience: {self.patience_counter}/{self.config.patience}")
+
+            # Check for early stopping
+            if self.patience_counter >= self.config.patience:
+                logger.info("Training stopped due to early stopping")
+                break
+
+        # Save final model
+        self.save_model("final_model")
+
+        logger.info("Training completed!")
+        logger.info(f"Best validation loss: {self.best_val_loss:.4f}")
+        wandb.finish()
+
+
+def main():
+    """Main function"""
+    config = TrainingConfig()
+
+    # Optimized for RTX 4090 with BF16
+    # You can override config here based on your VRAM usage:
+    # config.batch_size = 32  # RTX 4090 can handle larger batches with 24GB VRAM
+    # config.gradient_accumulation_steps = 1  # May not need accumulation
+    # config.learning_rate = 1e-5  # Sometimes better for fine-tuning
+
+    trainer = Trainer(config)
+    trainer.train()
+
+
+if __name__ == "__main__":
+    main()

training_config.json

@@ -0,0 +1,29 @@
+{
+    "model_name": "microsoft/deberta-v3-large",
+    "num_labels": 2,
+    "train_file": "train_windows.jsonl",
+    "val_file": "val_windows.jsonl",
+    "max_length": 512,
+    "batch_size": 8,
+    "gradient_accumulation_steps": 8,
+    "num_epochs": 3,
+    "learning_rate": 1e-06,
+    "warmup_ratio": 0.1,
+    "weight_decay": 0.01,
+    "max_grad_norm": 1.0,
+    "label_smoothing": 0.0,
+    "device": "cuda",
+    "num_workers": 0,
+    "seed": 42,
+    "bf16": true,
+    "logging_steps": 1,
+    "eval_steps": 5000,
+    "save_steps": 10000,
+    "output_dir": "./deberta_link_output",
+    "wandb_project": "deberta-link-classification",
+    "wandb_name": "deberta-v3-large-link-tokens",
+    "patience": 2,
+    "min_delta": 0.0001,
+    "max_checkpoints": 5,
+    "protect_latest_epoch_step": true
+}