ynuozhang

add inference

c91cc5c 11 months ago

10.1 kB

# non fouling as example
import matplotlib.pyplot as plt
import matplotlib as mpl
import numpy as np
import os
import pandas as pd
from lightning.pytorch import seed_everything
import torch
from tqdm import tqdm
from datasets import Dataset, DatasetDict, Features, Value, Sequence
from transformers import AutoModelForMaskedLM
import sys
from transformers import AutoTokenizer, EsmModel
from datasets import Dataset, DatasetDict
import tqdm

seed_everything(1986)
# -------------------------
# Process from source
# -------------------------
m1 = [
    '[PAD]','A','R','N','D','C','Q','E','G','H',
    'I','L','K','M','F','P','S','T','W','Y','V'
]
m2 = dict(zip(
    ['[PAD]','[UNK]','[CLS]','[SEP]','[MASK]','L',
    'A','G','V','E','S','I','K','R','D','T','P','N',
    'Q','F','Y','M','H','C','W','X','U','B','Z','O'],
    range(30)
))
# Create reverse mapping
reverse_m2 = {v: k for k, v in m2.items()}
sequences = []
labels = []

# Load and process positive sequences
print("Processing positive sequences...")
with np.load('nf-positive.npz') as pos:
    pos_data = pos['arr_0']
for seq in pos_data:
    sequence = ''.join(reverse_m2[token] for token in seq if token != 0)
    sequences.append(sequence)
    labels.append(1)

# Load and process negative sequences
print("Processing negative sequences...")
with np.load('nf-negative.npz') as neg:
    neg_data = neg['arr_0']
for seq in neg_data:
    sequence = ''.join(reverse_m2[token] for token in seq if token != 0)
    sequences.append(sequence)
    labels.append(0)
    
# Build a DataFrame and add stable IDs
ids = [f"seq_{i:06d}" for i in range(len(sequences))]
df = pd.DataFrame({
    "id": ids,
    "sequence": sequences,
    "label": labels,
})
print("Before dedup:", len(df))

df = (
    df
    .drop_duplicates(subset=["sequence"])   # keep first occurrence
    .reset_index(drop=True)
)

print("After dedup:", len(df))
# Save to CSV
df.to_csv("nf_all.csv", index=False)
print("Saved nf_all.csv")

# Save as FASTA for MMseqs
with open("nf_all.fasta", "w") as f:
    for seq_id, seq in zip(df["id"], df["sequence"]):
        f.write(f">{seq_id}\n{seq}\n")

print("Saved nf_all.fasta")

# -------------------------
# RUN MMSEQS IN TERMINAL
# -------------------------

# -------------------------
"""
mkdir -p mmseqs_tmp
mmseqs createdb nf_all.fasta nfDB
mmseqs cluster nfDB nfDB_clu mmseqs_tmp \
  --min-seq-id 0.3 -c 0.8 --cov-mode 0
mmseqs createtsv nfDB nfDB nfDB_clu clusters-nf.tsv
"""
# -------------------------


# -------------------------
# Split based on clusters
# -------------------------

train_fraction = 0.8
csv_path = "nf_all.csv"
clusters_tsv = "clusters-nf.tsv"
rng = np.random.default_rng()

df = pd.read_csv(csv_path)   # must contain: id, sequence, label

# Map id -> index
id_to_index = {sid: i for i, sid in enumerate(df["id"])}

# Read MMseqs clusters
cluster_map = {}  # member_id -> cluster_id (rep_id)
with open(clusters_tsv) as f:
    for line in f:
        if not line.strip():
            continue
        rep_id, member_id = line.strip().split('\t')
        cluster_map[member_id] = rep_id

# Handle singleton sequences (not clustered)
for sid in df["id"]:
    if sid not in cluster_map:
        cluster_map[sid] = sid

# Invert to cluster_id -> dataset indices
cluster_to_indices = {}
for sid, cid in cluster_map.items():
    idx = id_to_index[sid]
    cluster_to_indices.setdefault(cid, []).append(idx)

# Shuffle clusters
cluster_ids = list(cluster_to_indices.keys())
rng.shuffle(cluster_ids)

# Assign clusters to splits
total_n = len(df)
train_target = int(train_fraction * total_n)

train_indices = []
val_indices = []
current_train = 0

for cid in cluster_ids:
    indices = cluster_to_indices[cid]
    if current_train + len(indices) <= train_target:
        train_indices.extend(indices)
        current_train += len(indices)
    else:
        val_indices.extend(indices)

# Create split column
split = np.full(total_n, "val", dtype=object)
split[train_indices] = "train"

# ===== Master CSV with split =====
df_with_split = df.copy()
df_with_split["split"] = split
df_with_split.to_csv("nf_meta_with_split.csv", index=False)

# ===== Other CSVs =====
df_train = df_with_split[df_with_split["split"] == "train"].reset_index(drop=True)
df_val   = df_with_split[df_with_split["split"] == "val"].reset_index(drop=True)

df_train.to_csv("nf_train.csv", index=False)
df_val.to_csv("nf_val.csv", index=False)

# ===== Quick sanity output =====
print("Split counts:")
print(df_with_split["split"].value_counts())
print()
print(f"Train size: {len(df_train)}")
print(f"Val size:   {len(df_val)}")
print("Wrote:")
print("  - sol_meta_with_split.csv")
print("  - sol_train.csv")
print("  - sol_val.csv")


device = torch.device("cuda:0")
print(f"Using device: {device}")

meta_path = "./Classifier_Weight/training_data_cleaned/nf/nf_meta_with_split.csv"
save_path = "./Classifier_Weight/training_data_cleaned/nf/nf_wt_with_embeddings"

tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t33_650M_UR50D")
model = EsmModel.from_pretrained("facebook/esm2_t33_650M_UR50D")
model = model.to(device)
model.eval()


def compute_embeddings(sequences, batch_size=32):
    """Return numpy array of shape (N, hidden_dim)."""
    embeddings = []
    for i in tqdm.trange(0, len(sequences), batch_size):
        batch_sequences = sequences[i:i + batch_size]

        inputs = tokenizer(
            batch_sequences,
            padding=True,
            max_length=1022,
            truncation=True,
            return_tensors="pt"
        )
        inputs = {k: v.to(device) for k, v in inputs.items()}

        with torch.no_grad():
            outputs = model(**inputs)
            last_hidden_states = outputs.last_hidden_state  # (B, L, H)

            attention_mask = inputs["attention_mask"].unsqueeze(-1)
            masked_hidden_states = last_hidden_states * attention_mask
            sum_hidden_states = masked_hidden_states.sum(dim=1)
            seq_lengths = attention_mask.sum(dim=1)
            batch_embeddings = sum_hidden_states / seq_lengths  # (B, H)

        embeddings.append(batch_embeddings.cpu())

    return torch.cat(embeddings, dim=0).numpy()


def create_and_save_datasets():
    # Load sequences + labels + splits
    meta = pd.read_csv(meta_path)
    sequences = meta["sequence"].tolist()
    labels = meta["label"].tolist()
    splits = meta["split"].tolist()

    print(f"Total sequences: {len(sequences)}")
    print("Split counts:", pd.Series(splits).value_counts().to_dict())

    print("Computing ESM embeddings...")
    embeddings = compute_embeddings(sequences)  # shape (N, H)

    full_ds = Dataset.from_dict({
        "sequence": sequences,
        "embedding": embeddings,
        "label": labels,
        "split": splits,
    })

    # Split into train / val, then drop the 'split' column
    train_ds = full_ds.filter(lambda x: x["split"] == "train")
    val_ds   = full_ds.filter(lambda x: x["split"] == "val")

    train_ds = train_ds.remove_columns("split")
    val_ds   = val_ds.remove_columns("split")

    ds_dict = DatasetDict({
        "train": train_ds,
        "val": val_ds,
    })

    ds_dict.save_to_disk(save_path)
    print(f"Saved DatasetDict with train/val to: {save_path}")
    print("Train size:", len(ds_dict["train"]))
    print("Val size:", len(ds_dict["val"]))

    return ds_dict


ds = create_and_save_datasets()

ex = ds["train"][0]
print("\nExample from train:")
print("Sequence:", ex["sequence"])
print("Embedding shape:", np.array(ex["embedding"]).shape)
print("Label:", ex["label"])

torch.cuda.empty_cache()

meta_path = "./Classifier_Weight/training_data_cleaned/nf/nf_meta_with_split.csv"
save_path = "./Classifier_Weight/training_data_cleaned/nf/nf_wt_with_embeddings_unpooled"

tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t33_650M_UR50D")
model = EsmModel.from_pretrained("facebook/esm2_t33_650M_UR50D", add_pooling_layer=False).to(device).eval()

cls_id = tokenizer.cls_token_id
eos_id = tokenizer.eos_token_id

@torch.no_grad()
def embed_one(seq, max_length=1022):
    inputs = tokenizer(seq, padding=False, truncation=True, max_length=max_length, return_tensors="pt")
    inputs = {k: v.to(device) for k, v in inputs.items()}
    out = model(**inputs)
    h = out.last_hidden_state[0]                # (L, H)
    attn = inputs["attention_mask"][0].bool()   # (L,)
    ids = inputs["input_ids"][0]

    keep = attn.clone()
    if cls_id is not None:
        keep &= (ids != cls_id)
    if eos_id is not None:
        keep &= (ids != eos_id)

    hb = h[keep].detach().cpu().to(torch.float16).numpy()  # (Li, H)
    return hb

H = 1280
features = Features({
    "sequence": Value("string"),
    "label": Value("int64"),
    "embedding": Sequence(Sequence(Value("float16"), length=H)),  # (Li, H) as nested lists
    "attention_mask": Sequence(Value("int8")),                    # (Li,)
    "length": Value("int64"),
})

def make_generator(df):
    for seq, lab in tqdm.tqdm(zip(df["sequence"].tolist(), df["label"].astype(int).tolist()), total=len(df)):
        emb = embed_one(seq)  # (Li, H) float16
        emb_list = emb.tolist()
        li = len(emb_list)
        yield {
            "sequence": seq,
            "label": int(lab),
            "embedding": emb_list,
            "attention_mask": [1] * li,
            "length": li,
        }

def build_and_save_split(df, out_dir):
    ds = Dataset.from_generator(make_generator, gen_kwargs={"df": df}, features=features)
    # small batches when writing prevents the 2GB offset overflow
    ds.save_to_disk(out_dir, max_shard_size="1GB")
    return ds

meta = pd.read_csv(meta_path)
train_df = meta[meta["split"] == "train"].reset_index(drop=True)
val_df   = meta[meta["split"] == "val"].reset_index(drop=True)

train_dir = os.path.join(save_path, "train")
val_dir   = os.path.join(save_path, "val")
os.makedirs(save_path, exist_ok=True)

train_ds = build_and_save_split(train_df, train_dir)
val_ds   = build_and_save_split(val_df, val_dir)

ds_dict = DatasetDict({"train": train_ds, "val": val_ds})
ds_dict.save_to_disk(save_path)
print(ds_dict)

	# non fouling as example
	import matplotlib.pyplot as plt
	import matplotlib as mpl
	import numpy as np
	import os
	import pandas as pd
	from lightning.pytorch import seed_everything
	import torch
	from tqdm import tqdm
	from datasets import Dataset, DatasetDict, Features, Value, Sequence
	from transformers import AutoModelForMaskedLM
	import sys
	from transformers import AutoTokenizer, EsmModel
	from datasets import Dataset, DatasetDict
	import tqdm

	seed_everything(1986)
	# -------------------------
	# Process from source
	# -------------------------
	m1 = [
	'[PAD]','A','R','N','D','C','Q','E','G','H',
	'I','L','K','M','F','P','S','T','W','Y','V'
	]
	m2 = dict(zip(
	['[PAD]','[UNK]','[CLS]','[SEP]','[MASK]','L',
	'A','G','V','E','S','I','K','R','D','T','P','N',
	'Q','F','Y','M','H','C','W','X','U','B','Z','O'],
	range(30)
	))
	# Create reverse mapping
	reverse_m2 = {v: k for k, v in m2.items()}
	sequences = []
	labels = []

	# Load and process positive sequences
	print("Processing positive sequences...")
	with np.load('nf-positive.npz') as pos:
	pos_data = pos['arr_0']
	for seq in pos_data:
	sequence = ''.join(reverse_m2[token] for token in seq if token != 0)
	sequences.append(sequence)
	labels.append(1)

	# Load and process negative sequences
	print("Processing negative sequences...")
	with np.load('nf-negative.npz') as neg:
	neg_data = neg['arr_0']
	for seq in neg_data:
	sequence = ''.join(reverse_m2[token] for token in seq if token != 0)
	sequences.append(sequence)
	labels.append(0)

	# Build a DataFrame and add stable IDs
	ids = [f"seq_{i:06d}" for i in range(len(sequences))]
	df = pd.DataFrame({
	"id": ids,
	"sequence": sequences,
	"label": labels,
	})
	print("Before dedup:", len(df))

	df = (
	df
	.drop_duplicates(subset=["sequence"]) # keep first occurrence
	.reset_index(drop=True)
	)

	print("After dedup:", len(df))
	# Save to CSV
	df.to_csv("nf_all.csv", index=False)
	print("Saved nf_all.csv")

	# Save as FASTA for MMseqs
	with open("nf_all.fasta", "w") as f:
	for seq_id, seq in zip(df["id"], df["sequence"]):
	f.write(f">{seq_id}\n{seq}\n")

	print("Saved nf_all.fasta")

	# -------------------------
	# RUN MMSEQS IN TERMINAL
	# -------------------------

	# -------------------------
	"""
	mkdir -p mmseqs_tmp

	mmseqs createdb nf_all.fasta nfDB

	mmseqs cluster nfDB nfDB_clu mmseqs_tmp \
	--min-seq-id 0.3 -c 0.8 --cov-mode 0

	mmseqs createtsv nfDB nfDB nfDB_clu clusters-nf.tsv
	"""
	# -------------------------


	# -------------------------
	# Split based on clusters
	# -------------------------

	train_fraction = 0.8
	csv_path = "nf_all.csv"
	clusters_tsv = "clusters-nf.tsv"
	rng = np.random.default_rng()

	df = pd.read_csv(csv_path) # must contain: id, sequence, label

	# Map id -> index
	id_to_index = {sid: i for i, sid in enumerate(df["id"])}

	# Read MMseqs clusters
	cluster_map = {} # member_id -> cluster_id (rep_id)
	with open(clusters_tsv) as f:
	for line in f:
	if not line.strip():
	continue
	rep_id, member_id = line.strip().split('\t')
	cluster_map[member_id] = rep_id

	# Handle singleton sequences (not clustered)
	for sid in df["id"]:
	if sid not in cluster_map:
	cluster_map[sid] = sid

	# Invert to cluster_id -> dataset indices
	cluster_to_indices = {}
	for sid, cid in cluster_map.items():
	idx = id_to_index[sid]
	cluster_to_indices.setdefault(cid, []).append(idx)

	# Shuffle clusters
	cluster_ids = list(cluster_to_indices.keys())
	rng.shuffle(cluster_ids)

	# Assign clusters to splits
	total_n = len(df)
	train_target = int(train_fraction * total_n)

	train_indices = []
	val_indices = []
	current_train = 0

	for cid in cluster_ids:
	indices = cluster_to_indices[cid]
	if current_train + len(indices) <= train_target:
	train_indices.extend(indices)
	current_train += len(indices)
	else:
	val_indices.extend(indices)

	# Create split column
	split = np.full(total_n, "val", dtype=object)
	split[train_indices] = "train"

	# ===== Master CSV with split =====
	df_with_split = df.copy()
	df_with_split["split"] = split
	df_with_split.to_csv("nf_meta_with_split.csv", index=False)

	# ===== Other CSVs =====
	df_train = df_with_split[df_with_split["split"] == "train"].reset_index(drop=True)
	df_val = df_with_split[df_with_split["split"] == "val"].reset_index(drop=True)

	df_train.to_csv("nf_train.csv", index=False)
	df_val.to_csv("nf_val.csv", index=False)

	# ===== Quick sanity output =====
	print("Split counts:")
	print(df_with_split["split"].value_counts())
	print()
	print(f"Train size: {len(df_train)}")
	print(f"Val size: {len(df_val)}")
	print("Wrote:")
	print(" - sol_meta_with_split.csv")
	print(" - sol_train.csv")
	print(" - sol_val.csv")


	device = torch.device("cuda:0")
	print(f"Using device: {device}")

	meta_path = "./Classifier_Weight/training_data_cleaned/nf/nf_meta_with_split.csv"
	save_path = "./Classifier_Weight/training_data_cleaned/nf/nf_wt_with_embeddings"

	tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t33_650M_UR50D")
	model = EsmModel.from_pretrained("facebook/esm2_t33_650M_UR50D")
	model = model.to(device)
	model.eval()


	def compute_embeddings(sequences, batch_size=32):
	"""Return numpy array of shape (N, hidden_dim)."""
	embeddings = []
	for i in tqdm.trange(0, len(sequences), batch_size):
	batch_sequences = sequences[i:i + batch_size]

	inputs = tokenizer(
	batch_sequences,
	padding=True,
	max_length=1022,
	truncation=True,
	return_tensors="pt"
	)
	inputs = {k: v.to(device) for k, v in inputs.items()}

	with torch.no_grad():
	outputs = model(**inputs)
	last_hidden_states = outputs.last_hidden_state # (B, L, H)

	attention_mask = inputs["attention_mask"].unsqueeze(-1)
	masked_hidden_states = last_hidden_states * attention_mask
	sum_hidden_states = masked_hidden_states.sum(dim=1)
	seq_lengths = attention_mask.sum(dim=1)
	batch_embeddings = sum_hidden_states / seq_lengths # (B, H)

	embeddings.append(batch_embeddings.cpu())

	return torch.cat(embeddings, dim=0).numpy()


	def create_and_save_datasets():
	# Load sequences + labels + splits
	meta = pd.read_csv(meta_path)
	sequences = meta["sequence"].tolist()
	labels = meta["label"].tolist()
	splits = meta["split"].tolist()

	print(f"Total sequences: {len(sequences)}")
	print("Split counts:", pd.Series(splits).value_counts().to_dict())

	print("Computing ESM embeddings...")
	embeddings = compute_embeddings(sequences) # shape (N, H)

	full_ds = Dataset.from_dict({
	"sequence": sequences,
	"embedding": embeddings,
	"label": labels,
	"split": splits,
	})

	# Split into train / val, then drop the 'split' column
	train_ds = full_ds.filter(lambda x: x["split"] == "train")
	val_ds = full_ds.filter(lambda x: x["split"] == "val")

	train_ds = train_ds.remove_columns("split")
	val_ds = val_ds.remove_columns("split")

	ds_dict = DatasetDict({
	"train": train_ds,
	"val": val_ds,
	})

	ds_dict.save_to_disk(save_path)
	print(f"Saved DatasetDict with train/val to: {save_path}")
	print("Train size:", len(ds_dict["train"]))
	print("Val size:", len(ds_dict["val"]))

	return ds_dict


	ds = create_and_save_datasets()

	ex = ds["train"][0]
	print("\nExample from train:")
	print("Sequence:", ex["sequence"])
	print("Embedding shape:", np.array(ex["embedding"]).shape)
	print("Label:", ex["label"])

	torch.cuda.empty_cache()

	meta_path = "./Classifier_Weight/training_data_cleaned/nf/nf_meta_with_split.csv"
	save_path = "./Classifier_Weight/training_data_cleaned/nf/nf_wt_with_embeddings_unpooled"

	tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t33_650M_UR50D")
	model = EsmModel.from_pretrained("facebook/esm2_t33_650M_UR50D", add_pooling_layer=False).to(device).eval()

	cls_id = tokenizer.cls_token_id
	eos_id = tokenizer.eos_token_id

	@torch.no_grad()
	def embed_one(seq, max_length=1022):
	inputs = tokenizer(seq, padding=False, truncation=True, max_length=max_length, return_tensors="pt")
	inputs = {k: v.to(device) for k, v in inputs.items()}
	out = model(**inputs)
	h = out.last_hidden_state[0] # (L, H)
	attn = inputs["attention_mask"][0].bool() # (L,)
	ids = inputs["input_ids"][0]

	keep = attn.clone()
	if cls_id is not None:
	keep &= (ids != cls_id)
	if eos_id is not None:
	keep &= (ids != eos_id)

	hb = h[keep].detach().cpu().to(torch.float16).numpy() # (Li, H)
	return hb

	H = 1280
	features = Features({
	"sequence": Value("string"),
	"label": Value("int64"),
	"embedding": Sequence(Sequence(Value("float16"), length=H)), # (Li, H) as nested lists
	"attention_mask": Sequence(Value("int8")), # (Li,)
	"length": Value("int64"),
	})

	def make_generator(df):
	for seq, lab in tqdm.tqdm(zip(df["sequence"].tolist(), df["label"].astype(int).tolist()), total=len(df)):
	emb = embed_one(seq) # (Li, H) float16
	emb_list = emb.tolist()
	li = len(emb_list)
	yield {
	"sequence": seq,
	"label": int(lab),
	"embedding": emb_list,
	"attention_mask": [1] * li,
	"length": li,
	}

	def build_and_save_split(df, out_dir):
	ds = Dataset.from_generator(make_generator, gen_kwargs={"df": df}, features=features)
	# small batches when writing prevents the 2GB offset overflow
	ds.save_to_disk(out_dir, max_shard_size="1GB")
	return ds

	meta = pd.read_csv(meta_path)
	train_df = meta[meta["split"] == "train"].reset_index(drop=True)
	val_df = meta[meta["split"] == "val"].reset_index(drop=True)

	train_dir = os.path.join(save_path, "train")
	val_dir = os.path.join(save_path, "val")
	os.makedirs(save_path, exist_ok=True)

	train_ds = build_and_save_split(train_df, train_dir)
	val_ds = build_and_save_split(val_df, val_dir)

	ds_dict = DatasetDict({"train": train_ds, "val": val_ds})
	ds_dict.save_to_disk(save_path)
	print(ds_dict)