|
import torch
|
|
|
|
from tqdm import tqdm
|
|
from torch.nn.functional import normalize
|
|
from transformers import BertConfig, BertModel, BertTokenizer
|
|
|
|
|
|
class TextEncoder(torch.nn.Module):
|
|
def __init__(self,
|
|
config_path: str,
|
|
out_dim: int,
|
|
load_pretrained: bool = True,
|
|
gradient_checkpointing: bool = False):
|
|
"""
|
|
Args:
|
|
config_path: Path to the config file
|
|
|
|
out_dim: Output dimension of the text representation
|
|
|
|
load_pretrained: Whether to load pretrained weights
|
|
|
|
gradient_checkpointing: Whether to enable gradient checkpointing
|
|
"""
|
|
super().__init__()
|
|
config = BertConfig.from_pretrained(config_path)
|
|
if load_pretrained:
|
|
self.model = BertModel.from_pretrained(config_path, add_pooling_layer=False)
|
|
else:
|
|
self.model = BertModel(config, add_pooling_layer=False)
|
|
self.out = torch.nn.Linear(config.hidden_size, out_dim)
|
|
|
|
|
|
self.model.encoder.gradient_checkpointing = gradient_checkpointing
|
|
|
|
self.tokenizer = BertTokenizer.from_pretrained(config_path)
|
|
|
|
def get_repr(self, texts: list, batch_size: int = 64, verbose: bool = False) -> torch.Tensor:
|
|
"""
|
|
Compute text representation for the given texts
|
|
Args:
|
|
texts: A list of strings
|
|
batch_size: Batch size for inference
|
|
verbose: Whether to print progress
|
|
"""
|
|
device = next(self.parameters()).device
|
|
|
|
text_repr = []
|
|
if verbose:
|
|
iterator = tqdm(range(0, len(texts), batch_size), desc="Computing text embeddings")
|
|
else:
|
|
iterator = range(0, len(texts), batch_size)
|
|
|
|
for i in iterator:
|
|
text_inputs = self.tokenizer.batch_encode_plus(texts[i: i+batch_size],
|
|
return_tensors="pt",
|
|
truncation=True,
|
|
max_length=512,
|
|
padding=True)
|
|
text_inputs = {k: v.to(device) for k, v in text_inputs.items()}
|
|
output = self(text_inputs)
|
|
|
|
text_repr.append(output)
|
|
|
|
text_repr = torch.cat(text_repr, dim=0)
|
|
return normalize(text_repr, dim=-1)
|
|
|
|
def forward(self, inputs: dict):
|
|
"""
|
|
Encode text into text representation
|
|
Args:
|
|
inputs: A dictionary containing the following keys:
|
|
- input_ids: [batch, seq_len]
|
|
- attention_mask: [batch, seq_len]
|
|
- token_type_ids: [batch, seq_len]
|
|
|
|
Returns:
|
|
text_repr: [batch, text_repr_dim]
|
|
"""
|
|
reprs = self.model(**inputs).last_hidden_state[:, 0, :]
|
|
reprs = self.out(reprs)
|
|
return reprs |
