Spaces:
Sleeping
Sleeping
File size: 4,528 Bytes
fc4a129 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 | 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)
|