pyserini/encode/_aggretriever.py

#
# Pyserini: Reproducible IR research with sparse and dense representations
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

from typing import Optional
import numpy as np
import torch
from torch import Tensor
import torch.nn as nn
if torch.cuda.is_available():
 from torch.cuda.amp import autocast

from transformers import DistilBertConfig, BertConfig
from transformers import AutoModelForMaskedLM, AutoTokenizer, PreTrainedModel
from pyserini.encode import DocumentEncoder, QueryEncoder

class BERTAggretrieverEncoder(PreTrainedModel):
 config_class = BertConfig
 base_model_prefix = 'encoder'
 load_tf_weights = None

 def __init__(self, config: BertConfig):
 super().__init__(config)
 self.config = config
 self.softmax = nn.Softmax(dim=-1)
 self.encoder = AutoModelForMaskedLM.from_config(config)
 self.tok_proj = torch.nn.Linear(config.hidden_size, 1)
 self.cls_proj = torch.nn.Linear(config.hidden_size, 128)
 self.init_weights()

 # Copied from https://github.com/castorini/dhr/blob/main/tevatron/Aggretriever/utils.py
 def cal_remove_dim(self, dims, vocab_size=30522):
 remove_dims = vocab_size % dims
 if remove_dims > 1000: # the first 1000 tokens in BERT are useless
 remove_dims -= dims
 return remove_dims

 # Copied from https://github.com/castorini/dhr/blob/main/tevatron/Aggretriever/utils.py
 def aggregate(self,
 lexical_reps: Tensor,
 dims: int = 640, 
 remove_dims: int = -198, 
 full: bool = True
 ):
 if full:
 remove_dims = self.cal_remove_dim(dims*2)
 batch_size = lexical_reps.shape[0]
 if remove_dims >= 0:
 lexical_reps = lexical_reps[:, remove_dims:].view(batch_size, -1, dims*2)
 else:
 lexical_reps = torch.nn.functional.pad(lexical_reps, (0, -remove_dims), "constant", 0).view(batch_size, -1, dims*2)
 tok_reps, _ = lexical_reps.max(1)
 positive_tok_reps = tok_reps[:, 0:2*dims:2]
 negative_tok_reps = tok_reps[:, 1:2*dims:2]
 positive_mask = positive_tok_reps > negative_tok_reps
 negative_mask = positive_tok_reps <= negative_tok_reps
 tok_reps = positive_tok_reps * positive_mask - negative_tok_reps * negative_mask
 else:
 remove_dims = self.cal_remove_dim(dims)
 batch_size = lexical_reps.shape[0]
 lexical_reps = lexical_reps[:, remove_dims:].view(batch_size, -1, dims)
 tok_reps, index_reps = lexical_reps.max(1)
 return tok_reps

 # Copied from transformers.models.bert.modeling_bert.BertPreTrainedModel._init_weights
 def _init_weights(self, module):
 """ Initialize the weights """
 if isinstance(module, (torch.nn.Linear, torch.nn.Embedding)):
 # Slightly different from the TF version which uses truncated_normal for initialization
 # cf https://github.com/pytorch/pytorch/pull/5617
 module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
 elif isinstance(module, torch.nn.LayerNorm):
 module.bias.data.zero_()
 module.weight.data.fill_(1.0)
 if isinstance(module, torch.nn.Linear) and module.bias is not None:
 module.bias.data.zero_()

 def init_weights(self):
 self.encoder.init_weights()
 self.tok_proj.apply(self._init_weights)
 self.cls_proj.apply(self._init_weights)

 def forward(
 self,
 input_ids: torch.Tensor,
 attention_mask: Optional[torch.Tensor] = None,
 token_type_ids: torch.Tensor = None,
 skip_mlm: bool = False
 ):
 seq_out = self.encoder(input_ids=input_ids, attention_mask=attention_mask, return_dict=True)
 seq_hidden = seq_out.hidden_states[-1] 
 cls_hidden = seq_hidden[:,0] # get [CLS] embeddings
 term_weights = self.tok_proj(seq_hidden[:,1:]) # batch, seq, 1
 if not skip_mlm:
 logits = seq_out.logits[:,1:] # batch, seq-1, vocab
 logits = self.softmax(logits)
 attention_mask = attention_mask[:,1:].unsqueeze(-1)
 lexical_reps = torch.max((logits * term_weights) * attention_mask, dim=-2).values
 else:
 # w/o MLM
 lexical_reps = torch.zeros(seq_hidden.shape[0], seq_hidden.shape[1], 30522, dtype=seq_hidden.dtype, device=seq_hidden.device) # (batch, len, vocab)
 lexical_reps = torch.scatter(lexical_reps, dim=-1, index=input_ids[:,1:,None], src=term_weights)
 lexical_reps = lexical_reps.max(-2).values

 lexical_reps = self.aggregate(lexical_reps, 640)
 semantic_reps = self.cls_proj(cls_hidden)
 return torch.cat((semantic_reps, lexical_reps), -1)


class DistlBERTAggretrieverEncoder(BERTAggretrieverEncoder):
 config_class = DistilBertConfig
 base_model_prefix = 'encoder'
 load_tf_weights = None


class AggretrieverDocumentEncoder(DocumentEncoder):
 def __init__(self, model_name: str, tokenizer_name=None, device='cuda:0'):
 self.device = device
 if 'distilbert' in model_name.lower():
 self.model = DistlBERTAggretrieverEncoder.from_pretrained(model_name)
 else:
 self.model = BERTAggretrieverEncoder.from_pretrained(model_name)
 self.model.to(self.device)
 self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_name or model_name)

 def encode(self, texts, titles=None, fp16=False, max_length=512, **kwargs):
 if titles is not None:
 texts = [f'{title} {text}' for title, text in zip(titles, texts)]
 else:
 texts = [text for text in texts]
 inputs = self.tokenizer(
 texts,
 max_length=max_length,
 padding="longest",
 truncation=True,
 add_special_tokens=True,
 return_tensors='pt'
 )
 inputs.to(self.device)
 if fp16:
 with autocast():
 with torch.no_grad():
 outputs = self.model(**inputs)
 else:
 outputs = self.model(**inputs)
 return outputs.detach().cpu().numpy()


class AggretrieverQueryEncoder(QueryEncoder):
 def __init__(self, model_name: str, tokenizer_name=None, device='cuda:0'):
 self.device = device
 if 'distilbert' in model_name.lower():
 self.model = DistlBERTAggretrieverEncoder.from_pretrained(model_name)
 else:
 self.model = BERTAggretrieverEncoder.from_pretrained(model_name)
 self.model.to(self.device)
 self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_name or model_name)

 def encode(self, texts, fp16=False, max_length=32, **kwargs):
 texts = [text for text in texts]
 inputs = self.tokenizer(
 texts,
 max_length=max_length,
 padding="longest",
 truncation=True,
 add_special_tokens=True,
 return_tensors='pt'
 )
 inputs.to(self.device)
 if fp16:
 with autocast():
 with torch.no_grad():
 outputs = self.model(**inputs)
 else:
 outputs = self.model(**inputs)
 return outputs.detach().cpu().numpy()