<jupyter_start><jupyter_text>Load Assay TSVs<jupyter_code>assays = [x for x in os.listdir(assays_path) if x.endswith('.tsv')]
print(assays)
assay_name = assays[2] #pick which assay
assay_name
def read_assay(assay_path):
    assay = {}
    with open(assay_path, "r") as f:
        lines = [x.strip().split('\t') for x in f.readlines()]
        for i, column in enumerate(lines[0]):
            assay[column] = [line[i] for line in lines[1:]]
    return assay
assay = read_assay(assays_path + assay_name)
assay.keys()
result = np.array([float(result) for result in assay['result']], dtype=np.float32)
binding = np.array([int(result) for result in assay['binding']], dtype=np.int32)
binding.sum()
num_instances = len(assay['CanonicalSMILES'])
print(num_instances)<jupyter_output>7462<jupyter_text>Load Pretrained Transformer<jupyter_code>import torch
import torch.nn as nn
import torch.nn.functional as F
import itertools

import sys
sys.path.insert(1, '../') #make parent folder visible
from transformer import Transformer, create_masks, nopeak_mask
checkpoint_dir = "../checkpoints/"
MAX_LEN = 256
MODEL_DIM = 512
N_LAYERS = 6
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(DEVICE)
TRANSFORMER_DEVICE = DEVICE #torch.device("cpu")
PRINTABLE_ASCII_CHARS = 95

_extra_chars = ["seq_start", "seq_end", "pad"]
EXTRA_CHARS = {key: chr(PRINTABLE_ASCII_CHARS + i) for i, key in enumerate(_extra_chars)}
ALPHABET_SIZE = PRINTABLE_ASCII_CHARS + len(EXTRA_CHARS)
def find_ckpts(*args, **kwargs):
    ckpts = os.listdir(checkpoint_dir)
    str_args = [str(x) for x in itertools.chain(args, kwargs.values())]
    return [checkpoint_dir + ckpt for ckpt in ckpts if all([arg in ckpt.replace(".", "_").split("_") for arg in str_args])]
def encode_char(c):
    return ord(c) - 32

def decode_char(n):
    return chr(n + 32)
def encode_string(string, start_char=chr(0)):
    return torch.tensor([ord(start_char)] + [encode_char(c) for c in string])

def encode_string_np(string, start_char=chr(0), pad_char=chr(0)):
    if len(string) > 255:
        string = string[:255]
        
    arr = np.full((256,), ord(pad_char), dtype=np.float32)
    arr[:len(string)+1] = np.array([ord(start_char)] + [encode_char(c) for c in string])
    return arr
def pad_tensors(tensors, pad_char=chr(0), max_len=None):
    if not max_len:
        max_len = max([t.shape[0] for t in tensors]) + 1
        
    padded_tensors = torch.full((len(tensors), max_len), ord(pad_char), dtype=torch.long)
    for i, tensor in enumerate(tensors):
        padded_tensors[i, 0:tensor.shape[0]] = tensor
        
    return padded_tensors<jupyter_output><empty_output><jupyter_text>Select Weights Checkpoint<jupyter_code>found = find_ckpts()
print(found)
load_path = found[0] if len(found) > 0 else ""
model = Transformer(ALPHABET_SIZE, MODEL_DIM, N_LAYERS)
model = nn.DataParallel(model)
model = model.to(TRANSFORMER_DEVICE)
checkpoint = torch.load(load_path)
model.load_state_dict(checkpoint['state_dict'])
model = model.eval()<jupyter_output><empty_output><jupyter_text>Create H5PY Dataset File<jupyter_code>import h5py
transformer_epoch = 2
assay_name = assay_name.replace(".tsv", "_" + str(transformer_epoch) + ".hdf5")
assay_path = embeddings_path + assay_name
print(assay_path)
f = h5py.File(assay_path, 'w-')
embeddings = f.create_dataset("embeddings", (num_instances, 256, 512), dtype=np.float32)
result_dset = f.create_dataset("result", (num_instances,), dtype=np.float32)
binding_dset = f.create_dataset("binding", (num_instances,), dtype=np.int32)
smiles_enc = f.create_dataset("smiles", (num_instances, 256), dtype=np.float32)
with torch.no_grad():
    for i, smiles in enumerate(assay['CanonicalSMILES']):
        encoded = encode_string(smiles, start_char=EXTRA_CHARS['seq_start']).unsqueeze(0).to(TRANSFORMER_DEVICE)
        encoded = encoded[:,:MAX_LEN]
        mask = create_masks(encoded)
        embedding = model.module.encoder(encoded, mask)[0].cpu().numpy()
        embeddings[i,:embedding.shape[0],:] = embedding
        result_dset[i] = result[i]
        binding_dset[i] = binding[i]
        
        encoded = encode_string_np(smiles, start_char=EXTRA_CHARS['seq_start'], pad_char=EXTRA_CHARS['pad'])
        encoded = encoded / ALPHABET_SIZE
        smiles_enc[i,:] = encoded
        
        if i % 1000 == 0:
            print(i)
f.close()<jupyter_output><empty_output>
StarcoderdataJupyterStructuredCleanDedup