Update README.md

Add environment/dependencies, document layer-truncation architecture

README.md

@@ -30,7 +30,6 @@ tags:
 
 QRRanker is a lightweight reranking framework that leverages **Query-focused Retrieval (QR) heads** to produce continuous relevance scores, enabling effective listwise reranking with small-scale models.
 
-
 ## Model Description
 
 Built upon the existing analysis of retrieval heads in large language models, QRRanker trains models to estimate passage–query relevance using the attention scores of selected **Query-focused Retrieval (QR) heads**. These heads are identified through QR score computation on seed data and are particularly effective at capturing query-document relevance signals.
@@ -42,293 +41,158 @@ Our approach provides a **listwise solution** that leverages the holistic inform
 - **Listwise Reranking**: Leverages holistic information within the entire candidate shortlist during ranking
 - **Continuous Relevance Scores**: Enables training on arbitrary retrieval datasets without requiring Likert-scale supervision
 - **Selective Head Usage**: Focuses on top-performing QR attention heads
+- **Layer Truncation**: Only the first 25 of 36 layers are retained — all QR heads fall within layers 17–24, so deeper layers are unnecessary
 - **Memory Enhancement**: Optional contextual summaries for improved accuracy on long narratives and dialogues
 
-## Quick Start
+## Architecture
 
+This model is a **layer-truncated** version of Qwen3-4B-Instruct-2507. The original model has 36 transformer layers, but only the first **25 layers** are retained. The top-performing QR heads (layers 17–24) all fall within this range — deeper layers contribute no useful QR signal but consume extra computation and memory.
 
-### Basic Usage
+Key design choices in `modeling_qwen3_qr.py`:
 
-```python
-import torch
-from transformers import AutoModel, AutoConfig, AutoTokenizer
+- **`Qwen3ConfigGating`**: Extends `Qwen3Config` with `qr_start_layer`, `qr_end_layer`, `qr_head_list`, and `qr_head_list_mapped` (head indices remapped relative to `qr_start_layer`)
+- **Layer construction**: Only instantiates `qr_end_layer` (25) layers instead of all `num_hidden_layers` (36)
+- **No final norm**: Skips `self.norm(hidden_states)` since we only need intermediate KV/query caches, not the final hidden state
+- **`DynamicCacheWithQuery`**: Custom KV-cache that additionally stores query states at specified token positions during the forward pass
 
-# Load model
-config = AutoConfig.from_pretrained("MindscapeRAG/QRRanker", trust_remote_code=True)
-model = AutoModel.from_pretrained(
-    "MindscapeRAG/QRRanker",
-    config=config,
-    torch_dtype=torch.float16,
-    trust_remote_code=True,
-)
-model.eval()
+### Default Top-16 QR Heads
 
-# Load tokenizer
-tokenizer = AutoTokenizer.from_pretrained("MindscapeRAG/QRRanker", trust_remote_code=True)
 ```
-
-## Input Data Format
-
-Input data should be in JSON format. Each sample contains the following fields:
-
-```json
-{
-    "id": "sample_001",
-    "question": "What is the capital of France?",
-    "answer": "Paris",
-    "paragraphs": [
-        {
-            "idx": 0,
-            "title": "France",
-            "paragraph_text": "Paris is the capital and largest city of France...",
-            "is_supporting": true
-        },
-        {
-            "idx": 1,
-            "title": "Germany",
-            "paragraph_text": "Berlin is the capital of Germany...",
-            "is_supporting": false
-        }
-    ],
-    "summary": "Optional summary text..."
-}
+Layer-Head: 20-15, 21-11, 17-27, 23-10, 22-4, 21-10, 21-8, 21-18,
+            18-15, 18-19, 17-25, 17-17, 24-13, 17-4, 19-12, 21-31
 ```
 
-### Field Description
-
-| Field | Type | Required | Description |
-|-------|------|----------|-------------|
-| `id` | string | Yes | Unique sample identifier |
-| `question` | string | Yes | User query/question |
-| `answer` | string | No | Ground truth answer (for evaluation) |
-| `paragraphs` | list | Yes | List of candidate paragraphs |
-| `paragraphs[].idx` | int | Yes | Paragraph index |
-| `paragraphs[].title` | string | No | Paragraph title |
-| `paragraphs[].paragraph_text` | string | Yes | Paragraph content |
-| `paragraphs[].is_supporting` | bool | No | Whether it's a supporting paragraph (for evaluation) |
-| `summary` | string | No | Optional summary information |
+All selected heads fall within layers 17–24, which is why truncation to 25 layers is safe.
 
-## Core Algorithm
+### Model Configuration
 
-### 0. DynamicCacheWithQuery (Custom Cache Class)
+| Parameter | Value | Description |
+|-----------|-------|-------------|
+| `qr_start_layer` | 17 | First layer containing QR heads |
+| `qr_end_layer` | 25 | Layers 0–24 are retained; layers 25–35 are removed |
+| `qr_head_list` | 16 (layer, head) pairs | Top QR heads using original layer indices |
+| `qr_head_list_mapped` | 16 (layer, head) pairs | QR heads with layer indices remapped relative to `qr_start_layer` |
+| `num_hidden_layers` | 36 | Original full model depth (config only, not instantiated) |
+| `num_attention_heads` | 32 | Attention heads per layer |
+| `num_key_value_heads` | 8 | GQA key-value heads per layer |
 
-This custom cache class is essential for QRRanker. It extends the standard `DynamicCache` to also store query states at specified positions.
+## Quick Start
 
+### Loading the Model
 
 ```python
-from typing import Any, Dict, Optional, Tuple
-from transformers.cache_utils import DynamicCache
 import torch
+from transformers import AutoModel, AutoConfig, AutoTokenizer
 
+# Load model — trust_remote_code loads the layer-truncated Qwen3Model
+# and Qwen3ConfigGating automatically via auto_map in config.json
+config = AutoConfig.from_pretrained("MindscapeRAG/QRRanker", trust_remote_code=True)
+model = AutoModel.from_pretrained(
+    "MindscapeRAG/QRRanker",
+    config=config,
+    torch_dtype=torch.float16,
+    trust_remote_code=True,
+).cuda().eval()
 
-class DynamicCacheWithQuery(DynamicCache):
-    """
-    Custom cache class for QRRanker that stores both key/value states and query states.
-    The query states are extracted at specified token positions for attention computation.
-    """
-
-    def __init__(self, query_indices=[]) -> None:
-        super().__init__()
-        self._query_indices = query_indices  # Token indices where query states should be saved
-        self.query_cache = []
-
-    def update(
-            self,
-            key_states: torch.Tensor,
-            value_states: torch.Tensor,
-            layer_idx: int,
-            cache_kwargs: Optional[Dict[str, Any]] = None,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        """
-        Updates the cache with new key_states, value_states, and optionally query_states.
-
-        Parameters:
-            key_states: New key states to cache [batch, num_kv_heads, seq_len, head_dim]
-            value_states: New value states to cache [batch, num_kv_heads, seq_len, head_dim]
-            layer_idx: Index of the layer
-            cache_kwargs: Optional dict containing 'query_states' to cache
-
-        Returns:
-            Tuple of (updated_key_states, updated_value_states)
-        """
-        # Update seen tokens count
-        if layer_idx == 0:
-            self._seen_tokens += key_states.shape[-2]
-
-        # Update key/value cache
-        if key_states is not None:
-            if len(self.key_cache) <= layer_idx:
-                for _ in range(len(self.key_cache), layer_idx):
-                    self.key_cache.append(torch.tensor([]))
-                    self.value_cache.append(torch.tensor([]))
-                self.key_cache.append(key_states)
-                self.value_cache.append(value_states)
-            elif not self.key_cache[layer_idx].numel():
-                self.key_cache[layer_idx] = key_states
-                self.value_cache[layer_idx] = value_states
-            else:
-                self.key_cache[layer_idx] = torch.cat(
-                    [self.key_cache[layer_idx], key_states], dim=-2
-                )
-                self.value_cache[layer_idx] = torch.cat(
-                    [self.value_cache[layer_idx], value_states], dim=-2
-                )
-
-        # Update query cache if query_states provided
-        if cache_kwargs is not None:
-            query_states = cache_kwargs.get("query_states", None)
-        else:
-            query_states = None
-            
-        if query_states is not None:
-            if len(self.query_cache) <= layer_idx:
-                self.query_cache.append(query_states)
-            else:
-                self.query_cache[layer_idx] = torch.cat(
-                    [self.query_cache[layer_idx], query_states], dim=-2
-                )
-                
-        return self.key_cache[layer_idx], self.value_cache[layer_idx]
+tokenizer = AutoTokenizer.from_pretrained("MindscapeRAG/QRRanker")
 ```
 
-### 1. Attention Weight Computation
+### QR Score Computation
+
+After a forward pass, QR scores are computed from the cached query and key states:
 
 ```python
 import math
-import torch
 
-def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
-    """Expand key/value states to match the number of query heads."""
-    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+def repeat_kv(hidden_states, n_rep):
+    """Expand KV heads to match query heads (GQA)."""
+    batch, num_kv_heads, slen, head_dim = hidden_states.shape
     if n_rep == 1:
         return hidden_states
-    hidden_states = hidden_states[:, :, None, :, :].expand(
-        batch, num_key_value_heads, n_rep, slen, head_dim
-    )
-    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
-
-
-def get_causal_mask(attn_weights):
-    """Generate causal attention mask."""
-    query_len, seq_len = attn_weights.size(-2), attn_weights.size(-1)
-    causal_mask = torch.ones_like(attn_weights.transpose(-1, -2).squeeze(0))
-    causal_mask = torch.triu(causal_mask, diagonal=-(seq_len - query_len))
-    causal_mask = causal_mask.transpose(-1, -2)
-    causal_mask = (1 - causal_mask) * torch.finfo(causal_mask.dtype).min
-    return causal_mask
+    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_kv_heads, n_rep, slen, head_dim)
+    return hidden_states.reshape(batch, num_kv_heads * n_rep, slen, head_dim)
 
 
 def get_attn_weights(key_states, query_states):
-    """Compute attention weights between query and key states."""
+    """Compute softmax attention weights with causal mask."""
     bsz, num_heads, q_len, head_dim = query_states.size()
-    num_key_value_heads = key_states.size(1)
-    num_key_value_groups = num_heads // num_key_value_heads
-    kv_seq_len = key_states.size(-2)
-
-    # Expand key states to match query heads
-    key_states = repeat_kv(key_states, num_key_value_groups)
+    num_kv_heads = key_states.size(1)
+    key_states = repeat_kv(key_states, num_heads // num_kv_heads)
 
-    # Scaled dot-product attention
     scale = 1.0 / math.sqrt(head_dim)
-    scaled_queries = query_states * scale
-    attn_weights = torch.matmul(scaled_queries, key_states.transpose(2, 3))
-
-    # Apply causal mask
-    causal_mask = get_causal_mask(attn_weights).to(attn_weights.device)
-    attn_weights += causal_mask.unsqueeze(0)
-    
-    # Softmax normalization
+    attn_weights = torch.matmul(query_states * scale, key_states.transpose(2, 3))
+
+    # Causal mask
+    seq_len = attn_weights.size(-1)
+    causal_mask = torch.ones(num_heads, q_len, seq_len, device=attn_weights.device)
+    causal_mask = torch.triu(causal_mask.transpose(-1, -2), diagonal=-(seq_len - q_len)).transpose(-1, -2)
+    attn_weights += ((1 - causal_mask) * torch.finfo(attn_weights.dtype).min).unsqueeze(0)
+
     attn_lses = torch.logsumexp(attn_weights, dim=-1, keepdim=True)
-    attn_weights = torch.exp(attn_weights - attn_lses)
-    
-    return attn_weights
-```
+    return torch.exp(attn_weights - attn_lses)
 
-### 2. QRRanker Score Computation
 
-```python
-def compute_qr_scores(
-    query_cache,
-    key_cache,
-    qr_head_list,
-    chunk_ranges,
-    query_upper_bound,
-):
+def compute_qr_scores(query_cache, key_cache, qr_head_list, chunk_ranges, query_upper_bound):
     """
-    Compute QRRanker attention scores for document chunks.
-    
+    Compute QRRanker relevance scores for document chunks.
+
     Args:
-        query_cache: List of query states from each layer
-        key_cache: List of key states from each layer
-        qr_head_list: String of QR heads, e.g., "20-15,21-11,17-27,..."
-        chunk_ranges: List of [start, end] token positions for each chunk
-        query_upper_bound: Upper bound token position for query
-    
+        query_cache: List[Tensor] — query states per layer from DynamicCacheWithQuery
+        key_cache: List[Tensor] — key states per layer
+        qr_head_list: str — e.g. "20-15,21-11,17-27,..."
+        chunk_ranges: List[[start, end]] — token ranges for each chunk
+        query_upper_bound: int — upper bound of query token positions
+
     Returns:
-        scores: Tensor of shape [num_chunks] with relevance scores
+        scores: Tensor of shape [num_chunks]
     """
     all_head_scores = []
-    
     for key_state, query_state in zip(key_cache, query_cache):
-        # Compute attention weights
-        attn_weights = get_attn_weights(
-            key_state[:, :, :query_upper_bound, :], 
-            query_state
+        attn_weights = get_attn_weights(key_state[:, :, :query_upper_bound, :], query_state)
+        attn_weights = attn_weights.mean(dim=-2)  # average over query positions
+        chunk_scores = torch.stack(
+            [attn_weights[:, :, s:e].sum(dim=-1) for s, e in chunk_ranges], dim=2
         )
-        # Average over query positions
-        attn_weights = attn_weights.mean(dim=-2)
-        
-        # Aggregate scores for each chunk
-        chunk_scores = []
-        for start, end in chunk_ranges:
-            chunk_scores.append(attn_weights[:, :, start:end].sum(dim=-1))
-        chunk_scores = torch.stack(chunk_scores, dim=2)
         all_head_scores.append(chunk_scores)
-    
-    # Stack all layers: [batch, num_layers, num_heads, num_chunks]
+
+    # [batch, num_layers, num_heads, num_chunks]
     all_head_scores = torch.stack(all_head_scores, dim=1).float()
-    
+
     # Select specific QR heads
     if qr_head_list is not None:
         head_set = [tuple(map(int, h.split('-'))) for h in qr_head_list.split(',')]
-        indices = torch.tensor(head_set).to(all_head_scores.device)
-        layers, heads = indices[:, 0], indices[:, 1]
-        all_head_scores = all_head_scores[:, layers, heads, :]
-    
-    # Sum over selected heads
-    scores = all_head_scores.sum(dim=1).squeeze(0)
-    
-    return scores
+        indices = torch.tensor(head_set, device=all_head_scores.device)
+        all_head_scores = all_head_scores[:, indices[:, 0], indices[:, 1], :]
+
+    return all_head_scores.sum(dim=1).squeeze(0)
 ```
 
-### 3. Complete Inference Pipeline
+### Complete Inference Pipeline
 
 ```python
 from custom_cache_new import DynamicCacheWithQuery
 
 def rerank_documents(model, tokenizer, question, paragraphs, qr_head_list, device):
     """
-    Rerank documents based on QRRanker scores.
-    
+    Rerank candidate paragraphs by QRRanker relevance scores.
+
     Args:
-        model: QRRanker model
-        tokenizer: Tokenizer
+        model: QRRanker model (loaded with trust_remote_code=True)
+        tokenizer: Corresponding tokenizer
         question: Query string
-        paragraphs: List of paragraph dicts with 'idx' and 'paragraph_text'
-        qr_head_list: QR head list string (e.g., "20-15,21-11,17-27,...")
+        paragraphs: List of dicts with 'idx', 'title', 'paragraph_text'
+        qr_head_list: str — e.g. "20-15,21-11,17-27,..."
         device: torch device
-    
+
     Returns:
-        ranked_ids: List of paragraph IDs sorted by relevance
-        scores: Corresponding relevance scores
+        ranked_ids: Paragraph indices sorted by descending relevance
+        ranked_scores: Corresponding scores
     """
-    # Build input sequence
-    prompt_prefix = '<|im_start|>user\n'
-    retrieval_instruction = "Here are some retrieved chunks:\n\n"
-    
-    chunk_part = prompt_prefix + retrieval_instruction
+    # Build input: [chunks] + [query]
+    prompt_prefix = '<|im_start|>user\nHere are some retrieved chunks:\n\n'
+    chunk_part = prompt_prefix
     chunk_ranges = []
-    
+
     for i, p in enumerate(paragraphs):
         text = p.get('title', '') + ': ' + p['paragraph_text']
         chunk_part += f"[{i+1}]"
@@ -337,133 +201,110 @@ def rerank_documents(model, tokenizer, question, paragraphs, qr_head_list, devic
         end = len(chunk_part)
         chunk_ranges.append([start, end])
         chunk_part += '\n\n'
-    
+
     query_part = f"Use the retrieved chunks to answer the user's query.\n\nQuery: {question}"
     full_seq = chunk_part + query_part
-    
-    # Tokenize with offset mapping
-    inputs = tokenizer(
-        full_seq, 
-        max_length=262144, 
-        truncation=True, 
-        return_tensors='pt',
-        return_offsets_mapping=True,
-        add_special_tokens=False
-    )
-    
+
+    # Tokenize
+    inputs = tokenizer(full_seq, max_length=262144, truncation=True,
+                       return_tensors='pt', return_offsets_mapping=True, add_special_tokens=False)
     input_ids = inputs['input_ids'].to(device)
     attention_mask = inputs['attention_mask'].to(device)
     offset_mapping = inputs['offset_mapping'][0]
-    
-    # Build character-to-token mapping
+
+    # Character-to-token mapping
     char_to_token = {}
     for i, (s, e) in enumerate(offset_mapping):
         for j in range(s, e):
             char_to_token[j] = i
-    
-    # Map chunk character ranges to token ranges
-    token_chunk_ranges = []
-    for start, end in chunk_ranges:
-        token_start = char_to_token.get(start, 0)
-        token_end = char_to_token.get(end - 1, 0) + 1
-        token_chunk_ranges.append([token_start, token_end])
-    
-    # Get query token positions
-    query_start_char = full_seq.index(question)
-    query_end_char = query_start_char + len(question) - 1
+
+    token_chunk_ranges = [
+        [char_to_token.get(s, 0), char_to_token.get(e - 1, 0) + 1]
+        for s, e in chunk_ranges
+    ]
+
+    query_start = full_seq.index(question)
     query_positions = list(range(
-        char_to_token[query_start_char],
-        char_to_token[query_end_char] + 1
+        char_to_token[query_start],
+        char_to_token[query_start + len(question) - 1] + 1
     ))
     query_upper_bound = query_positions[-1] + 1
-    
-    # Forward pass with custom cache
+
+    # Forward pass
     with torch.no_grad():
-        # Initialize cache with query token positions
         past_kv = DynamicCacheWithQuery(query_indices=query_positions)
-        
-        # Run model forward pass
         output = model(input_ids, attention_mask, past_key_values=past_kv)
-        
-        # Extract query and key states from cache
-        query_cache = output.past_key_values.query_cache
-        key_cache = output.past_key_values.key_cache
-        
-        # Compute relevance scores
         scores = compute_qr_scores(
-            query_cache, key_cache,
+            output.past_key_values.query_cache,
+            output.past_key_values.key_cache,
             qr_head_list, token_chunk_ranges, query_upper_bound
         )
-    
-    # Sort by scores (descending)
-    sorted_indices = torch.argsort(scores, descending=True).cpu().tolist()
-    ranked_ids = [paragraphs[i]['idx'] for i in sorted_indices]
-    ranked_scores = [float(scores[i]) for i in sorted_indices]
-    
-    return ranked_ids, ranked_scores
-```
-
-## Model Configuration
-
-The model configuration includes the following QRRanker-specific parameters:
-
-| Parameter | Description |
-|-----------|-------------|
-| `qr_start_layer` | Starting layer index for QR heads |
-| `qr_end_layer` | Ending layer index for QR heads |
-| `qr_head_list` | List of (layer, head) tuples for top QR heads |
-
-### Default Top-16 QR Heads
 
-```
-20-15, 21-11, 17-27, 23-10, 22-4, 21-10, 21-8, 21-18,
-18-15, 18-19, 17-25, 17-17, 24-13, 17-4, 19-12, 21-31
+    sorted_idx = torch.argsort(scores, descending=True).cpu().tolist()
+    return [paragraphs[i]['idx'] for i in sorted_idx], [float(scores[i]) for i in sorted_idx]
 ```
 
-## Command Line Usage
+## Input Data Format
 
-```bash
-# Basic inference
-python qr_ranker_inference.py \
-    --base_model MindscapeRAG/QRRanker \
-    --data_path /path/to/data.json \
-    --mode top16
-
-# With summary
-python qr_ranker_inference.py \
-    --base_model MindscapeRAG/QRRanker \
-    --data_path /path/to/data.json \
-    --mode top16 \
-    --use_summary
+```json
+{
+    "id": "sample_001",
+    "question": "What is the capital of France?",
+    "answer": "Paris",
+    "paragraphs": [
+        {
+            "idx": 0,
+            "title": "France",
+            "paragraph_text": "Paris is the capital and largest city of France...",
+            "is_supporting": true
+        }
+    ],
+    "summary": "Optional summary text..."
+}
 ```
 
-### Arguments
+| Field | Type | Required | Description |
+|-------|------|----------|-------------|
+| `id` | string | Yes | Unique sample identifier |
+| `question` | string | Yes | User query/question |
+| `answer` | string | No | Ground truth answer (for evaluation) |
+| `paragraphs` | list | Yes | List of candidate paragraphs |
+| `paragraphs[].idx` | int | Yes | Paragraph index |
+| `paragraphs[].title` | string | No | Paragraph title |
+| `paragraphs[].paragraph_text` | string | Yes | Paragraph content |
+| `paragraphs[].is_supporting` | bool | No | Whether it's a supporting paragraph (for evaluation) |
+| `summary` | string | No | Optional summary information |
 
-| Argument | Type | Default | Description |
-|----------|------|---------|-------------|
-| `--base_model` | str | required | Path to QRRanker model |
-| `--data_path` | str | required | Path to input data file |
-| `--output_dir` | str | `./outputs` | Output directory |
-| `--mode` | str | `top16` | Mode: `full` (all heads) or `top16` (selected heads) |
-| `--qr_head_list` | str | None | Custom QR head list |
-| `--use_summary` | flag | False | Use summary field in data |
+## Environment
 
+| Package | Version |
+|---------|---------|
+| Python | 3.10 |
+| torch | 2.7.1 |
+| transformers | 4.53.0 |
+| flash-attn | (required for `flash_attention_2`) |
+| safetensors | 0.5.3 |
+| tokenizers | 0.21.2 |
 
+```bash
+pip install torch==2.7.1 transformers==4.53.0 safetensors
+pip install flash-attn --no-build-isolation
+```
 
-If you use our QRRanker, please kindly cite:
+## Citation
 
 ```bibtex
 @misc{li2026queryfocusedmemoryawarererankerlong,
-      title={Query-focused and Memory-aware Reranker for Long Context Processing}, 
+      title={Query-focused and Memory-aware Reranker for Long Context Processing},
       author={Yuqing Li and Jiangnan Li and Mo Yu and Guoxuan Ding and Zheng Lin and Weiping Wang and Jie Zhou},
       year={2026},
       eprint={2602.12192},
       archivePrefix={arXiv},
       primaryClass={cs.CL},
-      url={https://arxiv.org/abs/2602.12192}, 
+      url={https://arxiv.org/abs/2602.12192},
 }
 ```
 
 ## License
 
-This project is licensed under the Apache 2.0 License.
+This project is licensed under the Apache 2.0 License.