readme: ---
  language: "en"
  tags:
  - document-retrieval
  - knowledge-distillation
  datasets:
  - ms_marco
  ---
  
  # Intra-Document Cascading (IDCM)
  
  We provide a retrieval trained IDCM model. Our model is trained on MSMARCO-Document with up to 2000 tokens.
  
  This instance can be used to **re-rank a candidate set** of long documents. The base BERT architecure is a 6-layer DistilBERT.
  
  If you want to know more about our intra document cascading model & training procedure using knowledge distillation check out our paper: https://arxiv.org/abs/2105.09816 🎉
  
  For more information, training data, source code, and a minimal usage example please visit: https://github.com/sebastian-hofstaetter/intra-document-cascade
  
  ## Configuration
  
  - Trained with fp16 mixed precision
  - We select the top 4 windows of size (50 + 2*7 overlap words) with our fast CK model and score them with BERT
  - The published code here is only usable for inference (we removed the training code)
  
  ## Model Code
  
  ````python
  from transformers import AutoTokenizer,AutoModel, PreTrainedModel,PretrainedConfig
  from typing import Dict
  import torch
  from torch import nn as nn
  
  class IDCM_InferenceOnly(PreTrainedModel):
      '''
      IDCM is a neural re-ranking model for long documents, it creates an intra-document cascade between a fast (CK) and a slow module (BERT_Cat)
      This code is only usable for inference (we removed the training mechanism for simplicity)
      '''
  
      config_class = IDCM_Config
      base_model_prefix = "bert_model"
  
      def __init__(self,
                   cfg) -> None:
          super().__init__(cfg)
  
          #
          # bert - scoring
          #
          if isinstance(cfg.bert_model, str):
              self.bert_model = AutoModel.from_pretrained(cfg.bert_model)
          else:
              self.bert_model = cfg.bert_model
  
          #
          # final scoring (combination of bert scores)
          #
          self._classification_layer = torch.nn.Linear(self.bert_model.config.hidden_size, 1)
          self.top_k_chunks = cfg.top_k_chunks
          self.top_k_scoring = nn.Parameter(torch.full([1,self.top_k_chunks], 1, dtype=torch.float32, requires_grad=True))
  
          #
          # local self attention
          #
          self.padding_idx= cfg.padding_idx
          self.chunk_size = cfg.chunk_size
          self.overlap = cfg.overlap
          self.extended_chunk_size = self.chunk_size + 2 * self.overlap
  
          #
          # sampling stuff
          #
          self.sample_n = cfg.sample_n
          self.sample_context = cfg.sample_context
  
          if self.sample_context == "ck":
              i = 3
              self.sample_cnn3 = nn.Sequential(
                          nn.ConstantPad1d((0,i - 1), 0),
                          nn.Conv1d(kernel_size=i, in_channels=self.bert_model.config.dim, out_channels=self.bert_model.config.dim),
                          nn.ReLU()
                          ) 
          elif self.sample_context == "ck-small":
              i = 3
              self.sample_projector = nn.Linear(self.bert_model.config.dim,384)
              self.sample_cnn3 = nn.Sequential(
                          nn.ConstantPad1d((0,i - 1), 0),
                          nn.Conv1d(kernel_size=i, in_channels=384, out_channels=128),
                          nn.ReLU()
                          ) 
  
          self.sampling_binweights = nn.Linear(11, 1, bias=True)
          torch.nn.init.uniform_(self.sampling_binweights.weight, -0.01, 0.01)
          self.kernel_alpha_scaler = nn.Parameter(torch.full([1,1,11], 1, dtype=torch.float32, requires_grad=True))
  
          self.register_buffer("mu",nn.Parameter(torch.tensor([1.0, 0.9, 0.7, 0.5, 0.3, 0.1, -0.1, -0.3, -0.5, -0.7, -0.9]), requires_grad=False).view(1, 1, 1, -1))
          self.register_buffer("sigma", nn.Parameter(torch.tensor([0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1]), requires_grad=False).view(1, 1, 1, -1))
          
  
      def forward(self,
                  query: Dict[str, torch.LongTensor],
                  document: Dict[str, torch.LongTensor],
                  use_fp16:bool = True,
                  output_secondary_output: bool = False):
  
          #
          # patch up documents - local self attention
          #
          document_ids = document["input_ids"][:,1:]
          if document_ids.shape[1] > self.overlap:
              needed_padding = self.extended_chunk_size - (((document_ids.shape[1]) % self.chunk_size)  - self.overlap)
          else:
              needed_padding = self.extended_chunk_size - self.overlap - document_ids.shape[1]
          orig_doc_len = document_ids.shape[1]
  
          document_ids = nn.functional.pad(document_ids,(self.overlap, needed_padding),value=self.padding_idx)
          chunked_ids = document_ids.unfold(1,self.extended_chunk_size,self.chunk_size)
  
          batch_size = chunked_ids.shape[0]
          chunk_pieces = chunked_ids.shape[1]
  
  
          chunked_ids_unrolled=chunked_ids.reshape(-1,self.extended_chunk_size)
          packed_indices = (chunked_ids_unrolled[:,self.overlap:-self.overlap] != self.padding_idx).any(-1)
          orig_packed_indices 