sugiv/stablebridge-pruner-highlighter

┌─────────────────────────────────────────────────────────────┐
│  Input: "[CLS] query [SEP] sentence_1 sentence_2 ..."       │
│                            │                                │
│                            ▼                                │
│  ┌──────────────────────────────────────────────────────┐   │
│  │     FROZEN: BAAI/bge-reranker-v2-m3 (568M params)    │   │
│  │     Output: Token embeddings [batch, seq_len, 1024]  │   │
│  └──────────────────────────┬───────────────────────────┘   │
│                             │                               │
│                             ▼                               │
│  ┌──────────────────────────────────────────────────────┐   │
│  │     TRAINABLE: PruningHead MLP (525K params)         │   │
│  │     Linear(1024→512) → GELU → Dropout → Linear(512→1)│   │
│  └──────────────────────────┬───────────────────────────┘   │
│                             │                               │
│                             ▼                               │
│  ┌──────────────────────────────────────────────────────┐   │
│  │     AGGREGATION: Mean pooling over sentence tokens   │   │
│  │     Output: Per-sentence relevance scores [0, 1]     │   │
│  └──────────────────────────────────────────────────────┘   │
└─────────────────────────────────────────────────────────────┘

# Threshold 0.6 → 90% recall, removes ~20% noise
pruned_sentences = [s for s, score in zip(sentences, scores) if score >= 0.6]

# Threshold 0.9 → 48% recall, 27% precision
# Use only as "best guess", not authoritative
highlight_indices = [i for i, score in enumerate(scores) if score >= 0.9]

pip install "transformers>=4.41" "huggingface_hub>=0.27" torch

pip install "transformers>=4.41" "huggingface_hub>=0.27" torch

import torch
import torch.nn as nn
from transformers import AutoTokenizer, AutoModel
from huggingface_hub import hf_hub_download

class PruningHead(nn.Module):
    """Trainable classification head for sentence relevance."""
    def __init__(self, hidden_size=1024, intermediate_size=512, dropout=0.2):
        super().__init__()
        self.classifier = nn.Sequential(
            nn.Linear(hidden_size, intermediate_size),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(intermediate_size, 1),
        )
    
    def forward(self, embeddings):
        return self.classifier(embeddings).squeeze(-1)


class StableBridgePruner:
    """Unified Pruner and Highlighter for regulatory documents."""
    
    def __init__(self, device: str = 'cuda'):
        self.device = torch.device(device if torch.cuda.is_available() else 'cpu')
        
        # Load encoder (frozen)
        self.tokenizer = AutoTokenizer.from_pretrained('BAAI/bge-reranker-v2-m3')
        self.encoder = AutoModel.from_pretrained('BAAI/bge-reranker-v2-m3')
        self.encoder.to(self.device).eval()
        
        # Load pruning head
        checkpoint_path = hf_hub_download(
            repo_id="sugiv/stablebridge-pruner-highlighter",
            filename="best.pt"
        )
        self.head = PruningHead().to(self.device)
        checkpoint = torch.load(checkpoint_path, map_location=self.device, weights_only=False)
        self.head.load_state_dict(checkpoint['model_state_dict'])
        self.head.eval()
        
        # Default thresholds
        self.prune_threshold = 0.6
        self.highlight_threshold = 0.9
    
    @torch.no_grad()
    def get_sentence_scores(self, query: str, sentences: list[str]) -> list[float]:
        """
        Get relevance scores for each sentence given a query.
        
        Args:
            query: User's search query
            sentences: List of sentences from the document
        
        Returns:
            List of relevance scores in [0, 1] for each sentence
        """
        # Build document text with sentence boundaries tracked
        doc_text = ' '.join(sentences)
        text = f"{query} [SEP] {doc_text}"
        
        # Tokenize
        encoding = self.tokenizer(
            text,
            max_length=8192,
            truncation=True,
            return_tensors='pt',
            return_offsets_mapping=True
        )
        
        # Get token embeddings from encoder
        input_ids = encoding['input_ids'].to(self.device)
        attention_mask = encoding['attention_mask'].to(self.device)
        
        outputs = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
        token_embeddings = outputs.last_hidden_state
        
        # Get token-level scores
        token_logits = self.head(token_embeddings)[0]  # [seq_len]
        token_scores = torch.sigmoid(token_logits).cpu().numpy()
        
        # Aggregate to sentence scores via mean pooling
        # (Simplified: uses character offset mapping for proper aggregation)
        offsets = encoding['offset_mapping'][0].numpy()
        
        sentence_scores = []
        char_pos = len(query) + 6  # Skip "[SEP] "
        
        for sent in sentences:
            sent_start = char_pos
            sent_end = char_pos + len(sent)
            
            # Find tokens within this sentence
            token_indices = []
            for idx, (start, end) in enumerate(offsets):
                if start >= sent_start and end <= sent_end and start != end:
                    token_indices.append(idx)
            
            if token_indices:
                sent_score = float(token_scores[token_indices].mean())
            else:
                sent_score = 0.5  # Default for unmapped sentences
            
            sentence_scores.append(sent_score)
            char_pos = sent_end + 1  # +1 for space
        
        return sentence_scores
    
    def prune(self, query: str, sentences: list[str], threshold: float = None) -> list[str]:
        """
        Remove low-relevance sentences from document.
        
        Args:
            query: User's search query
            sentences: List of sentences from the document
            threshold: Minimum score to keep (default: 0.6)
        
        Returns:
            List of sentences with score >= threshold
        """
        threshold = threshold or self.prune_threshold
        scores = self.get_sentence_scores(query, sentences)
        return [s for s, score in zip(sentences, scores) if score >= threshold]
    
    def highlight(self, query: str, sentences: list[str], threshold: float = None) -> list[int]:
        """
        Get indices of high-relevance sentences for highlighting.
        
        Args:
            query: User's search query
            sentences: List of sentences from the document
            threshold: Minimum score to highlight (default: 0.9)
        
        Returns:
            List of sentence indices with score >= threshold
        """
        threshold = threshold or self.highlight_threshold
        scores = self.get_sentence_scores(query, sentences)
        return [i for i, score in enumerate(scores) if score >= threshold]
    
    def score_document(self, query: str, sentences: list[str]) -> list[tuple[str, float]]:
        """
        Get all sentences with their relevance scores.
        
        Returns:
            List of (sentence, score) tuples
        """
        scores = self.get_sentence_scores(query, sentences)
        return list(zip(sentences, scores))


# Example usage
if __name__ == "__main__":
    # Initialize model
    model = StableBridgePruner(device='cuda')
    
    # Example regulatory text
    query = "What are the licensing requirements for stablecoin issuers?"
    sentences = [
        "The GENIUS Act establishes a comprehensive framework for payment stablecoin regulation.",
        "All payment stablecoin issuers must obtain a federal license from the appropriate regulator.",
        "The legislation was introduced in the 118th Congress.",
        "Issuers must maintain reserves equal to 100% of outstanding stablecoins.",
        "The weather in Washington D.C. was cloudy during the vote.",
        "State-chartered banks may issue stablecoins under state supervision.",
    ]
    
    # Get scores
    scores = model.get_sentence_scores(query, sentences)
    print("Sentence Scores:")
    for sent, score in zip(sentences, scores):
        print(f"  [{score:.3f}] {sent[:60]}...")
    
    # Prune (remove irrelevant)
    pruned = model.prune(query, sentences)
    print(f"\nPruned: {len(sentences)} → {len(pruned)} sentences")
    
    # Highlight (mark important)
    highlights = model.highlight(query, sentences)
    print(f"Highlighted indices: {highlights}")

model:
  base_model: BAAI/bge-reranker-v2-m3
  freeze_encoder: true
  max_length: 8192
  max_sentences: 500

training:
  epochs: 10
  batch_size: 8
  gradient_accumulation_steps: 8  # Effective batch: 64
  learning_rate: 1e-4
  weight_decay: 0.02
  warmup_steps: 200
  scheduler: CosineAnnealingLR

loss:
  type: BCEWithLogitsLoss
  pos_weight: 70.0  # Compensate for 105:1 imbalance

├── best.pt              # Model checkpoint (6.3MB)
├── config.yaml          # Training configuration
├── README.md            # This model card
└── requirements.txt     # Dependencies

@misc{stablebridge-pruner-2026,
  title={StableBridge Unified Pruner and Highlighter},
  author={Sugi Venugeethan},
  year={2026},
  publisher={Hugging Face},
  url={https://huggingface.co/sugiv/stablebridge-pruner-highlighter}
}