text_plain_output_1.png
from esm.model.esm2 import ESM2

from multiprocess import Pool

from tqdm import tqdm

from tqdm import tqdm

from tqdm import tqdm

import esm

import gc

import gc

import numpy as np

import numpy as np

import os

import os

import os

import pandas as pd

import pandas as pd

import re

import tensorflow as tf

import torch

import torch.distributed as dist

import warnings

import pandas as pd
import numpy as np
import plotly.express as px
from plotly.subplots import make_subplots
import gc
from tqdm import tqdm
import pickle
import os
import warnings
warnings.filterwarnings('ignore')
train_clinical_data = pd.read_csv('../input/amp-parkinsons-disease-progression-prediction/train_clinical_data.csv')
train_peptides = pd.read_csv('../input/amp-parkinsons-disease-progression-prediction/train_peptides.csv')
train_proteins = pd.read_csv('../input/amp-parkinsons-disease-progression-prediction/train_proteins.csv')
supplemental_clinical_data = pd.read_csv('../input/amp-parkinsons-disease-progression-prediction/supplemental_clinical_data.csv')

peptideIndexLabels = ['visit_id', 'patient_id', 'UniProt']
train_peptides.set_index(peptideIndexLabels, drop=True)
peptideData = train_peptides.merge(train_proteins, on=['visit_id', 'visit_month', 'patient_id', 'UniProt'], how='inner').drop_duplicates()

peptides_per_patient = peptideData.groupby('patient_id')['Peptide'].nunique()

import numpy as np
import pandas as pd
from tqdm import tqdm
from multiprocess import Pool
import os
import tensorflow as tf
peptideData['Peptide oxidation score'] = 0
peptideData['Protein oxidation score'] = 0
peptideData['Peptide carbamidomethylation score'] = 0
peptideData['Protein carbamidomethylation score'] = 0
def preprocessPeptides(args):
    (patientID, proteinID), currentPeptides = args
    peptideCount = len(currentPeptides)
    if peptideCount == 0:
        return currentPeptides
    def deserializeUniMod(peptide, scoreAndRemove=['UniMod_35', 'UniMod_4']):
        oxidationScore, carbamidomethylationScore = (0, 0)
        for string in scoreAndRemove:
            while string in peptide:
                peptide = peptide.replace(f'({string})', '', 1)
                if string == 'UniMod_35':
                    oxidationScore += np.log(len(peptide))
                elif string == 'UniMod_4':
                    carbamidomethylationScore += np.log(len(peptide))
        return pd.Series([peptide, oxidationScore, carbamidomethylationScore])
    currentPeptides[['Peptide', 'Peptide oxidation score', 'Peptide carbamidomethylation score']] = currentPeptides['Peptide'].apply(deserializeUniMod)
    normalized_carbamidomethylation_score = currentPeptides['Peptide carbamidomethylation score'].sum() / peptideCount
    normalized_oxidation_score = currentPeptides['Peptide oxidation score'].sum() / peptideCount
    currentPeptides['Protein oxidation score'] = normalized_oxidation_score
    currentPeptides['Protein carbamidomethylation score'] = normalized_carbamidomethylation_score
    return currentPeptides
grouped_peptides = peptideData.groupby(['patient_id', 'UniProt'])
with Pool(os.cpu_count()) as pool:
    preprocessed_peptides = pool.map(preprocessPeptides, [(group, data) for group, data in grouped_peptides])
peptideData = pd.concat(preprocessed_peptides).reset_index(drop=True)

len([(row[0], row[0]) for row in peptideData.groupby(['Peptide'])])

import multiprocess
import gc
import os
import torch
import torch.distributed as dist
import esm
from esm.model.esm2 import ESM2
from tqdm import tqdm
import re
def upgrade_state_dict(state_dict):
    """Removes prefixes 'model.encoder.sentence_encoder.' and 'model.encoder.'."""
    prefixes = ['encoder.sentence_encoder.', 'encoder.']
    pattern = re.compile('^' + '|'.join(prefixes))
    state_dict = {pattern.sub('', name): param for name, param in state_dict.items()}
    return state_dict
model_path = '/kaggle/input/fair-esm/esm2_t33_650M_UR50D.pt'
token_layer = 33
url = 'tcp://localhost:23456'
dist.init_process_group(backend='gloo', init_method=url, world_size=1, rank=0)
model_data = torch.load(str(model_path), map_location='cpu')
cfg = model_data['cfg']['model']
state_dict = model_data['model']
state_dict = upgrade_state_dict(state_dict)
vocab = esm.data.Alphabet.from_architecture('ESM-1b')
model = ESM2(num_layers=cfg.encoder_layers, embed_dim=cfg.encoder_embed_dim, attention_heads=cfg.encoder_attention_heads, alphabet=vocab, token_dropout=cfg.token_dropout)
batch_converter = vocab.get_batch_converter()
data = [(row[0], row[0]) for row in peptideData.groupby(['Peptide'])]
batch_labels, batch_strs, batch_tokens = batch_converter(data)
batch_lens = (batch_tokens != vocab.padding_idx).sum(1)
if torch.cuda.is_available():
    model.cuda().eval()
    batch_tokens = batch_tokens.cuda()
else:
    model.eval()
with torch.no_grad():
    results = model(batch_tokens, repr_layers=[token_layer], return_contacts=False)
token_representations = results['representations'][token_layer].to('cpu')
if torch.cuda.is_available():
    torch.cuda.empty_cache()
gc.collect()
dist.destroy_process_group()
sequence_representations = {}
for (i, tokens_len), peptide, label in tqdm(zip(enumerate(batch_lens), token_representations, data), unit='peptide', total=len(batch_lens)):
    sequence_representations[label[0]] = tf.convert_to_tensor(token_representations[i, 1:tokens_len - 1].numpy())
del token_representations
del batch_lens
del model
del batch_strs
del batch_converter
del batch_tokens
del results
del state_dict
del cfg
del model_data
del data
gc.collect()

indexLabels = ['visit_id', 'visit_month', 'patient_id']
clinicalData = train_clinical_data
allData = peptideData.merge(clinicalData, on=indexLabels)
peptidesByRowIndex = allData['Peptide'].copy()

allData