pCoMole / cas9 /data /cas9_pre_batching.py

Maximilian Holsman

Claude Opus 5

Add Cas9 task

12fea4a about 1 month ago

5.19 kB

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}")

	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}")