| from transformers import PreTrainedModel | |
| from transformers.models.bert.modeling_bert import BertOnlyMLMHead | |
| from peptriever.model.bert_embedding import BertEmbeddingConfig, BertForEmbedding | |
| class BiEncoderConfig(BertEmbeddingConfig): | |
| max_length1: int | |
| max_length2: int | |
| class BiEncoder(PreTrainedModel): | |
| config_class = BiEncoderConfig | |
| def __init__(self, config: BiEncoderConfig): | |
| super().__init__(config) | |
| config1 = _replace_max_length(config, "max_length1") | |
| self.bert1 = BertForEmbedding(config1) | |
| config2 = _replace_max_length(config, "max_length2") | |
| self.bert2 = BertForEmbedding(config2) | |
| self.post_init() | |
| def forward(self, x1, x2): | |
| y1 = self.forward1(x1) | |
| y2 = self.forward2(x2) | |
| return {"y1": y1, "y2": y2} | |
| def forward2(self, x2): | |
| y2 = self.bert2(input_ids=x2["input_ids"]) | |
| return y2 | |
| def forward1(self, x1): | |
| y1 = self.bert1(input_ids=x1["input_ids"]) | |
| return y1 | |
| class BiEncoderWithMaskedLM(PreTrainedModel): | |
| config_class = BiEncoderConfig | |
| def __init__(self, config: BiEncoderConfig): | |
| super().__init__(config=config) | |
| config1 = _replace_max_length(config, "max_length1") | |
| self.bert1 = BertForEmbedding(config1) | |
| self.lm_head1 = BertOnlyMLMHead(config=config1) | |
| config2 = _replace_max_length(config, "max_length2") | |
| self.bert2 = BertForEmbedding(config2) | |
| self.lm_head2 = BertOnlyMLMHead(config=config2) | |
| self.post_init() | |
| def forward(self, x1, x2): | |
| y1, state1 = self.bert1.forward_with_state(input_ids=x1["input_ids"]) | |
| y2, state2 = self.bert2.forward_with_state(input_ids=x2["input_ids"]) | |
| scores1 = self.lm_head1(state1) | |
| scores2 = self.lm_head2(state2) | |
| outputs = {"y1": y1, "y2": y2, "scores1": scores1, "scores2": scores2} | |
| return outputs | |
| def _replace_max_length(config, length_key): | |
| c1 = config.__dict__.copy() | |
| c1["max_position_embeddings"] = c1.pop(length_key) | |
| config1 = BertEmbeddingConfig(**c1) | |
| return config1 | |