Spaces:
Sleeping
Sleeping
| import torch | |
| from torch import nn | |
| import numpy as np | |
| import os | |
| from transformers import RobertaModel | |
| from transformers import RobertaTokenizer | |
| from transformers import RobertaConfig | |
| def l2norm(X, dim, eps=1e-8): | |
| """L2-normalize columns of X | |
| """ | |
| norm = torch.pow(X, 2).sum(dim=dim, keepdim=True).sqrt() + eps | |
| X = torch.div(X, norm) | |
| return X | |
| class SmilesModel(nn.Module): | |
| def __init__(self, | |
| roberta_model_path, | |
| roberta_tokenizer_path, | |
| smiles_maxlen=300, | |
| vocab_size=181, | |
| max_position_embeddings=505, | |
| num_attention_heads=12, | |
| num_hidden_layers=6, | |
| type_vocab_size=1, | |
| no_txtnorm=False, | |
| feature_dim=768, | |
| embed_size=1024, | |
| **kwargs | |
| ): | |
| super(SmilesModel, self).__init__(**kwargs) | |
| self.smiles_maxlen = smiles_maxlen | |
| self.feature_dim = feature_dim | |
| self.t_prime = nn.Parameter(torch.tensor(np.log(10.0))) | |
| self.bias = nn.Parameter(torch.tensor(-10.0)) | |
| if roberta_tokenizer_path is not None: | |
| self.smiles_tokenizer = RobertaTokenizer.from_pretrained( | |
| roberta_tokenizer_path, max_len=self.smiles_maxlen) | |
| new_tokens = ["@", "/", "\\"] | |
| self.smiles_tokenizer.add_tokens(new_tokens) | |
| if roberta_model_path is None or not os.path.exists(roberta_model_path): | |
| self.smiles_config = RobertaConfig( | |
| vocab_size=vocab_size, | |
| max_position_embeddings=max_position_embeddings, | |
| num_attention_heads=num_attention_heads, | |
| num_hidden_layers=num_hidden_layers, | |
| type_vocab_size=type_vocab_size, | |
| hidden_size=self.feature_dim | |
| ) | |
| self.smiles_model = RobertaModel(config=self.smiles_config) | |
| else: | |
| self.smiles_config = RobertaConfig.from_pretrained( | |
| roberta_model_path) | |
| self.smiles_model = RobertaModel.from_pretrained( | |
| roberta_model_path) | |
| self.smiles_model.resize_token_embeddings(len(self.smiles_tokenizer)) | |
| self.model = self.smiles_model | |
| # ESA | |
| self.linear = nn.Linear(768, embed_size) | |
| self.linear1 = nn.Linear(embed_size, embed_size) | |
| self.dropout = nn.Dropout(0.1) | |
| self.no_txtnorm = no_txtnorm | |
| self.init_weights() | |
| def encode(self, input, lengths): | |
| lengths = lengths -1 | |
| input_ids, attention_mask = input | |
| hidden_states = self.model(input_ids, attention_mask)[0] | |
| hidden_states = hidden_states[:, 1:, :] | |
| cap_emb = self.linear(hidden_states) | |
| cap_emb = self.dropout(cap_emb) | |
| max_len = int(lengths.max()) | |
| mask = torch.arange(max_len).expand(lengths.size(0), max_len).to(lengths.device) | |
| mask = (mask < lengths.long().unsqueeze(1)).unsqueeze(-1) | |
| cap_emb = cap_emb[:, :int(lengths.max()), :] | |
| features_in = self.linear1(cap_emb) | |
| features_in = features_in.masked_fill(mask == 0,-10000) | |
| features_k_softmax = nn.Softmax(dim=1)(features_in-torch.max(features_in,dim=1)[0].unsqueeze(1)) | |
| attn = features_k_softmax.masked_fill(mask == 0,0) | |
| feature_cap = torch.sum(attn * cap_emb,dim=1) | |
| if not self.no_txtnorm: | |
| feature_cap = l2norm(feature_cap, dim=-1) | |
| return feature_cap | |
| def init_weights(self): | |
| """Xavier initialization for the fully connected layer | |
| """ | |
| for m in self.children(): | |
| if isinstance(m, nn.Linear): | |
| r = np.sqrt(6.) / np.sqrt(m.in_features + m.out_features) | |
| m.weight.data.uniform_(-r, r) | |
| m.bias.data.fill_(0) | |
| elif isinstance(m, nn.BatchNorm1d): | |
| m.weight.data.fill_(1) | |
| m.bias.data.zero_() | |
| def load_weights(self, path): | |
| if path is not None: | |
| model_dict = torch.load(path, map_location=torch.device('cpu')) | |
| self.load_state_dict(model_dict) | |
| def get_parameter_number(self): | |
| total_num = sum(p.numel() for p in self.parameters()) | |
| trainable_num = sum(p.numel() for p in self.parameters() if p.requires_grad) | |
| return 'Total parameters: {}, Trainable parameters: {}'.format(total_num, trainable_num) | |