Delete utils

utils/data_loader.py

@@ -1,43 +0,0 @@
-import pandas as pd
-import torch
-from torch.utils.data import Dataset, DataLoader
-from torch.nn.utils.rnn import pad_sequence
-from esm_utils import get_latents, load_esm2_model
-import config
-
-class ProteinDataset(Dataset):
-    def __init__(self, csv_file, tokenizer, model):
-        self.data = pd.read_csv(csv_file).head(4)
-        self.tokenizer = tokenizer
-        self.model = model
-
-    def __len__(self):
-        return len(self.data)
-
-    def __getitem__(self, idx):
-        sequence = self.data.iloc[idx]['Sequence']
-        latents = get_latents(self.model, self.tokenizer, sequence)
-
-        attention_mask = torch.ones_like(latents)
-        attention_mask = torch.mean(attention_mask, dim=-1)
-
-        return latents, attention_mask
-
-def collate_fn(batch):
-    latents, attention_mask = zip(*batch)
-    latents_padded = pad_sequence([torch.tensor(latent) for latent in latents], batch_first=True, padding_value=0)
-    attention_mask_padded = pad_sequence([torch.tensor(mask) for mask in attention_mask], batch_first=True, padding_value=0)
-    return latents_padded, attention_mask_padded
-
-def get_dataloaders(config):
-    tokenizer, masked_model, embedding_model = load_esm2_model(config.MODEL_NAME)
-    
-    train_dataset = ProteinDataset(config.Loader.DATA_PATH + "/train.csv", tokenizer, embedding_model)
-    val_dataset = ProteinDataset(config.Loader.DATA_PATH + "/val.csv", tokenizer, embedding_model)
-    test_dataset = ProteinDataset(config.Loader.DATA_PATH + "/test.csv", tokenizer, embedding_model)
-    
-    train_loader = DataLoader(train_dataset, batch_size=config.Loader.BATCH_SIZE, num_workers=0, shuffle=True, collate_fn=collate_fn)
-    val_loader = DataLoader(val_dataset, batch_size=config.Loader.BATCH_SIZE, num_workers=0, shuffle=False, collate_fn=collate_fn)
-    test_loader = DataLoader(test_dataset, batch_size=config.Loader.BATCH_SIZE, num_workers=0, shuffle=False, collate_fn=collate_fn)
-    
-    return train_loader, val_loader, test_loader

utils/esm_utils.py

@@ -1,15 +0,0 @@
-import torch
-import config
-from transformers import AutoTokenizer, AutoModel, AutoModelForMaskedLM
-
-def load_esm2_model(model_name):
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    masked_model = AutoModelForMaskedLM.from_pretrained(model_name)
-    embedding_model = AutoModel.from_pretrained(model_name)
-    return tokenizer, masked_model, embedding_model
-
-def get_latents(model, tokenizer, sequence):
-    inputs = tokenizer(sequence, return_tensors="pt")
-    with torch.no_grad():
-        outputs = model(**inputs)
-    return outputs.last_hidden_state.squeeze(0)