| from transformers.models.bert import modeling_bert |
| from open_clip import CustomTextCLIP |
| from open_clip.hf_model import HFTextEncoder |
| import torch.nn.functional as F |
| from torch import TensorType |
|
|
| def patch_encode_text(): |
|
|
| def encode_text_patched(self, text, normalize: bool = False, output_attentions = False, output_tokens = False): |
| if output_attentions: |
| features, attn_scores = self.text(text, output_attentions = output_attentions, output_tokens = output_tokens) |
| features = F.normalize(features, dim=-1) if normalize else features |
| return features, attn_scores |
| else: |
| features = self.text(text, output_attentions = output_attentions, output_tokens = output_tokens) |
| return F.normalize(features, dim=-1) if normalize else features |
|
|
| def HFText_encoder_patched(self, x: TensorType, output_attentions=False, output_tokens=False): |
| self.output_tokens = output_tokens |
| attn_mask = (x != self.config.pad_token_id).long() |
| out = self.transformer(input_ids=x, attention_mask=attn_mask, output_attentions=output_attentions) |
| if self.output_tokens: |
| tokens = self.proj(out[0]) |
| if output_attentions: |
| return tokens, out[1] |
| else: |
| return tokens |
| else: |
| pooled_out = self.pooler(out, attn_mask) |
| projected = self.proj(pooled_out) |
|
|
| return projected |
|
|
| CustomTextCLIP.encode_text = encode_text_patched |
| HFTextEncoder.forward = HFText_encoder_patched |
|
|
|
|
|
|
