Upload model_lilt.py with huggingface_hub

model_lilt.py

@@ -0,0 +1,1491 @@
+"""
+DEXTR with LiLT - 3-Step Document Extraction.
+
+Architecture:
+- Document encoder: LiLT (XLM-RoBERTa + Layout Transformer), extended to 1024 tokens
+- Query encoder: Sentence Transformer (frozen, for semantic similarity)
+- Step 1: Token Classification (Q/A/T/H/O) - TRAINED
+- Step 2: Query-Question Matching (ZERO-SHOT) - NO TRAINING
+- Step 3: Table Head (hierarchical with attention-based column assignment) - TRAINED
+
+Key insight: Query→Question matching is semantic similarity. Sentence transformers
+provide proper semantic embeddings for matching queries to question text.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import LiltModel, LiltConfig
+from sentence_transformers import SentenceTransformer
+from typing import Optional, Dict, List, Tuple
+
+
+# Label mappings for Step 1
+LABEL2ID = {
+    "O": 0,
+    "B-Q": 1, "I-Q": 2,      # Question
+    "B-A": 3, "I-A": 4,      # Answer
+    "B-H": 5, "I-H": 6,      # Header
+    "B-TABLE": 7, "I-TABLE": 8,  # Table
+}
+ID2LABEL = {v: k for k, v in LABEL2ID.items()}
+NUM_LABELS = len(LABEL2ID)
+
+
+class TokenClassificationHead(nn.Module):
+    """
+    Step 1: Token Classification Head.
+    Predicts Q/A/H/T/O labels for each token (FUNSD-style).
+    """
+
+    def __init__(self, hidden_size: int = 768, num_labels: int = NUM_LABELS, dropout: float = 0.1):
+        super().__init__()
+        self.classifier = nn.Sequential(
+            nn.Dropout(dropout),
+            nn.Linear(hidden_size, 256),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(256, num_labels),
+        )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.classifier(hidden_states)
+
+
+class QALinker(nn.Module):
+    """
+    Step 1.5: Q-A Linker.
+    Predicts which Answer span links to which Question span.
+    Uses both semantic similarity and spatial features.
+    """
+
+    def __init__(self, hidden_size: int = 768, dropout: float = 0.1):
+        super().__init__()
+        self.hidden_size = hidden_size
+
+        # Project Q and A to same space for semantic matching
+        self.q_proj = nn.Sequential(
+            nn.Linear(hidden_size, 256),
+            nn.LayerNorm(256),
+            nn.GELU(),
+            nn.Dropout(dropout),
+        )
+        self.a_proj = nn.Sequential(
+            nn.Linear(hidden_size, 256),
+            nn.LayerNorm(256),
+            nn.GELU(),
+            nn.Dropout(dropout),
+        )
+
+        # Spatial feature scorer
+        # Features: x_dist, y_dist, is_right, is_below, is_same_line, width_ratio, height_ratio
+        self.spatial_scorer = nn.Sequential(
+            nn.Linear(7, 32),
+            nn.GELU(),
+            nn.Linear(32, 1),
+        )
+
+        # Learnable temperature
+        self.log_temp = nn.Parameter(torch.tensor(1.0))
+
+    def compute_spatial_features(
+        self,
+        q_bboxes: torch.Tensor,  # (num_q, 4) - [x1, y1, x2, y2]
+        a_bboxes: torch.Tensor,  # (num_a, 4)
+    ) -> torch.Tensor:
+        """
+        Compute spatial features for all Q-A pairs.
+        Returns: (num_q, num_a, 7) feature tensor
+        """
+        num_q = q_bboxes.shape[0]
+        num_a = a_bboxes.shape[0]
+        device = q_bboxes.device
+
+        # Compute centers and sizes
+        q_cx = (q_bboxes[:, 0] + q_bboxes[:, 2]) / 2  # (num_q,)
+        q_cy = (q_bboxes[:, 1] + q_bboxes[:, 3]) / 2
+        q_w = q_bboxes[:, 2] - q_bboxes[:, 0]
+        q_h = q_bboxes[:, 3] - q_bboxes[:, 1]
+
+        a_cx = (a_bboxes[:, 0] + a_bboxes[:, 2]) / 2  # (num_a,)
+        a_cy = (a_bboxes[:, 1] + a_bboxes[:, 3]) / 2
+        a_w = a_bboxes[:, 2] - a_bboxes[:, 0]
+        a_h = a_bboxes[:, 3] - a_bboxes[:, 1]
+
+        # Expand for pairwise computation
+        q_cx = q_cx.unsqueeze(1).expand(num_q, num_a)  # (num_q, num_a)
+        q_cy = q_cy.unsqueeze(1).expand(num_q, num_a)
+        q_x2 = q_bboxes[:, 2].unsqueeze(1).expand(num_q, num_a)
+        q_y2 = q_bboxes[:, 3].unsqueeze(1).expand(num_q, num_a)
+        q_w = q_w.unsqueeze(1).expand(num_q, num_a)
+        q_h = q_h.unsqueeze(1).expand(num_q, num_a)
+
+        a_cx = a_cx.unsqueeze(0).expand(num_q, num_a)
+        a_cy = a_cy.unsqueeze(0).expand(num_q, num_a)
+        a_x1 = a_bboxes[:, 0].unsqueeze(0).expand(num_q, num_a)
+        a_y1 = a_bboxes[:, 1].unsqueeze(0).expand(num_q, num_a)
+        a_w = a_w.unsqueeze(0).expand(num_q, num_a)
+        a_h = a_h.unsqueeze(0).expand(num_q, num_a)
+
+        # Compute features (all normalized to [0, 1] range roughly)
+        x_dist = (a_x1 - q_x2) / 1000.0  # Horizontal distance (positive = A right of Q)
+        y_dist = (a_cy - q_cy).abs() / 1000.0  # Vertical distance
+        is_right = (a_x1 > q_x2).float()  # A is to the right of Q
+        is_below = (a_y1 > q_y2).float()  # A is below Q
+        is_same_line = (y_dist < 0.03).float()  # Same line (small y distance)
+        width_ratio = (a_w / (q_w + 1e-6)).clamp(0, 5) / 5.0  # Relative width
+        height_ratio = (a_h / (q_h + 1e-6)).clamp(0, 5) / 5.0  # Relative height
+
+        # Stack features: (num_q, num_a, 7)
+        features = torch.stack([
+            x_dist, y_dist, is_right, is_below, is_same_line, width_ratio, height_ratio
+        ], dim=-1)
+
+        return features
+
+    def forward(
+        self,
+        q_embeds: torch.Tensor,   # (num_q, hidden)
+        a_embeds: torch.Tensor,   # (num_a, hidden)
+        q_bboxes: torch.Tensor,   # (num_q, 4)
+        a_bboxes: torch.Tensor,   # (num_a, 4)
+    ) -> torch.Tensor:
+        """
+        Compute Q-A link scores.
+        Returns: (num_q, num_a) score matrix
+        """
+        # Semantic similarity
+        q_proj = self.q_proj(q_embeds)  # (num_q, 256)
+        a_proj = self.a_proj(a_embeds)  # (num_a, 256)
+
+        q_proj = F.normalize(q_proj, dim=-1)
+        a_proj = F.normalize(a_proj, dim=-1)
+
+        semantic_scores = torch.matmul(q_proj, a_proj.t())  # (num_q, num_a)
+
+        # Spatial scores
+        spatial_feats = self.compute_spatial_features(q_bboxes, a_bboxes)  # (num_q, num_a, 7)
+        spatial_scores = self.spatial_scorer(spatial_feats).squeeze(-1)  # (num_q, num_a)
+
+        # Combine with learnable temperature
+        temperature = self.log_temp.exp().clamp(min=0.1, max=10.0)
+        combined_scores = (semantic_scores + spatial_scores) * temperature
+
+        return combined_scores
+
+
+class QuestionPredictor(nn.Module):
+    """
+    Predicts Question embedding from Answer embedding.
+    Used when explicit Question tokens are missing in the document.
+
+    Input: LiLT A embedding (768 dim)
+    Output: Predicted Q embedding in sentence transformer space (384 dim)
+
+    Training:
+    - For Q-A pairs with Q: randomly mask Q and train to predict it
+    - Learns what Q "should look like" given the answer
+    """
+
+    def __init__(self, input_dim: int = 768, output_dim: int = 384, dropout: float = 0.1):
+        super().__init__()
+
+        # Input: LiLT answer embedding (768)
+        # Output: predicted question embedding in sentence transformer space (384)
+        self.predictor = nn.Sequential(
+            nn.Linear(input_dim, input_dim),
+            nn.LayerNorm(input_dim),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(input_dim, output_dim),
+            nn.LayerNorm(output_dim),
+        )
+
+    def forward(self, answer_embed: torch.Tensor) -> torch.Tensor:
+        """
+        Predict question embedding from answer embedding.
+        Args:
+            answer_embed: (batch, hidden) or (hidden,) - LiLT answer span embedding
+        Returns:
+            predicted_q_embed: (batch, output_dim) or (output_dim,) - in sentence transformer space
+        """
+        return self.predictor(answer_embed)
+
+
+class QueryAnswerMatcher(nn.Module):
+    """
+    Step 2: Query-Answer Matching.
+    Given answer span embeddings and query embedding, scores which span matches.
+
+    Features:
+    - Start+End pooling: spans use [h_start; h_end] instead of mean (captures boundaries)
+    - Bbox features: spatial position helps distinguish similar text
+    - Learnable temperature for score scaling
+    - Cosine similarity with projection to lower dim
+    """
+
+    def __init__(
+        self,
+        hidden_size: int = 768,
+        proj_dim: int = 256,
+        dropout: float = 0.1,
+        use_bbox_features: bool = True,
+        num_bbox_features: int = 6,  # x1, y1, x2, y2, width, height
+    ):
+        super().__init__()
+        self.proj_dim = proj_dim
+        self.use_bbox_features = use_bbox_features
+        self.hidden_size = hidden_size
+
+        # Span input: [start; end] = 2*hidden, optionally + bbox
+        span_input_dim = 2 * hidden_size
+        if use_bbox_features:
+            span_input_dim += num_bbox_features
+
+        self.span_proj = nn.Sequential(
+            nn.Linear(span_input_dim, proj_dim),
+            nn.LayerNorm(proj_dim),
+            nn.GELU(),
+            nn.Dropout(dropout),
+        )
+        self.query_proj = nn.Sequential(
+            nn.Linear(hidden_size, proj_dim),
+            nn.LayerNorm(proj_dim),
+            nn.GELU(),
+            nn.Dropout(dropout),
+        )
+        # Learnable temperature (initialized to sqrt(proj_dim) like standard attention)
+        self.log_temp = nn.Parameter(torch.tensor(proj_dim ** 0.5).log())
+
+    def forward(
+        self,
+        span_embeddings: torch.Tensor,   # (batch, num_spans, span_input_dim)
+        query_embedding: torch.Tensor,    # (batch, hidden)
+        span_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        span_proj = self.span_proj(span_embeddings)
+        query_proj = self.query_proj(query_embedding).unsqueeze(1)
+
+        # Normalize for cosine similarity-like behavior
+        span_proj = F.normalize(span_proj, dim=-1)
+        query_proj = F.normalize(query_proj, dim=-1)
+
+        # Scaled dot product with learnable temperature
+        temperature = self.log_temp.exp().clamp(min=0.1, max=100.0)
+        scores = torch.bmm(query_proj, span_proj.transpose(1, 2)).squeeze(1) * temperature
+
+        if span_mask is not None:
+            scores = scores.masked_fill(~span_mask.bool(), float('-inf'))
+        return scores
+
+
+class TableHead(nn.Module):
+    """
+    Step 3: Hierarchical Table Extraction Head.
+
+    Sub-steps:
+    1. Table detection (TABLE/O per token)
+    2. Header detection (HEADER/CELL per token)
+    3. Row segmentation (B-ROW/I-ROW/O)
+    4. Column assignment (cell → header via attention)
+    """
+
+    def __init__(self, hidden_size: int = 768, dropout: float = 0.1):
+        super().__init__()
+
+        self.table_detector = nn.Sequential(
+            nn.Dropout(dropout),
+            nn.Linear(hidden_size, 256),
+            nn.GELU(),
+            nn.Linear(256, 2),
+        )
+
+        self.header_detector = nn.Sequential(
+            nn.Dropout(dropout),
+            nn.Linear(hidden_size, 256),
+            nn.GELU(),
+            nn.Linear(256, 2),
+        )
+
+        self.row_tagger = nn.Sequential(
+            nn.Dropout(dropout),
+            nn.Linear(hidden_size, 256),
+            nn.GELU(),
+            nn.Linear(256, 3),
+        )
+
+        self.col_key_proj = nn.Linear(hidden_size, 128)
+        self.col_query_proj = nn.Linear(hidden_size, 128)
+
+    def forward(self, hidden_states: torch.Tensor) -> Dict[str, torch.Tensor]:
+        table_logits = self.table_detector(hidden_states)
+        header_logits = self.header_detector(hidden_states)
+        row_logits = self.row_tagger(hidden_states)
+
+        col_keys = self.col_key_proj(hidden_states)
+        col_queries = self.col_query_proj(hidden_states)
+        col_scores = torch.bmm(col_queries, col_keys.transpose(1, 2)) / (128 ** 0.5)
+
+        return {
+            "table_logits": table_logits,
+            "header_logits": header_logits,
+            "row_logits": row_logits,
+            "col_scores": col_scores,
+        }
+
+
+class DEXTRLiLT(nn.Module):
+    """
+    DEXTR with LiLT: Multi-Step Document Extraction with Zero-Shot Query Matching.
+
+    Architecture:
+    - Step 1: Token Classification (Q/A/H/T/O) - TRAINED
+    - Step 1.5: Q-A Linker (links Question spans to Answer spans) - TRAINED
+    - Step 2: Query → Question Matching (ZERO-SHOT with Sentence Transformer)
+    - Step 3: Table Head - TRAINED
+
+    Uses SEPARATE encoders:
+    - Document encoder: LiLT (XLM-RoBERTa + Layout Transformer) for layout-aware encoding
+    - Query encoder: Sentence Transformer (frozen) for semantic similarity
+
+    Zero-shot matching:
+    - Query text encoded with frozen Sentence Transformer
+    - Question text (from predicted Q regions) encoded with frozen Sentence Transformer
+    - Direct cosine similarity (no trainable projections)
+    - Fallback: if no Q regions, use QuestionPredictor on A text
+    """
+
+    def __init__(
+        self,
+        model_name: str = "nielsr/lilt-xlm-roberta-base",
+        query_model_name: str = "paraphrase-multilingual-MiniLM-L12-v2",
+        max_seq_len: int = 1024,
+        hidden_size: int = 768,
+        dropout: float = 0.1,
+        q_mask_prob: float = 0.3,  # Probability of masking Q during training (for Q-A linker)
+    ):
+        super().__init__()
+
+        self.max_seq_len = max_seq_len
+        self.hidden_size = hidden_size
+        self.q_mask_prob = q_mask_prob
+
+        # Document encoder: LiLT (layout-aware)
+        if max_seq_len > 512:
+            self.encoder = self._load_lilt_extended(model_name, max_seq_len)
+        else:
+            self.encoder = LiltModel.from_pretrained(model_name)
+
+        # Query encoder: Sentence Transformer (frozen, zero-shot)
+        # Provides proper semantic similarity for query→question matching
+        self.query_encoder = SentenceTransformer(query_model_name)
+        self.query_encoder.requires_grad_(False)  # Freeze for zero-shot
+        self.query_embed_dim = self.query_encoder.get_sentence_embedding_dimension()
+        print(f"Loaded frozen query encoder: {query_model_name} (dim={self.query_embed_dim})")
+
+        # Step 1: Token Classification Head
+        self.token_classifier = TokenClassificationHead(hidden_size, NUM_LABELS, dropout)
+
+        # Step 1.5: Q-A Linker (links Q spans to A spans)
+        self.qa_linker = QALinker(hidden_size, dropout)
+
+        # Step 2: QuestionPredictor for documents without explicit Q labels (e.g., receipts)
+        # Input: LiLT A embedding (768 dim), Output: Q embedding in sentence transformer space (384 dim)
+        self.question_predictor = QuestionPredictor(hidden_size, self.query_embed_dim, dropout)
+
+        # Step 3: Table Head
+        self.table_head = TableHead(hidden_size, dropout)
+
+    def _load_lilt_extended(self, model_name: str, max_seq_len: int) -> LiltModel:
+        """Load LiLT with extended position embeddings."""
+        model = LiltModel.from_pretrained(model_name)
+        original_max_pos = model.config.max_position_embeddings
+        required_positions = max_seq_len + 2
+
+        if required_positions <= original_max_pos:
+            return model
+
+        # 1. Extend text position embeddings
+        old_pos_emb = model.embeddings.position_embeddings.weight.data
+        hidden_size = old_pos_emb.shape[1]
+        padding_idx = model.embeddings.position_embeddings.padding_idx
+
+        new_pos_emb = nn.Embedding(required_positions, hidden_size, padding_idx=padding_idx)
+        new_pos_emb.weight.data[:original_max_pos] = old_pos_emb
+        new_pos_emb.weight.data[original_max_pos:].normal_(mean=0.0, std=0.02)
+
+        model.embeddings.position_embeddings = new_pos_emb
+
+        # 2. Extend layout box_position_embeddings (same sequence length limit)
+        old_box_pos = model.layout_embeddings.box_position_embeddings.weight.data
+        box_hidden = old_box_pos.shape[1]  # 192
+
+        new_box_pos = nn.Embedding(required_positions, box_hidden)
+        new_box_pos.weight.data[:original_max_pos] = old_box_pos
+        new_box_pos.weight.data[original_max_pos:].normal_(mean=0.0, std=0.02)
+
+        model.layout_embeddings.box_position_embeddings = new_box_pos
+
+        # 3. Update config
+        model.config.max_position_embeddings = required_positions
+
+        # 4. Extend position_ids buffer
+        new_position_ids = torch.arange(required_positions).unsqueeze(0)
+        model.embeddings.register_buffer("position_ids", new_position_ids, persistent=False)
+
+        print(f"Extended LiLT positions: {original_max_pos} → {required_positions}")
+        return model
+
+    def encode(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor,
+        bbox: torch.Tensor,
+    ) -> torch.Tensor:
+        """Encode tokens with layout using shared LiLT encoder."""
+        outputs = self.encoder(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            bbox=bbox,
+        )
+        return outputs.last_hidden_state
+
+    def encode_texts(
+        self,
+        texts: List[str],
+        device: torch.device,
+    ) -> torch.Tensor:
+        """
+        Encode text strings using frozen Sentence Transformer (zero-shot).
+
+        Used for encoding Q text extracted from predicted regions.
+
+        Args:
+            texts: List of text strings to encode
+            device: Device to put tensors on
+
+        Returns:
+            embeddings: (num_texts, embed_dim) tensor
+        """
+        if not texts:
+            return torch.zeros(0, self.query_embed_dim, device=device)
+
+        with torch.no_grad():
+            embeddings = self.query_encoder.encode(
+                texts,
+                convert_to_tensor=True,
+                device=device,
+            )
+            # Clone to exit inference mode (needed for autograd compatibility)
+            embeddings = embeddings.clone()
+
+        return embeddings  # (num_texts, embed_dim)
+
+    def pool_spans(
+        self,
+        hidden_states: torch.Tensor,
+        span_indices: List[List[Tuple[int, int]]],
+        bbox: Optional[torch.Tensor] = None,
+        use_bbox_features: bool = True,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Pool hidden states for answer spans using start+end tokens.
+
+        Returns [h_start; h_end; bbox_features] for each span instead of mean pooling.
+        This preserves boundary information which is crucial for extraction tasks.
+
+        Args:
+            hidden_states: (batch, seq, hidden) - document encodings
+            span_indices: List of (start, end) tuples per batch item
+            bbox: (batch, seq, 4) - bounding boxes [x1, y1, x2, y2] normalized to [0, 1000]
+            use_bbox_features: whether to include bbox in span representation
+
+        Returns:
+            span_embeddings: (batch, max_spans, span_dim) where span_dim = 2*hidden + 6 (if bbox)
+            span_mask: (batch, max_spans) bool mask
+        """
+        batch_size = hidden_states.shape[0]
+        hidden_size = hidden_states.shape[2]
+        device = hidden_states.device
+
+        max_spans = max(len(spans) for spans in span_indices) if span_indices else 1
+        max_spans = max(max_spans, 1)
+
+        # Output dimension: [h_start; h_end] + optionally [x1, y1, x2, y2, width, height]
+        span_dim = 2 * hidden_size
+        if use_bbox_features and bbox is not None:
+            span_dim += 6  # normalized bbox features
+
+        span_embeddings = torch.zeros(batch_size, max_spans, span_dim, device=device)
+        span_mask = torch.zeros(batch_size, max_spans, dtype=torch.bool, device=device)
+
+        for b, spans in enumerate(span_indices):
+            for s, (start, end) in enumerate(spans):
+                if s >= max_spans:
+                    break
+                if start >= hidden_states.shape[1] or end > hidden_states.shape[1]:
+                    continue
+
+                # Start+End pooling: [h_start; h_end]
+                h_start = hidden_states[b, start, :]
+                h_end = hidden_states[b, end - 1, :]  # end is exclusive
+                span_repr = torch.cat([h_start, h_end], dim=0)
+
+                # Add bbox features if available
+                if use_bbox_features and bbox is not None:
+                    # Get bbox of span (union of start and end token bboxes)
+                    bbox_start = bbox[b, start, :]  # [x1, y1, x2, y2]
+                    bbox_end = bbox[b, end - 1, :]
+
+                    # Compute span bounding box (min x1/y1, max x2/y2)
+                    span_x1 = torch.min(bbox_start[0], bbox_end[0])
+                    span_y1 = torch.min(bbox_start[1], bbox_end[1])
+                    span_x2 = torch.max(bbox_start[2], bbox_end[2])
+                    span_y2 = torch.max(bbox_start[3], bbox_end[3])
+
+                    # Normalize to [0, 1] and add width/height
+                    bbox_feat = torch.tensor([
+                        span_x1 / 1000.0,
+                        span_y1 / 1000.0,
+                        span_x2 / 1000.0,
+                        span_y2 / 1000.0,
+                        (span_x2 - span_x1) / 1000.0,  # width
+                        (span_y2 - span_y1) / 1000.0,  # height
+                    ], device=device)
+
+                    span_repr = torch.cat([span_repr, bbox_feat], dim=0)
+
+                span_embeddings[b, s] = span_repr
+                span_mask[b, s] = True
+
+        return span_embeddings, span_mask
+
+    def pool_single_span(
+        self,
+        hidden_states: torch.Tensor,  # (seq, hidden) - single sample
+        span: Tuple[int, int],
+        bbox: Optional[torch.Tensor] = None,  # (seq, 4)
+    ) -> torch.Tensor:
+        """
+        Pool a single span to get its embedding.
+        Returns mean of start and end tokens (hidden_size,).
+        """
+        start, end = span
+        if start >= hidden_states.shape[0] or end > hidden_states.shape[0]:
+            return torch.zeros(self.hidden_size, device=hidden_states.device)
+
+        h_start = hidden_states[start, :]
+        h_end = hidden_states[end - 1, :]
+
+        # Mean of start and end for Q embedding (simpler than concat for matching)
+        return (h_start + h_end) / 2
+
+    def get_span_bbox(
+        self,
+        bbox: torch.Tensor,  # (seq, 4)
+        span: Tuple[int, int],
+    ) -> torch.Tensor:
+        """Get bounding box for a span (union of start and end token bboxes)."""
+        start, end = span
+        bbox_start = bbox[start, :]
+        bbox_end = bbox[end - 1, :]
+
+        span_bbox = torch.stack([
+            torch.min(bbox_start[0], bbox_end[0]),  # x1
+            torch.min(bbox_start[1], bbox_end[1]),  # y1
+            torch.max(bbox_start[2], bbox_end[2]),  # x2
+            torch.max(bbox_start[3], bbox_end[3]),  # y2
+        ])
+        return span_bbox
+
+    def extract_span_text(
+        self,
+        span: Tuple[int, int],
+        tokens: List[str],
+        subword_to_word: Dict[int, int],
+    ) -> str:
+        """
+        Extract text from a span using tokens and subword-to-word mapping.
+
+        Args:
+            span: (start, end) subword indices (exclusive end)
+            tokens: List of word tokens
+            subword_to_word: Dict mapping subword idx -> word idx
+
+        Returns:
+            Text string for the span
+        """
+        start, end = span
+        # Get word indices for this span
+        word_indices = set()
+        for subword_idx in range(start, end):
+            if subword_idx in subword_to_word:
+                word_indices.add(subword_to_word[subword_idx])
+
+        if not word_indices:
+            return ""
+
+        # Get contiguous word range
+        min_word = min(word_indices)
+        max_word = max(word_indices)
+
+        # Extract and join tokens
+        span_tokens = tokens[min_word:max_word + 1]
+        return " ".join(span_tokens)
+
+    def extract_qa_regions(
+        self,
+        token_logits: torch.Tensor,  # (batch, seq, num_labels)
+        attention_mask: torch.Tensor,  # (batch, seq)
+    ) -> Tuple[List[List[Tuple[int, int]]], List[List[Tuple[int, int]]]]:
+        """
+        Extract Q and A regions from Step 1 token predictions.
+
+        This enables the cascading architecture where Step 2 uses
+        Step 1's predictions instead of ground truth spans.
+
+        Returns:
+            q_regions: List of Q spans per batch item [(start, end), ...]
+            a_regions: List of A spans per batch item [(start, end), ...]
+        """
+        batch_size = token_logits.shape[0]
+        preds = token_logits.argmax(dim=-1)  # (batch, seq)
+
+        q_regions = []
+        a_regions = []
+
+        for b in range(batch_size):
+            sample_q_regions = []
+            sample_a_regions = []
+
+            seq_len = attention_mask[b].sum().item()
+            pred_seq = preds[b, :int(seq_len)].cpu().tolist()
+
+            # Extract Q spans (B-Q=1, I-Q=2)
+            current_span = None
+            for i, label in enumerate(pred_seq):
+                if label == 1:  # B-Q
+                    if current_span is not None:
+                        sample_q_regions.append(current_span)
+                    current_span = (i, i + 1)
+                elif label == 2:  # I-Q
+                    if current_span is not None:
+                        current_span = (current_span[0], i + 1)
+                else:
+                    if current_span is not None:
+                        sample_q_regions.append(current_span)
+                        current_span = None
+            if current_span is not None:
+                sample_q_regions.append(current_span)
+
+            # Extract A spans (B-A=3, I-A=4)
+            current_span = None
+            for i, label in enumerate(pred_seq):
+                if label == 3:  # B-A
+                    if current_span is not None:
+                        sample_a_regions.append(current_span)
+                    current_span = (i, i + 1)
+                elif label == 4:  # I-A
+                    if current_span is not None:
+                        current_span = (current_span[0], i + 1)
+                else:
+                    if current_span is not None:
+                        sample_a_regions.append(current_span)
+                        current_span = None
+            if current_span is not None:
+                sample_a_regions.append(current_span)
+
+            q_regions.append(sample_q_regions)
+            a_regions.append(sample_a_regions)
+
+        return q_regions, a_regions
+
+    def match_regions_to_gt(
+        self,
+        pred_regions: List[Tuple[int, int]],
+        gt_regions: List[Tuple[int, int]],
+    ) -> Tuple[List[int], List[int]]:
+        """
+        Match predicted regions to GT regions by overlap.
+
+        Returns:
+            gt_to_pred: For each GT region, index of best matching pred region (-1 if none)
+            pred_to_gt: For each pred region, index of best matching GT region (-1 if none)
+        """
+        gt_to_pred = []
+        for gt_start, gt_end in gt_regions:
+            best_pred_idx = -1
+            best_overlap = 0
+            for pred_idx, (pred_start, pred_end) in enumerate(pred_regions):
+                overlap_start = max(gt_start, pred_start)
+                overlap_end = min(gt_end, pred_end)
+                overlap = max(0, overlap_end - overlap_start)
+                if overlap > best_overlap:
+                    best_overlap = overlap
+                    best_pred_idx = pred_idx
+            gt_to_pred.append(best_pred_idx)
+
+        pred_to_gt = []
+        for pred_start, pred_end in pred_regions:
+            best_gt_idx = -1
+            best_overlap = 0
+            for gt_idx, (gt_start, gt_end) in enumerate(gt_regions):
+                overlap_start = max(gt_start, pred_start)
+                overlap_end = min(gt_end, pred_end)
+                overlap = max(0, overlap_end - overlap_start)
+                if overlap > best_overlap:
+                    best_overlap = overlap
+                    best_gt_idx = gt_idx
+            pred_to_gt.append(best_gt_idx)
+
+        return gt_to_pred, pred_to_gt
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor,
+        bbox: torch.Tensor,
+        tokens: Optional[List[List[str]]] = None,
+        subword_to_word: Optional[List[Dict[int, int]]] = None,
+        query_texts: Optional[List[str]] = None,
+        gt_answer_spans: Optional[List[List[Tuple[int, int]]]] = None,
+        gt_question_spans: Optional[List[List[Optional[Tuple[int, int]]]]] = None,
+        target_field_idx: Optional[List[int]] = None,
+        training: bool = True,
+    ) -> Dict[str, torch.Tensor]:
+        """
+        Forward pass with CASCADING architecture and ZERO-SHOT query matching.
+
+        Key features:
+        - Step 1: Token classification to predict Q/A regions
+        - Step 1.5: Q-A linking using LiLT embeddings + spatial features
+        - Step 2: ZERO-SHOT query→question matching using Sentence Transformer
+        - GT spans are only used to compute labels (which predicted region is correct)
+
+        Args:
+            input_ids: (batch, seq) token IDs
+            attention_mask: (batch, seq) attention mask
+            bbox: (batch, seq, 4) bounding boxes
+            tokens: List of word tokens per batch item (for zero-shot Q text encoding)
+            subword_to_word: List of dicts mapping subword idx to word idx
+            query_texts: List of raw query strings (for sentence transformer)
+            gt_answer_spans: GT answer spans per batch item (for loss labels only)
+            gt_question_spans: GT question spans per batch item (for loss labels only)
+            target_field_idx: Index of GT field that query should match (for loss labels)
+            training: whether in training mode (affects Q masking for Q-A linker)
+        """
+        batch_size = input_ids.shape[0]
+        device = input_ids.device
+
+        # Encode document
+        hidden_states = self.encode(input_ids, attention_mask, bbox)
+
+        # Step 1: Token Classification
+        token_logits = self.token_classifier(hidden_states)
+
+        # Step 3: Table Head
+        table_outputs = self.table_head(hidden_states)
+
+        outputs = {
+            "hidden_states": hidden_states,
+            "token_logits": token_logits,
+            **table_outputs,
+        }
+
+        # Step 1.5 + Step 2: Q-A Linking and Query-Question Matching
+        # CASCADING: Extract Q/A regions from Step 1 PREDICTIONS
+        if query_texts is not None:
+            # Extract predicted Q and A regions from Step 1 output
+            pred_q_regions, pred_a_regions = self.extract_qa_regions(token_logits, attention_mask)
+            outputs["pred_q_regions"] = pred_q_regions
+            outputs["pred_a_regions"] = pred_a_regions
+
+            # Encode query using sentence transformer
+            query_emb = self.encode_texts(query_texts, device)
+            outputs["query_embedding"] = query_emb
+
+            # Process each sample in batch using PREDICTED regions
+            all_q_embeds_lilt = []  # Q embeddings from LiLT (768 dim) for QA Linker
+            all_q_embeds_st = []    # Q embeddings from sentence transformer (384 dim) for query matching
+            all_a_embeds = []       # A embeddings from predicted regions (768 dim)
+            all_q_bboxes = []       # Q bboxes for QA linker
+            all_a_bboxes = []       # A bboxes for QA linker
+            all_pred_q_from_a = []  # Predicted Q from A (for QuestionPredictor loss)
+            all_real_q_from_pred = []  # Real Q from predicted regions (for Q masking loss)
+            all_real_q_embeds = []  # Real Q embeds from GT (for aux loss)
+            all_gt_to_pred_a_idx = []  # Maps GT field idx to predicted A region idx
+
+            for b in range(batch_size):
+                sample_pred_q = pred_q_regions[b]  # Predicted Q spans from Step 1
+                sample_pred_a = pred_a_regions[b]  # Predicted A spans from Step 1
+
+                # Skip if no predicted regions
+                if len(sample_pred_a) == 0:
+                    all_q_embeds_lilt.append(None)
+                    all_q_embeds_st.append(None)
+                    all_a_embeds.append(None)
+                    all_q_bboxes.append(None)
+                    all_a_bboxes.append(None)
+                    all_pred_q_from_a.append(None)
+                    all_real_q_from_pred.append(None)
+                    all_real_q_embeds.append(None)
+                    all_gt_to_pred_a_idx.append(None)
+                    continue
+
+                # Pool A embeddings from predicted regions
+                sample_a_embeds = []
+                sample_a_bboxes = []
+                for a_span in sample_pred_a:
+                    a_emb = self.pool_single_span(hidden_states[b], a_span, bbox[b])
+                    sample_a_embeds.append(a_emb)
+                    sample_a_bboxes.append(self.get_span_bbox(bbox[b], a_span))
+
+                # Two Q representations:
+                # 1. q_embeds_lilt (768) - LiLT pooled, for QA Linker
+                # 2. q_embeds_st (384) - Sentence transformer, for query matching
+                sample_q_embeds_lilt = []  # For QA Linker (768 dim)
+                sample_q_embeds_st = []    # For query matching (384 dim)
+                sample_q_bboxes = []
+                sample_real_q_from_pred = []  # Real Q embeds (ST) for QuestionPredictor loss
+                sample_pred_q_from_a = []     # Predicted Q embeds from A (for loss)
+
+                if len(sample_pred_q) > 0:
+                    # We have predicted Q regions
+                    # Get Q texts for sentence transformer encoding
+                    q_texts = []
+                    if tokens is not None and subword_to_word is not None:
+                        for q_span in sample_pred_q:
+                            q_text = self.extract_span_text(q_span, tokens[b], subword_to_word[b])
+                            q_texts.append(q_text if q_text else "unknown")
+                        # Batch encode all Q texts with sentence transformer
+                        real_q_embeds_st = self.encode_texts(q_texts, device)  # (num_q, 384)
+                    else:
+                        real_q_embeds_st = None
+
+                    for i, q_span in enumerate(sample_pred_q):
+                        # LiLT embedding for QA Linker
+                        q_emb_lilt = self.pool_single_span(hidden_states[b], q_span, bbox[b])
+                        sample_q_embeds_lilt.append(q_emb_lilt)
+
+                        q_bbox = self.get_span_bbox(bbox[b], q_span)
+                        sample_q_bboxes.append(q_bbox)
+
+                        # Sentence transformer embedding for query matching
+                        if real_q_embeds_st is not None:
+                            real_q_emb_st = real_q_embeds_st[i]  # 384 dim
+                        else:
+                            real_q_emb_st = None
+
+                        # Q MASKING: during training, randomly mask Q and use predictor
+                        should_mask = training and (torch.rand(1).item() < self.q_mask_prob)
+
+                        if should_mask and i < len(sample_a_embeds):
+                            # Mask: use QuestionPredictor instead of real Q for query matching
+                            pred_q_emb = self.question_predictor(sample_a_embeds[i])  # 384 dim
+                            sample_q_embeds_st.append(pred_q_emb)
+                            sample_real_q_from_pred.append(real_q_emb_st)  # Real Q for loss
+                            sample_pred_q_from_a.append(pred_q_emb)  # Predicted Q for loss
+                        else:
+                            # No mask: use real Q (sentence transformer encoded)
+                            sample_q_embeds_st.append(real_q_emb_st if real_q_emb_st is not None else torch.zeros(self.query_embed_dim, device=device))
+                            sample_real_q_from_pred.append(real_q_emb_st)
+                            sample_pred_q_from_a.append(None)  # No prediction, no loss
+                else:
+                    # No Q regions predicted - use QuestionPredictor on all A
+                    for a_emb in sample_a_embeds:
+                        pred_q = self.question_predictor(a_emb)  # 384 dim
+                        sample_q_embeds_st.append(pred_q)
+                        sample_pred_q_from_a.append(pred_q)
+                        sample_real_q_from_pred.append(None)  # No real Q from prediction
+                    # No LiLT Q embeds when no Q predicted
+                    sample_q_embeds_lilt = []
+                    # Use A bboxes as proxy for Q bboxes
+                    sample_q_bboxes = sample_a_bboxes.copy()
+
+                # Store real Q embeds from GT (for aux loss when Q was masked)
+                # Encode GT Q text with sentence transformer (384 dim)
+                sample_real_q = []
+                if gt_question_spans is not None and b < len(gt_question_spans) and tokens is not None and subword_to_word is not None:
+                    gt_q_texts = []
+                    gt_q_valid_indices = []
+                    for idx, gt_q_span in enumerate(gt_question_spans[b]):
+                        if gt_q_span is not None:
+                            q_text = self.extract_span_text(gt_q_span, tokens[b], subword_to_word[b])
+                            gt_q_texts.append(q_text if q_text else "unknown")
+                            gt_q_valid_indices.append(idx)
+
+                    # Batch encode GT Q texts
+                    if gt_q_texts:
+                        gt_q_embeds = self.encode_texts(gt_q_texts, device)  # (num_valid, 384)
+                        embed_idx = 0
+                        for idx, gt_q_span in enumerate(gt_question_spans[b]):
+                            if gt_q_span is not None:
+                                sample_real_q.append(gt_q_embeds[embed_idx])
+                                embed_idx += 1
+                            else:
+                                sample_real_q.append(None)
+                    else:
+                        sample_real_q = [None] * len(gt_question_spans[b])
+
+                # Match GT A spans to predicted A regions (for label computation)
+                gt_to_pred_a = None
+                if gt_answer_spans is not None and b < len(gt_answer_spans):
+                    gt_a_spans = gt_answer_spans[b]
+                    gt_to_pred_a, _ = self.match_regions_to_gt(sample_pred_a, gt_a_spans)
+
+                all_q_embeds_lilt.append(torch.stack(sample_q_embeds_lilt) if sample_q_embeds_lilt else None)
+                all_q_embeds_st.append(torch.stack(sample_q_embeds_st) if sample_q_embeds_st else None)
+                all_a_embeds.append(torch.stack(sample_a_embeds) if sample_a_embeds else None)
+                all_q_bboxes.append(torch.stack(sample_q_bboxes) if sample_q_bboxes else None)
+                all_a_bboxes.append(torch.stack(sample_a_bboxes) if sample_a_bboxes else None)
+                # For Q predictor loss: parallel lists with None for non-masked positions
+                all_pred_q_from_a.append(sample_pred_q_from_a if sample_pred_q_from_a else None)
+                all_real_q_from_pred.append(sample_real_q_from_pred if sample_real_q_from_pred else None)
+                all_real_q_embeds.append(sample_real_q if sample_real_q else None)
+                all_gt_to_pred_a_idx.append(gt_to_pred_a)
+
+            # Filter out None entries and compute outputs
+            # Use LiLT Q embeds for valid check (QA Linker needs them)
+            valid_indices_lilt = [i for i, q in enumerate(all_q_embeds_lilt) if q is not None]
+            # Also track valid ST Q embeds for query matching
+            valid_indices_st = [i for i, q in enumerate(all_q_embeds_st) if q is not None]
+
+            if valid_indices_lilt:
+                # Get max spans for padding (use LiLT Q for QA Linker)
+                max_q_spans_lilt = max(all_q_embeds_lilt[i].shape[0] for i in valid_indices_lilt)
+                max_a_spans = max(all_a_embeds[i].shape[0] for i in valid_indices_lilt)
+
+                # Create padded tensors for Q (LiLT, 768 dim) - for QA Linker
+                q_embeds_lilt_padded = torch.zeros(batch_size, max_q_spans_lilt, self.hidden_size, device=device)
+                q_bboxes_padded = torch.zeros(batch_size, max_q_spans_lilt, 4, device=device)
+                q_span_mask = torch.zeros(batch_size, max_q_spans_lilt, dtype=torch.bool, device=device)
+
+                # Create padded tensors for A (LiLT, 768 dim)
+                a_embeds_padded = torch.zeros(batch_size, max_a_spans, self.hidden_size, device=device)
+                a_bboxes_padded = torch.zeros(batch_size, max_a_spans, 4, device=device)
+                a_span_mask = torch.zeros(batch_size, max_a_spans, dtype=torch.bool, device=device)
+
+                for b in valid_indices_lilt:
+                    nq = all_q_embeds_lilt[b].shape[0]
+                    q_embeds_lilt_padded[b, :nq] = all_q_embeds_lilt[b]
+                    q_bboxes_padded[b, :nq] = all_q_bboxes[b]
+                    q_span_mask[b, :nq] = True
+
+                    na = all_a_embeds[b].shape[0]
+                    a_embeds_padded[b, :na] = all_a_embeds[b]
+                    a_bboxes_padded[b, :na] = all_a_bboxes[b]
+                    a_span_mask[b, :na] = True
+
+                # Step 1.5: Q-A Linker - predict which Q links to which A (uses LiLT embeds)
+                qa_link_scores_list = []
+                for b in valid_indices_lilt:
+                    nq = all_q_embeds_lilt[b].shape[0]
+                    na = all_a_embeds[b].shape[0]
+                    if nq > 0 and na > 0:
+                        link_scores = self.qa_linker(
+                            all_q_embeds_lilt[b],  # (num_q, 768) LiLT
+                            all_a_embeds[b],       # (num_a, 768) LiLT
+                            all_q_bboxes[b],       # (num_q, 4)
+                            all_a_bboxes[b],       # (num_a, 4)
+                        )
+                        qa_link_scores_list.append(link_scores)
+                    else:
+                        qa_link_scores_list.append(None)
+
+                outputs["qa_link_scores"] = qa_link_scores_list
+
+            # Create padded tensors for Q (sentence transformer, 384 dim) - for query matching
+            if valid_indices_st:
+                max_q_spans_st = max(all_q_embeds_st[i].shape[0] for i in valid_indices_st)
+                q_embeds_st_padded = torch.zeros(batch_size, max_q_spans_st, self.query_embed_dim, device=device)
+                q_span_mask_st = torch.zeros(batch_size, max_q_spans_st, dtype=torch.bool, device=device)
+
+                for b in valid_indices_st:
+                    nq = all_q_embeds_st[b].shape[0]
+                    q_embeds_st_padded[b, :nq] = all_q_embeds_st[b]
+                    q_span_mask_st[b, :nq] = True
+
+                outputs["q_embeds_st"] = q_embeds_st_padded  # For query matching
+                outputs["q_span_mask_st"] = q_span_mask_st
+
+            # GT-based QA link scores (for training - uses GT spans directly, LiLT embeddings)
+            if training and gt_question_spans is not None and gt_answer_spans is not None:
+                gt_qa_link_scores_list = []
+                gt_valid_q_indices_list = []  # Track which Q indices are valid (not None)
+                for b in range(batch_size):
+                    gt_q_spans = gt_question_spans[b] if b < len(gt_question_spans) else []
+                    gt_a_spans = gt_answer_spans[b] if b < len(gt_answer_spans) else []
+
+                    # Filter valid Q spans (not None) and track indices
+                    valid_q_indices = [i for i, q in enumerate(gt_q_spans) if q is not None]
+                    valid_q_spans = [gt_q_spans[i] for i in valid_q_indices]
+
+                    if len(valid_q_spans) > 0 and len(gt_a_spans) > 0:
+                        # Pool embeddings from GT spans (LiLT, 768 dim)
+                        gt_q_embeds = torch.stack([
+                            self.pool_single_span(hidden_states[b], q_span, bbox[b])
+                            for q_span in valid_q_spans
+                        ])
+                        gt_a_embeds = torch.stack([
+                            self.pool_single_span(hidden_states[b], a_span, bbox[b])
+                            for a_span in gt_a_spans
+                        ])
+                        gt_q_bboxes = torch.stack([
+                            self.get_span_bbox(bbox[b], q_span)
+                            for q_span in valid_q_spans
+                        ])
+                        gt_a_bboxes = torch.stack([
+                            self.get_span_bbox(bbox[b], a_span)
+                            for a_span in gt_a_spans
+                        ])
+
+                        # Compute QA link scores on GT
+                        gt_link_scores = self.qa_linker(
+                            gt_q_embeds, gt_a_embeds, gt_q_bboxes, gt_a_bboxes
+                        )
+                        gt_qa_link_scores_list.append(gt_link_scores)
+                        gt_valid_q_indices_list.append(valid_q_indices)
+                    else:
+                        gt_qa_link_scores_list.append(None)
+                        gt_valid_q_indices_list.append(None)
+
+                outputs["gt_qa_link_scores"] = gt_qa_link_scores_list
+                outputs["gt_valid_q_indices"] = gt_valid_q_indices_list
+
+            # Step 2: Query-Question Matching (ZERO-SHOT with Sentence Transformer)
+            # Use pre-computed q_embeds_st for query matching
+            if valid_indices_st and "q_embeds_st" in outputs:
+                q_embeds_st_padded = outputs["q_embeds_st"]
+                q_span_mask_st = outputs["q_span_mask_st"]
+
+                # Compute cosine similarity between query and Q embeddings
+                query_norm = F.normalize(query_emb, dim=-1)
+                q_norm = F.normalize(q_embeds_st_padded, dim=-1)
+                match_scores = torch.bmm(q_norm, query_norm.unsqueeze(-1)).squeeze(-1)
+                match_scores = match_scores.masked_fill(~q_span_mask_st, float('-inf'))
+                outputs["match_scores"] = match_scores
+                outputs["match_span_mask"] = q_span_mask_st
+
+            # Store LiLT embeddings for QA Linker
+            if valid_indices_lilt:
+                outputs["q_span_mask"] = q_span_mask
+                outputs["a_span_mask"] = a_span_mask
+                outputs["q_embeds"] = q_embeds_lilt_padded
+                outputs["a_embeds"] = a_embeds_padded
+                outputs["q_bboxes"] = q_bboxes_padded
+                outputs["a_bboxes"] = a_bboxes_padded
+
+            # Store for loss computation (Q masking / QuestionPredictor)
+            outputs["pred_q_from_a"] = all_pred_q_from_a
+            outputs["real_q_from_pred"] = all_real_q_from_pred  # Real Q (ST) when masked
+            outputs["real_q_embeds"] = all_real_q_embeds  # GT Q embeds (ST)
+            outputs["gt_to_pred_a_idx"] = all_gt_to_pred_a_idx
+            outputs["valid_batch_indices_lilt"] = valid_indices_lilt
+            outputs["valid_batch_indices_st"] = valid_indices_st
+
+        return outputs
+
+
+def compute_loss(
+    outputs: Dict[str, torch.Tensor],
+    token_labels: torch.Tensor,
+    attention_mask: torch.Tensor,
+    table_labels: Optional[torch.Tensor] = None,
+    row_labels: Optional[torch.Tensor] = None,
+    header_labels: Optional[torch.Tensor] = None,
+    match_labels: Optional[torch.Tensor] = None,
+    col_labels: Optional[torch.Tensor] = None,
+    class_weights: Optional[torch.Tensor] = None,
+    qa_link_labels: Optional[List[torch.Tensor]] = None,  # NEW: Q-A link labels
+) -> Dict[str, torch.Tensor]:
+    """Compute joint loss for all steps including Q-A linking."""
+    device = token_labels.device
+    losses = {}
+
+    # Step 1: Token Classification
+    token_logits = outputs["token_logits"]
+    token_logits_flat = token_logits.view(-1, token_logits.shape[-1])
+    token_labels_flat = token_labels.view(-1)
+    attn_flat = attention_mask.view(-1).bool()
+
+    if class_weights is not None:
+        step1_loss = F.cross_entropy(
+            token_logits_flat[attn_flat],
+            token_labels_flat[attn_flat],
+            weight=class_weights,
+        )
+    else:
+        step1_loss = F.cross_entropy(
+            token_logits_flat[attn_flat],
+            token_labels_flat[attn_flat],
+        )
+    losses["step1_loss"] = step1_loss
+
+    # Step 1.5: Q-A Linker Loss (uses GT spans directly - no dependency on Step 1 predictions)
+    if "gt_qa_link_scores" in outputs and qa_link_labels is not None:
+        qa_link_losses = []
+        gt_scores_list = outputs["gt_qa_link_scores"]
+        gt_valid_q_indices = outputs.get("gt_valid_q_indices", [None] * len(gt_scores_list))
+        for scores, labels, valid_indices in zip(gt_scores_list, qa_link_labels, gt_valid_q_indices):
+            if scores is None or labels is None or valid_indices is None:
+                continue
+            # Filter labels to only include valid Q indices (matching scores shape)
+            filtered_labels = labels[valid_indices] if len(valid_indices) > 0 else labels
+            if scores.numel() > 0 and filtered_labels.numel() > 0:
+                # scores: (num_valid_q, num_gt_a), filtered_labels: (num_valid_q,)
+                num_a = scores.shape[1]
+                valid_mask = (filtered_labels >= 0) & (filtered_labels < num_a)
+                if valid_mask.any():
+                    link_loss = F.cross_entropy(scores[valid_mask], filtered_labels[valid_mask])
+                    qa_link_losses.append(link_loss)
+        if qa_link_losses:
+            losses["qa_link_loss"] = torch.stack(qa_link_losses).mean()
+        else:
+            losses["qa_link_loss"] = torch.tensor(0.0, device=device)
+    else:
+        losses["qa_link_loss"] = torch.tensor(0.0, device=device)
+
+    # Question Predictor Loss (MSE between predicted Q and real Q in Sentence Transformer space)
+    # QuestionPredictor learns: A_lilt_embed (768) → Q_st_embed (384)
+    # Training: when Q is masked, compare predicted Q with real Q (sentence transformer)
+    if "pred_q_from_a" in outputs and "real_q_from_pred" in outputs:
+        pred_q_from_a = outputs["pred_q_from_a"]  # List of lists
+        real_q_from_pred = outputs["real_q_from_pred"]  # List of lists
+
+        q_predict_losses = []
+        for batch_pred, batch_real in zip(pred_q_from_a, real_q_from_pred):
+            if batch_pred is None or batch_real is None:
+                continue
+            for pred_q, real_q in zip(batch_pred, batch_real):
+                if pred_q is not None and real_q is not None:
+                    # MSE loss to make predicted Q similar to real Q
+                    mse = F.mse_loss(pred_q, real_q.detach())
+                    q_predict_losses.append(mse)
+
+        if q_predict_losses:
+            losses["q_predict_loss"] = torch.stack(q_predict_losses).mean()
+        else:
+            losses["q_predict_loss"] = torch.tensor(0.0, device=device)
+    else:
+        losses["q_predict_loss"] = torch.tensor(0.0, device=device)
+
+    # Step 2: Query-Question Matching - ZERO-SHOT (no loss)
+    # Zero-shot Q text matching - no training needed
+    # QuestionPredictor IS trained (for A→Q prediction when no Q regions)
+    losses["step2_loss"] = torch.tensor(0.0, device=device)
+
+    # Step 3: Table Losses
+    if table_labels is not None:
+        table_logits = outputs["table_logits"]
+        table_logits_flat = table_logits.view(-1, 2)
+        table_labels_flat = table_labels.view(-1)
+        table_det_loss = F.cross_entropy(
+            table_logits_flat[attn_flat],
+            table_labels_flat[attn_flat],
+        )
+        losses["table_det_loss"] = table_det_loss
+    else:
+        losses["table_det_loss"] = torch.tensor(0.0, device=device)
+
+    if header_labels is not None and table_labels is not None:
+        header_logits = outputs["header_logits"]
+        table_mask = (table_labels == 1) & attention_mask.bool()
+        if table_mask.any():
+            header_logits_flat = header_logits.view(-1, 2)
+            header_labels_flat = header_labels.view(-1)
+            table_mask_flat = table_mask.view(-1)
+            header_loss = F.cross_entropy(
+                header_logits_flat[table_mask_flat],
+                header_labels_flat[table_mask_flat],
+            )
+            losses["header_loss"] = header_loss
+        else:
+            losses["header_loss"] = torch.tensor(0.0, device=device)
+    else:
+        losses["header_loss"] = torch.tensor(0.0, device=device)
+
+    if row_labels is not None and table_labels is not None:
+        row_logits = outputs["row_logits"]
+        table_mask = (table_labels == 1) & attention_mask.bool()
+        if table_mask.any():
+            row_logits_flat = row_logits.view(-1, 3)
+            row_labels_flat = row_labels.view(-1)
+            table_mask_flat = table_mask.view(-1)
+            row_loss = F.cross_entropy(
+                row_logits_flat[table_mask_flat],
+                row_labels_flat[table_mask_flat],
+            )
+            losses["row_loss"] = row_loss
+        else:
+            losses["row_loss"] = torch.tensor(0.0, device=device)
+    else:
+        losses["row_loss"] = torch.tensor(0.0, device=device)
+
+    if col_labels is not None and table_labels is not None and header_labels is not None:
+        col_scores = outputs["col_scores"]
+        cell_mask = (table_labels == 1) & (header_labels == 0) & attention_mask.bool()
+        if cell_mask.any():
+            col_scores_flat = col_scores.view(-1, col_scores.shape[-1])
+            col_labels_flat = col_labels.view(-1)
+            cell_mask_flat = cell_mask.view(-1)
+            col_loss = F.cross_entropy(
+                col_scores_flat[cell_mask_flat],
+                col_labels_flat[cell_mask_flat],
+            )
+            losses["col_loss"] = col_loss
+        else:
+            losses["col_loss"] = torch.tensor(0.0, device=device)
+    else:
+        losses["col_loss"] = torch.tensor(0.0, device=device)
+
+    # Total loss with weighted components
+    total_loss = (
+        losses["step1_loss"] +              # Token classification (main task)
+        losses["qa_link_loss"] * 0.3 +      # Q-A linker (reduced to not compete with step1)
+        losses["q_predict_loss"] * 0.2 +    # Q prediction aux (reduced)
+        losses["step2_loss"] * 0.5 +        # Query-Q matching (reduced)
+        losses["table_det_loss"] * 0.5 +    # Table detection
+        losses["header_loss"] * 0.3 +       # Header detection
+        losses["row_loss"] * 0.3 +          # Row segmentation
+        losses["col_loss"] * 0.3            # Column assignment
+    )
+    losses["loss"] = total_loss
+
+    return losses
+
+
+def compute_contrastive_loss(
+    query_embedding: torch.Tensor,  # (batch, hidden)
+    span_embeddings: torch.Tensor,  # (batch, num_spans, span_dim)
+    span_mask: torch.Tensor,        # (batch, num_spans)
+    match_labels: torch.Tensor,     # (batch,) - index of correct span
+    temperature: float = 0.07,
+    hard_negative_weight: float = 2.0,
+) -> torch.Tensor:
+    """
+    Contrastive loss for query-answer matching.
+
+    InfoNCE-style loss that:
+    - Pulls query towards correct answer span
+    - Pushes query away from all other spans (in-batch negatives)
+    - Applies extra weight to hard negatives (high-scoring wrong answers)
+
+    Args:
+        query_embedding: Query [CLS] embeddings
+        span_embeddings: Pooled span embeddings [h_start; h_end; bbox]
+        span_mask: Valid span mask
+        match_labels: Index of correct span for each query
+        temperature: Softmax temperature (lower = harder)
+        hard_negative_weight: Extra weight for hard negatives
+
+    Returns:
+        Contrastive loss scalar
+    """
+    batch_size = query_embedding.shape[0]
+    device = query_embedding.device
+
+    if batch_size == 0:
+        return torch.tensor(0.0, device=device)
+
+    # Project spans to same dimension as query if needed
+    # (This is already done by QueryAnswerMatcher, so we use the projected scores)
+    # Here we compute a simpler version using the match_scores from forward
+
+    # For now, use a margin-based contrastive approach
+    # We want: sim(q, correct_span) > sim(q, wrong_span) + margin
+
+    losses = []
+    for b in range(batch_size):
+        valid_spans = span_mask[b].sum().item()
+        if valid_spans <= 1:
+            continue
+
+        correct_idx = match_labels[b].item()
+        if correct_idx >= valid_spans:
+            continue
+
+        # Get span embeddings for this sample
+        spans = span_embeddings[b, :int(valid_spans), :]  # (num_valid, dim)
+        query = query_embedding[b]  # (hidden,)
+
+        # Simple approach: normalize and compute similarities
+        # The projection happens in QueryAnswerMatcher, so we compute raw similarities here
+        # This loss is auxiliary to the CE loss
+
+        # Normalize for cosine similarity
+        spans_norm = F.normalize(spans[:, :query.shape[0]], dim=-1)  # Use first hidden_size dims
+        query_norm = F.normalize(query, dim=-1)
+
+        # Compute similarities
+        sims = torch.matmul(spans_norm, query_norm) / temperature  # (num_valid,)
+
+        # Create target (one-hot)
+        target = torch.zeros(int(valid_spans), device=device)
+        target[correct_idx] = 1.0
+
+        # InfoNCE: -log(exp(sim_pos) / sum(exp(sim_all)))
+        # With hard negative weighting
+        weights = torch.ones(int(valid_spans), device=device)
+        weights[correct_idx] = 1.0
+
+        # Hard negatives: spans with high similarity but wrong
+        with torch.no_grad():
+            hard_neg_mask = (sims > sims[correct_idx] - 0.5) & (torch.arange(int(valid_spans), device=device) != correct_idx)
+            weights[hard_neg_mask] = hard_negative_weight
+
+        # Weighted softmax cross-entropy
+        log_probs = F.log_softmax(sims, dim=0)
+        loss = -log_probs[correct_idx]
+
+        # Add margin loss for hard negatives
+        for i in range(int(valid_spans)):
+            if i != correct_idx and hard_neg_mask[i] if hard_neg_mask.any() else False:
+                margin_loss = F.relu(sims[i] - sims[correct_idx] + 0.3)
+                loss = loss + 0.1 * margin_loss
+
+        losses.append(loss)
+
+    if not losses:
+        return torch.tensor(0.0, device=device)
+
+    return torch.stack(losses).mean()
+
+
+def compute_loss_with_contrastive(
+    outputs: Dict[str, torch.Tensor],
+    token_labels: torch.Tensor,
+    attention_mask: torch.Tensor,
+    table_labels: Optional[torch.Tensor] = None,
+    row_labels: Optional[torch.Tensor] = None,
+    header_labels: Optional[torch.Tensor] = None,
+    match_labels: Optional[torch.Tensor] = None,
+    col_labels: Optional[torch.Tensor] = None,
+    class_weights: Optional[torch.Tensor] = None,
+    qa_link_labels: Optional[List[torch.Tensor]] = None,
+    contrastive_weight: float = 0.5,
+) -> Dict[str, torch.Tensor]:
+    """
+    Compute joint loss with contrastive learning for Step 2.
+
+    Adds InfoNCE-style contrastive loss to help learn better span representations.
+    """
+    # Get base losses
+    losses = compute_loss(
+        outputs=outputs,
+        token_labels=token_labels,
+        attention_mask=attention_mask,
+        table_labels=table_labels,
+        row_labels=row_labels,
+        header_labels=header_labels,
+        match_labels=match_labels,
+        col_labels=col_labels,
+        class_weights=class_weights,
+        qa_link_labels=qa_link_labels,
+    )
+
+    # Contrastive loss for Step 2 is DISABLED for zero-shot matching
+    # Zero-shot matching - no contrastive loss needed
+    losses["contrastive_loss"] = torch.tensor(0.0, device=token_labels.device)
+
+    return losses
+
+
+if __name__ == "__main__":
+    print("Testing DEXTR LiLT model with CASCADING architecture...")
+    print("=" * 60)
+
+    model = DEXTRLiLT(
+        model_name="nielsr/lilt-xlm-roberta-base",
+        max_seq_len=1024,
+    )
+
+    total_params = sum(p.numel() for p in model.parameters())
+    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    print(f"Total parameters: {total_params:,}")
+    print(f"Trainable parameters: {trainable_params:,}")
+
+    # Test forward
+    batch_size = 2
+    seq_len = 128
+    query_len = 16
+
+    input_ids = torch.randint(0, 1000, (batch_size, seq_len))
+    attention_mask = torch.ones(batch_size, seq_len, dtype=torch.long)
+    # Create valid bboxes: [x1, y1, x2, y2] where x1<x2 and y1<y2, values in [0, 1000]
+    x1 = torch.randint(0, 500, (batch_size, seq_len))
+    y1 = torch.randint(0, 500, (batch_size, seq_len))
+    x2 = x1 + torch.randint(10, 500, (batch_size, seq_len))
+    y2 = y1 + torch.randint(10, 500, (batch_size, seq_len))
+    # Clamp to valid range [0, 1000]
+    bbox = torch.stack([x1, y1, x2.clamp(max=1000), y2.clamp(max=1000)], dim=-1)
+    query_input_ids = torch.randint(0, 1000, (batch_size, query_len))
+    query_attention_mask = torch.ones(batch_size, query_len, dtype=torch.long)
+
+    # GT spans (only used for loss labels in cascading architecture)
+    gt_answer_spans = [[(5, 10), (20, 25)], [(8, 12)]]
+    gt_question_spans = [[(2, 5), (17, 20)], [(5, 8)]]
+
+    print("\nRunning forward pass (CASCADING architecture)...")
+    print("  - Step 1: Token classification → predict Q/A/H/TABLE/O")
+    print("  - Step 1.5: Extract predicted regions, QALinker pairs them")
+    print("  - Step 2: Query matches to predicted Q embeddings")
+    model.eval()
+    with torch.no_grad():
+        outputs = model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            bbox=bbox,
+            query_input_ids=query_input_ids,
+            query_attention_mask=query_attention_mask,
+            gt_answer_spans=gt_answer_spans,
+            gt_question_spans=gt_question_spans,
+        )
+
+    print(f"\nOutput shapes:")
+    print(f"  token_logits: {outputs['token_logits'].shape}")
+    print(f"  table_logits: {outputs['table_logits'].shape}")
+    print(f"  header_logits: {outputs['header_logits'].shape}")
+    print(f"  row_logits: {outputs['row_logits'].shape}")
+    print(f"  col_scores: {outputs['col_scores'].shape}")
+    print(f"  query_embedding: {outputs['query_embedding'].shape}")
+
+    # Predicted regions from Step 1
+    print(f"\nPredicted regions from Step 1:")
+    print(f"  pred_q_regions: {outputs['pred_q_regions']}")
+    print(f"  pred_a_regions: {outputs['pred_a_regions']}")
+
+    # Match scores may not exist if no regions predicted
+    if "match_scores" in outputs:
+        print(f"  match_scores: {outputs['match_scores'].shape}")
+    else:
+        print("  match_scores: None (no regions predicted)")
+
+    # Test loss
+    print("\nTesting loss computation...")
+    token_labels = torch.randint(0, NUM_LABELS, (batch_size, seq_len))
+    table_labels = torch.randint(0, 2, (batch_size, seq_len))
+    row_labels = torch.randint(0, 3, (batch_size, seq_len))
+    header_labels = torch.randint(0, 2, (batch_size, seq_len))
+
+    # For cascading: match_labels should point to predicted Q region index
+    # This would normally be computed by matching GT fields to predicted regions
+    match_labels = None
+    if "match_scores" in outputs and outputs["match_scores"].shape[1] > 0:
+        match_labels = torch.zeros(batch_size, dtype=torch.long)
+
+    losses = compute_loss(
+        outputs,
+        token_labels=token_labels,
+        attention_mask=attention_mask,
+        table_labels=table_labels,
+        row_labels=row_labels,
+        header_labels=header_labels,
+        match_labels=match_labels,
+    )
+
+    print(f"\nLosses:")
+    for name, value in losses.items():
+        print(f"  {name}: {value.item():.4f}")
+
+    print("\n" + "=" * 60)
+    print("DEXTR LiLT CASCADING architecture test passed!")