File size: 2,681 Bytes
9f16c4c | 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 | import pandas as pd
import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader
from sklearn.model_selection import train_test_split
from collections import Counter
AA_VOCAB = {aa: i+1 for i, aa in enumerate('ARNDCQEGHILKMFPSTWYV')}
AA_VOCAB['PAD'] = 0
AA_VOCAB['X'] = 20
def encode_sequence(seq, max_len=200):
ids = [AA_VOCAB.get(c.upper(), AA_VOCAB['X']) for c in seq[:max_len]]
ids += [AA_VOCAB['PAD']] * (max_len - len(ids))
return ids
class PeptideDataset(Dataset):
def __init__(self, sequences, labels, max_len=200):
self.sequences = sequences
self.labels = labels
self.max_len = max_len
def __len__(self):
return len(self.sequences)
def __getitem__(self, idx):
seq = self.sequences[idx]
label = self.labels[idx]
ids = encode_sequence(seq, self.max_len)
return torch.tensor(ids, dtype=torch.long), torch.tensor(label, dtype=torch.long)
def load_genpept_data(csv_path='GenPept-Curated-2025/data/balanced_11000.csv'):
df = pd.read_csv(csv_path)
sequences = df['sequence'].values
labels = (df['label'].values == 'AMP').astype(np.int64)
return sequences, labels
def create_splits(sequences, labels, test_size=0.21, val_size=0.09, random_state=42):
X_temp, X_test, y_temp, y_test = train_test_split(
sequences, labels, test_size=test_size,
stratify=labels, random_state=random_state
)
val_ratio = val_size / (1 - test_size)
X_train, X_val, y_train, y_val = train_test_split(
X_temp, y_temp, test_size=val_ratio,
stratify=y_temp, random_state=random_state
)
return (X_train, y_train), (X_val, y_val), (X_test, y_test)
def get_dataloaders(sequences, labels, batch_size=64, max_len=200, num_workers=2):
(X_train, y_train), (X_val, y_val), (X_test, y_test) = create_splits(sequences, labels)
train_ds = PeptideDataset(X_train, y_train, max_len)
val_ds = PeptideDataset(X_val, y_val, max_len)
test_ds = PeptideDataset(X_test, y_test, max_len)
train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True, num_workers=num_workers)
val_loader = DataLoader(val_ds, batch_size=batch_size, shuffle=False, num_workers=num_workers)
test_loader = DataLoader(test_ds, batch_size=batch_size, shuffle=False, num_workers=num_workers)
return train_loader, val_loader, test_loader
if __name__ == '__main__':
seqs, labs = load_genpept_data()
print(f'Loaded {len(seqs)} sequences, {Counter(labs)}')
train_l, val_l, test_l = get_dataloaders(seqs, labs, batch_size=4)
for x, y in train_l:
print(f'Batch: x={x.shape}, y={y.shape}')
break
|