Spaces:
Running
Running
| 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) | |
| # Set gradient checkpointing | |
| 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 |