Upload modeling_paraformer.py with huggingface_hub

modeling_paraformer.py

@@ -0,0 +1,311 @@
+"""
+Paraformer model implementation for Hugging Face Transformers.
+
+This module implements the Paraformer model for legal document retrieval,
+based on the paper "Attentive Deep Neural Networks for Legal Document Retrieval".
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import List, Optional, Union, Tuple
+from transformers.modeling_outputs import BaseModelOutput, SequenceClassifierOutput
+from transformers.modeling_utils import PreTrainedModel
+from transformers.utils import logging
+
+try:
+    from .configuration_paraformer import ParaformerConfig
+except ImportError:
+    from configuration_paraformer import ParaformerConfig
+
+logger = logging.get_logger(__name__)
+
+
+def sparsemax(input_tensor, dim=-1):
+    """
+    Sparsemax activation function.
+    
+    Args:
+        input_tensor: Input tensor
+        dim: Dimension along which to apply sparsemax
+        
+    Returns:
+        Sparsemax output tensor
+    """
+    # Sort input in descending order
+    sorted_input, _ = torch.sort(input_tensor, dim=dim, descending=True)
+    
+    # Compute cumulative sum
+    input_cumsum = torch.cumsum(sorted_input, dim=dim) - 1
+    
+    # Create range tensor
+    k = torch.arange(1, input_tensor.size(dim) + 1, dtype=input_tensor.dtype, device=input_tensor.device)
+    if dim != -1:
+        shape = [1] * input_tensor.dim()
+        shape[dim] = -1
+        k = k.view(shape)
+    
+    # Compute support
+    support = k * sorted_input > input_cumsum
+    
+    # Find the largest k such that support[k] is True
+    support_cumsum = torch.cumsum(support.float(), dim=dim)
+    support_size = torch.sum(support.float(), dim=dim, keepdim=True)
+    
+    # Compute tau
+    tau_cumsum = torch.cumsum(sorted_input * support.float(), dim=dim)
+    tau = (tau_cumsum - 1) / support_size
+    
+    # Expand tau to match input shape
+    if dim != -1:
+        tau = tau.unsqueeze(dim)
+    
+    # Apply sparsemax
+    output = torch.clamp(input_tensor - tau, min=0)
+    
+    return output
+
+
+class ParaformerAttention(nn.Module):
+    """
+    Attention mechanism for Paraformer model.
+    
+    This implements a general attention mechanism with optional sparsemax activation.
+    """
+    
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.use_sparsemax = config.use_sparsemax
+        
+        # Attention layers
+        if config.attention_type == "general":
+            self.attention_weights = nn.Linear(config.hidden_size, 1, bias=False)
+        else:
+            raise ValueError(f"Unsupported attention type: {config.attention_type}")
+    
+    def forward(self, query_embedding, sentence_embeddings, attention_mask=None):
+        """
+        Apply attention mechanism.
+        
+        Args:
+            query_embedding: Query embedding tensor [batch_size, hidden_size]
+            sentence_embeddings: Sentence embeddings [batch_size, num_sentences, hidden_size]
+            attention_mask: Mask for padding sentences [batch_size, num_sentences]
+            
+        Returns:
+            attended_output: Weighted combination of sentence embeddings
+            attention_weights: Attention weights for interpretability
+        """
+        batch_size, num_sentences, hidden_size = sentence_embeddings.shape
+        
+        # Expand query embedding to match sentence embeddings
+        query_expanded = query_embedding.unsqueeze(1).expand(-1, num_sentences, -1)
+        
+        # Compute attention scores using general attention
+        # Combine query and sentence embeddings
+        combined = query_expanded * sentence_embeddings  # Element-wise multiplication
+        attention_scores = self.attention_weights(combined).squeeze(-1)  # [batch_size, num_sentences]
+        
+        # Apply attention mask if provided
+        if attention_mask is not None:
+            attention_scores = attention_scores.masked_fill(~attention_mask, float('-inf'))
+        
+        # Apply sparsemax or softmax
+        if self.use_sparsemax:
+            attention_weights = sparsemax(attention_scores, dim=-1)
+        else:
+            attention_weights = F.softmax(attention_scores, dim=-1)
+        
+        # Apply attention weights
+        attended_output = torch.sum(attention_weights.unsqueeze(-1) * sentence_embeddings.clone(), dim=1)
+        
+        return attended_output, attention_weights
+
+
+class ParaformerModel(PreTrainedModel):
+    """
+    Paraformer model for legal document retrieval.
+    
+    This model uses a hierarchical approach with attention mechanism to encode legal documents
+    and queries for relevance classification.
+    """
+
+    config_class = ParaformerConfig
+    base_model_prefix = "paraformer"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["ParaformerAttention"]
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.config = config
+
+        # Don't initialize SentenceTransformer in __init__ to avoid meta tensor issues
+        self._sentence_encoder = None
+        
+        # Attention mechanism
+        self.attention = ParaformerAttention(config)
+        
+        # Classifier
+        self.classifier = nn.Linear(config.hidden_size, config.num_labels)
+        self.dropout = nn.Dropout(config.dropout_prob)
+
+        # Initialize weights
+        self.post_init()
+
+    @property
+    def sentence_encoder(self):
+        """Lazy loading of SentenceTransformer to avoid meta tensor issues"""
+        if self._sentence_encoder is None:
+            from sentence_transformers import SentenceTransformer
+            self._sentence_encoder = SentenceTransformer(self.config.base_model_name)
+        return self._sentence_encoder
+
+    def forward(
+        self,
+        query_texts: Optional[List[str]] = None,
+        article_texts: Optional[List[List[str]]] = None,
+        labels: Optional[torch.Tensor] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs
+    ):
+        """
+        Forward pass of the Paraformer model.
+        
+        Args:
+            query_texts: List of query strings
+            article_texts: List of article sentence lists
+            labels: Optional labels for training
+            return_dict: Whether to return a dictionary
+            
+        Returns:
+            Model outputs including logits and optional loss
+        """
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        
+        if query_texts is None or article_texts is None:
+            raise ValueError("Both query_texts and article_texts must be provided")
+        
+        batch_size = len(query_texts)
+        device = next(self.parameters()).device
+        
+        # Encode queries
+        query_embeddings = self.sentence_encoder.encode(
+            query_texts, 
+            convert_to_tensor=True, 
+            device=device
+        ).clone()  # Clone to avoid inference tensor issues
+        
+        # Process articles
+        all_attended_outputs = []
+        all_attention_weights = []
+        
+        for i, article in enumerate(article_texts):
+            if not article:  # Handle empty articles
+                attended_output = torch.zeros(self.config.hidden_size, device=device)
+                attention_weights = torch.zeros(1, device=device)
+            else:
+                # Encode article sentences
+                sentence_embeddings = self.sentence_encoder.encode(
+                    article, 
+                    convert_to_tensor=True, 
+                    device=device
+                ).clone()  # Clone to avoid inference tensor issues
+                
+                # Add batch dimension if needed
+                if sentence_embeddings.dim() == 2:
+                    sentence_embeddings = sentence_embeddings.unsqueeze(0)
+                
+                # Apply attention
+                attended_output, attention_weights = self.attention(
+                    query_embeddings[i:i+1], 
+                    sentence_embeddings
+                )
+                attended_output = attended_output.squeeze(0)
+                attention_weights = attention_weights.squeeze(0)
+            
+            all_attended_outputs.append(attended_output)
+            all_attention_weights.append(attention_weights)
+        
+        # Stack outputs
+        attended_outputs = torch.stack(all_attended_outputs)
+        
+        # Apply dropout and classifier
+        attended_outputs = self.dropout(attended_outputs)
+        logits = self.classifier(attended_outputs)
+        
+        # Compute loss if labels provided
+        loss = None
+        if labels is not None:
+            loss_fct = nn.CrossEntropyLoss()
+            loss = loss_fct(logits.view(-1, self.config.num_labels), labels.view(-1))
+        
+        if not return_dict:
+            output = (logits,) + (all_attention_weights,)
+            return ((loss,) + output) if loss is not None else output
+        
+        return SequenceClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=None,
+            attentions=torch.stack([w.unsqueeze(0) for w in all_attention_weights]) if all_attention_weights else None,
+        )
+
+    def get_relevance_score(self, query: str, article: List[str]) -> float:
+        """
+        Get relevance score for a single query-article pair.
+        
+        Args:
+            query: Query string
+            article: List of article sentences
+            
+        Returns:
+            Relevance score between 0 and 1
+        """
+        self.eval()
+        with torch.no_grad():
+            outputs = self.forward(
+                query_texts=[query],
+                article_texts=[article],
+                return_dict=True
+            )
+            
+            probabilities = torch.softmax(outputs.logits, dim=-1)
+            relevance_score = probabilities[0, 1].item()  # Probability of being relevant
+            
+        return relevance_score
+
+    def predict_relevance(self, query: str, article: List[str]) -> int:
+        """
+        Predict binary relevance for a single query-article pair.
+        
+        Args:
+            query: Query string
+            article: List of article sentences
+            
+        Returns:
+            Binary prediction (0 = not relevant, 1 = relevant)
+        """
+        self.eval()
+        with torch.no_grad():
+            outputs = self.forward(
+                query_texts=[query],
+                article_texts=[article],
+                return_dict=True
+            )
+            
+            prediction = torch.argmax(outputs.logits, dim=-1).item()
+            
+        return prediction
+
+    def _init_weights(self, module):
+        """Initialize the weights"""
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.LayerNorm):
+            module.bias.data.zero_()
+            module.weight.data.fill_(1.0)
+