| import torch | |
| from transformers import PreTrainedModel, XLMRobertaConfig, XLMRobertaModel | |
| class MCLIPConfig(XLMRobertaConfig): | |
| model_type = "M-CLIP" | |
| def __init__(self, transformerDimSize=1024, imageDimSize=768, **kwargs): | |
| self.transformerDimensions = transformerDimSize | |
| self.numDims = imageDimSize | |
| super().__init__(**kwargs) | |
| class MultilingualCLIP(PreTrainedModel): | |
| config_class = MCLIPConfig | |
| def __init__(self, config, *args, **kwargs): | |
| super().__init__(config, *args, **kwargs) | |
| self.transformer = XLMRobertaModel(config) | |
| self.LinearTransformation = torch.nn.Linear( | |
| in_features=config.transformerDimensions, out_features=config.numDims | |
| ) | |
| def forward(self, input_ids, attention_mask): | |
| embs = self.transformer(input_ids=input_ids, attention_mask=attention_mask)[0] | |
| embs2 = (embs * attention_mask.unsqueeze(2)).sum(dim=1) / attention_mask.sum(dim=1)[:, None] | |
| return self.LinearTransformation(embs2), embs | |