File size: 5,187 Bytes
12fea4a | 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 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 | import os
import random
import pandas as pd
from datasets import Dataset, DatasetDict
# import selfies as sf
# from rdkit import Chem
# from rdkit import RDLogger
# RDLogger.DisableLog('rdApp.*')
import sys
from transformers import EsmTokenizer
# ---------------------------------------------------------------------
# CONFIG
# ---------------------------------------------------------------------
# Set PROJECT_ROOT or edit paths below
PROJECT_ROOT = os.environ.get('PROJECT_ROOT', 'path/to/project')
FASTA_PATH = os.path.join(PROJECT_ROOT, 'path/to/protein2pam_db_v1.0/cas9_full.fasta')
OUTPUT_DIR = os.path.join(PROJECT_ROOT, 'data/cas9/cas9_dataset_esm2_tokenized_under_1125')
MAX_TOKENS_PER_BATCH = 8192
TRAIN_RATIO, VAL_RATIO, TEST_RATIO = 0.8, 0.1, 0.1
# ------------------------------------------------------------
# 1) Load sequences from FASTA file
# ------------------------------------------------------------
from Bio import SeqIO
sequences = []
print("Loading sequences from FASTA and validating…")
for record in SeqIO.parse(FASTA_PATH, "fasta"):
s = str(record.seq)
# skip empty sequences
if not s:
continue
sequences.append(s)
# remove all duplicates
sequences = list(set(sequences))
## filter out any sequence that contains non-natural amino acids
print(f"len(sequences) before filtering: {len(sequences)}")
sequences = [
seq for seq in sequences
if all(aa in "ACDEFGHIKLMNPQRSTVWY" for aa in seq) and len(seq) <= 1125
]
print(f"len(sequences) after filtering: {len(sequences)}")
print(f"Average length of sequences: {sum(len(seq) for seq in sequences) / len(sequences)}")
# ---- Initialize ESM tokenizer ----
tokenizer = EsmTokenizer.from_pretrained("facebook/esm2_t33_650M_UR50D")
print(f"Vocabulary Size: {tokenizer.vocab_size}")
# shuffle before splitting
random.shuffle(sequences)
n = len(sequences)
n_train = int(n * TRAIN_RATIO)
n_val = int(n * VAL_RATIO)
train_sequences = sequences[:n_train]
val_sequences = sequences[n_train:n_train + n_val]
test_sequences = sequences[n_train + n_val:]
# ------------------------------------------------------------
# 2) helper: tokenize sequences using the ESM tokenizer
# ------------------------------------------------------------
def tokenize_sequences(sequences):
enc = {"input_ids": [], "attention_mask": []}
for s in sequences:
res = tokenizer(s, add_special_tokens=True)
enc["input_ids"].append(res["input_ids"])
enc["attention_mask"].append(res["attention_mask"])
return enc
# ------------------------------------------------------------
# 3) build batched dataset (token-based batching: groups sequences by length, packs up to max_tokens)
# ------------------------------------------------------------
from tqdm import tqdm
def build_batched_dataset(sequences, max_tokens=8192):
# tokenize all first
toks = tokenize_sequences(sequences)
input_ids_list = toks["input_ids"]
attn_mask_list = toks["attention_mask"]
# add lengths
items = []
for ids, mask in tqdm(zip(input_ids_list, attn_mask_list), total=len(input_ids_list), desc="Collecting items with length"):
items.append({
"input_ids": ids,
"attention_mask": mask,
"length": len(ids),
})
# sort by length
items.sort(key=lambda x: x["length"])
batched_input_ids = []
batched_attention_masks = []
batched_lengths = []
batched_batch_sizes = []
i = 0
n = len(items)
while i < n:
L = items[i]["length"]
# how many of length L can we pack?
max_bs = max_tokens // L
if max_bs < 1:
max_bs = 1
# collect up to max_bs items with same length
cur_ids = []
cur_masks = []
taken = 0
j = i
while j < n and items[j]["length"] == L and taken < max_bs:
cur_ids.append(items[j]["input_ids"]) # length L
cur_masks.append(items[j]["attention_mask"]) # length L
j += 1
taken += 1
# now this is one batch: shape (B, L)
batched_input_ids.append(cur_ids)
batched_attention_masks.append(cur_masks)
batched_lengths.append(L)
batched_batch_sizes.append(taken)
i = j
ds = Dataset.from_dict({
"input_ids": batched_input_ids, # list of (B, L)
"attention_mask": batched_attention_masks,
"seq_length": batched_lengths,
"batch_size": batched_batch_sizes,
})
return ds
# ------------------------------------------------------------
# 4) build batched datasets for each split
# ------------------------------------------------------------
train_ds = build_batched_dataset(train_sequences, max_tokens=MAX_TOKENS_PER_BATCH)
val_ds = build_batched_dataset(val_sequences, max_tokens=MAX_TOKENS_PER_BATCH)
test_ds = build_batched_dataset(test_sequences, max_tokens=MAX_TOKENS_PER_BATCH)
dsdict = DatasetDict(
{
"train": train_ds,
"validation": val_ds,
"test": test_ds,
}
)
os.makedirs(OUTPUT_DIR, exist_ok=True)
dsdict.save_to_disk(OUTPUT_DIR)
print(f"saved to {OUTPUT_DIR}")
|