|
|
import random |
|
|
import numpy as np |
|
|
import torch |
|
|
import torch.nn as nn |
|
|
from torch.nn import functional as F |
|
|
|
|
|
def set_seed(seed): |
|
|
random.seed(seed) |
|
|
np.random.seed(seed) |
|
|
torch.manual_seed(seed) |
|
|
torch.cuda.manual_seed_all(seed) |
|
|
|
|
|
def top_k_logits(logits, k): |
|
|
v, ix = torch.topk(logits, k) |
|
|
out = logits.clone() |
|
|
out[out < v[:, [-1]]] = -float('Inf') |
|
|
return out |
|
|
|
|
|
@torch.no_grad() |
|
|
def sample(model, x, steps, temperature=1.0, sample=False, top_k=None): |
|
|
block_size = model.get_block_size() |
|
|
model.eval() |
|
|
for k in range(steps): |
|
|
x_cond = x if x.size(1) <= block_size else x[:, -block_size:] |
|
|
logits, _ = model(x_cond) |
|
|
logits = logits[:, -1, :] / temperature |
|
|
if top_k is not None: |
|
|
logits = top_k_logits(logits, top_k) |
|
|
probs = F.softmax(logits, dim=-1) |
|
|
if sample: |
|
|
ix = torch.multinomial(probs, num_samples=1) |
|
|
else: |
|
|
_, ix = torch.topk(probs, k=1, dim=-1) |
|
|
x = torch.cat((x, ix), dim=1) |
|
|
|
|
|
return x |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
@torch.no_grad() |
|
|
def sample(model, x, steps, temperature=1.0,boundary=None): |
|
|
block_size = model.get_block_size() |
|
|
model.eval() |
|
|
for k in range(steps): |
|
|
x_cond = x if x.size(1) <= block_size else x[:, -block_size:] |
|
|
logits, _ = model(x_cond,boundary=boundary) |
|
|
logits = logits[:, -1, :] / temperature |
|
|
probs = F.softmax(logits, dim=-1) |
|
|
ix = torch.multinomial(probs, num_samples=1) |
|
|
x = torch.cat((x, ix), dim=1) |
|
|
return x |
|
|
|
|
|
'L_5*C(=O)NCc1cccc(OC)c1.*c1nsc2ccccc12COc1cccc(CNC(=O)c2cccc(NC(=O)c3nsc4ccccc34)c2)c1' |
|
|
|
|
|
|
|
|
def sample_L(i,option='string'): |
|
|
|
|
|
prefix = 'L_'+str(i) |
|
|
string_input = prefix + '*O=C1NN=Cc2c1cccc2.*O=C(C1CC1)N1CCNCC1' |
|
|
array_input = [vocab[a] for a in ['<bos>'] + list(string_input)] |
|
|
boundary = [len(array_input)] |
|
|
tensor_input = torch.tensor(array_input,device='cuda').unsqueeze(0).repeat(32,1) |
|
|
boundary = boundary*32 |
|
|
tensor_output = sample(model,tensor_input,250,boundary=boundary) |
|
|
strings_output = [] |
|
|
for j in range(tensor_output.shape[0]): |
|
|
list_string_output = [inv[a] for a in tensor_output[j,boundary[j]:].cpu().numpy() if a != vocab['<pad>']] |
|
|
|
|
|
|
|
|
if list_string_output[-1] == '<eos>': |
|
|
list_string_output = list_string_output[:-1] |
|
|
string_output = ''.join(list_string_output) |
|
|
strings_output.append(string_output) |
|
|
print(string_output) |
|
|
for j in range(tensor_output.shape[0]): |
|
|
if test_valid(strings_output[j]): |
|
|
print(1) |
|
|
else: |
|
|
print(0) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
['<bos>', 'L', '_', '5', '*', 'C', '(', '=', 'O', ')', 'N', 'C', 'c', '1', 'c', 'c', 'c', 'c', '(', 'O', 'C', ')', 'c', '1', '.', '*', 'c', '1', 'n', 's', 'c', '2', 'c', 'c', 'c', 'c', 'c', '1', '2', 'C', 'O', 'c', '1', 'c', 'c', 'c', 'c', '(', 'C', 'N', 'C', '(', '=', 'O', ')', 'c', '2', 'c', 'c', 'c', 'c', '(', 'N', 'C', '(', '=', 'O', ')', 'c', '3', 'n', 's', 'c', '4', 'c', 'c', 'c', 'c', 'c', '3', '4', ')', 'c', '2', ')', 'c', '1', '<eos>'] |
|
|
|
