{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":51294,"databundleVersionId":7331882,"sourceType":"competition"},{"sourceId":6822004,"sourceType":"datasetVersion","datasetId":3719560}],"dockerImageVersionId":30559,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Description\nWelcome to Stanford Ribonanza RNA Folding challenge. The task of this competition is predicting the chemical reactivity at each position of an RNA molecule. These data are extremely sensitive to the structure that each RNA forms, and an algorithm that could perfectly predict these chemical reactivities would need to have an implicit ‘understanding’ of RNA structure. Such an oracle could be then utilized to predictively model structures of novel RNA molecules. A better understanding of how to manipulate RNA could help usher in an age of programmable medicine, including first cures for pancreatic cancer and Alzheimer’s disease as well as much-needed antibiotics and new biotechnology approaches for climate change. \n\nThis notebook provides a simple baseline that may be used as a starting point for further experiments. Improvment of the baseline may include: \n1. Use of proper loss function to incorporate SN_filter = 0 samples as well as reactivity errors into training\n2. Model improvement and use of additional data, e.g. Ribonanza_bpp_files\n\nFinally, working on this competition keep in mind that train/public LB have different sequence length distribution from private LB, i.e. 115-206 vs. 207-457. Therefore, to avoid a strong shakeup one may need to look into performance vs. sequence length end ensure generalizability.","metadata":{}},{"cell_type":"markdown","source":"スタンフォード リボナンザ RNA フォールディング チャレンジへようこそ。このコンテストの課題は、RNA 分子の各位置での化学反応性を予測することです。これらのデータは、各 RNA が形成する構造に非常に敏感であり、これらの化学反応性を完全に予測できるアルゴリズムには、RNA 構造の暗黙の「理解」が必要です。このようなオラクルは、新しい RNA 分子の構造を予測的にモデル化するために使用できます。RNA の操作方法をより深く理解することで、膵臓がんやアルツハイマー病の最初の治療法、待望の抗生物質、気候変動に対する新しいバイオテクノロジーのアプローチなど、プログラム可能な医療の時代が到来する可能性があります。\n\nこのノートブックは、さらなる実験の出発点として使用できるシンプルなベースラインを提供します。ベースラインの改善には、次のことが含まれます。\n\nSN_filter = 0 のサンプルと反応性エラーをトレーニングに組み込むための適切な損失関数の使用\nモデルの改善と追加データの使用、例: Ribonanza_bpp_files\n最後に、このコンペティションに取り組む際は、トレーニング/パブリック LB のシーケンス長の分布がプライベート LB とは異なる (115 ～ 206 対 207 ～ 457) ことに留意してください。したがって、大きな変動を避けるには、パフォーマンスとシーケンス長を比較して一般化を保証する必要があるかもしれません。","metadata":{}},{"cell_type":"code","source":"import pandas as pd\nimport os, gc\nimport numpy as np\nfrom sklearn.model_selection import KFold\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader","metadata":{"_kg_hide-output":true,"_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-06-25T05:10:43.278545Z","iopub.execute_input":"2024-06-25T05:10:43.278818Z","iopub.status.idle":"2024-06-25T05:10:47.603211Z","shell.execute_reply.started":"2024-06-25T05:10:43.278794Z","shell.execute_reply":"2024-06-25T05:10:47.602339Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Fix fastai bug to enable fp16 training with dictionaries\n\nimport torch\nfrom fastai.vision.all import *\ndef flatten(o):\n    \"Concatenate all collections and items as a generator\"\n    for item in o:\n        if isinstance(o, dict): yield o[item]; continue\n        elif isinstance(item, str): yield item; continue\n        try: yield from flatten(item)\n        except TypeError: yield item\n\nfrom torch.cuda.amp import GradScaler, autocast\n@delegates(GradScaler)\nclass MixedPrecision(Callback):\n    \"Mixed precision training using Pytorch's `autocast` and `GradScaler`\"\n    order = 10\n    def __init__(self, **kwargs): self.kwargs = kwargs\n    def before_fit(self): \n        self.autocast,self.learn.scaler,self.scales = autocast(),GradScaler(**self.kwargs),L()\n    def before_batch(self): self.autocast.__enter__()\n    def after_pred(self):\n        if next(flatten(self.pred)).dtype==torch.float16: self.learn.pred = to_float(self.pred)\n    def after_loss(self): self.autocast.__exit__(None, None, None)\n    def before_backward(self): self.learn.loss_grad = self.scaler.scale(self.loss_grad)\n    def before_step(self):\n        \"Use `self` as a fake optimizer. `self.skipped` will be set to True `after_step` if gradients overflow. \"\n        self.skipped=True\n        self.scaler.step(self)\n        if self.skipped: raise CancelStepException()\n        self.scales.append(self.scaler.get_scale())\n    def after_step(self): self.learn.scaler.update()\n\n    @property \n    def param_groups(self): \n        \"Pretend to be an optimizer for `GradScaler`\"\n        return self.opt.param_groups\n    def step(self, *args, **kwargs): \n        \"Fake optimizer step to detect whether this batch was skipped from `GradScaler`\"\n        self.skipped=False\n    def after_fit(self): self.autocast,self.learn.scaler,self.scales = None,None,None\n        \nimport fastai\nfastai.callback.fp16.MixedPrecision = MixedPrecision\n","metadata":{"_kg_hide-output":true,"_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-06-25T05:10:47.604929Z","iopub.execute_input":"2024-06-25T05:10:47.605337Z","iopub.status.idle":"2024-06-25T05:10:49.579441Z","shell.execute_reply.started":"2024-06-25T05:10:47.605311Z","shell.execute_reply":"2024-06-25T05:10:49.578644Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def seed_everything(seed):\n    random.seed(seed)\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = True","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-06-25T05:10:49.580755Z","iopub.execute_input":"2024-06-25T05:10:49.581134Z","iopub.status.idle":"2024-06-25T05:10:49.586810Z","shell.execute_reply.started":"2024-06-25T05:10:49.581101Z","shell.execute_reply":"2024-06-25T05:10:49.585881Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"fname = 'example0'\nPATH = '/kaggle/input/stanford-ribonanza-rna-folding-converted/'\nOUT = './'\nbs = 256\nnum_workers = 2\nSEED = 2023\nnfolds = 4\ndevice = 'cuda' if torch.cuda.is_available() else 'cpu'","metadata":{"execution":{"iopub.status.busy":"2024-06-25T05:10:49.588039Z","iopub.execute_input":"2024-06-25T05:10:49.588344Z","iopub.status.idle":"2024-06-25T05:10:49.623015Z","shell.execute_reply.started":"2024-06-25T05:10:49.588320Z","shell.execute_reply":"2024-06-25T05:10:49.622241Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Data\n\nThe primary training data is provided in train_data.csv, which contains 821840 RNA sequences and the corresponding reactivity measurements with 2A3_MaP and DMS_MaP methods. The reactivity is reported in columns reactivity_0001 - reactivity_0206 and is set to NaN for the first 26 and the last 21 nucleotides as well as padding for sequences shorter than 206. For faster loading and effective RAM use, I converted the data into a float32 parquet file. \n\nEvaluation in this competition is performed only on samples with SN_filter = 1 for both measurement methods. In this example, I perform training only on samples wiht SN_filter = 1, which gives a noticeable CV boost but uses only 1/4 of the data (i.e. training on noisy SN_filter = 0 data degrades the performance). A proper consideration of all data as well as reactivity errors may boost the performance.\n\nIn this example, I use a simple CV Kfold split. However, given a mismatch in the RNA length between train/public LB vs. private LB data, **it may be important to verify the effect of the sequence length** to avoid a significant shakeup at the private LB.\n\nOne of the tricks, well known in NLP community, which I use here, is length matching batch sampling: composing batches of samples of approximately the same length to minimize the overhead caused by padding tokens.","metadata":{}},{"cell_type":"markdown","source":"データ\n主なトレーニング データは train_data.csv で提供され、821840 個の RNA シーケンスと、2A3_MaP および DMS_MaP メソッドによる対応する反応性測定が含まれています。反応性は列 reactivity_0001 - reactivity_0206 で報告され、最初の 26 ヌクレオチドと最後の 21 ヌクレオチドは NaN に設定され、206 より短いシーケンスはパディングされます。読み込みを高速化し、RAM を効率的に使用するために、データを float32 parquet ファイルに変換しました。\n\nこのコンペティションでの評価は、両方の測定方法の SN_filter = 1 のサンプルでのみ実行されます。この例では、SN_filter = 1 のサンプルでのみトレーニングを実行します。これにより、顕著な CV ブーストが得られますが、使用されるデータは 1/4 のみです (つまり、ノイズの多い SN_filter = 0 データでトレーニングすると、パフォーマンスが低下します)。すべてのデータと反応性エラーを適切に考慮すると、パフォーマンスが向上する可能性があります。\n\nこの例では、単純な CV Kfold 分割を使用しています。ただし、トレーニング/パブリック LB とプライベート LB データ間の RNA の長さが一致しない場合は、プライベート LB での大きな混乱を避けるために、シーケンスの長さの影響を確認することが重要になる場合があります。\n\nここで使用している、NLP コミュニティでよく知られているトリックの 1 つは、長さを一致させるバッチ サンプリングです。これは、ほぼ同じ長さのサンプルのバッチを構成して、パディング トークンによって発生するオーバーヘッドを最小限に抑えるものです。","metadata":{}},{"cell_type":"code","source":"class RNA_Dataset(Dataset):\n    def __init__(self, df, mode='train', seed=2023, fold=0, nfolds=4, \n                 mask_only=False, **kwargs):\n        self.seq_map = {'A':0,'C':1,'G':2,'U':3}\n        self.Lmax = 206\n        df['L'] = df.sequence.apply(len)\n        df_2A3 = df.loc[df.experiment_type=='2A3_MaP']\n        df_DMS = df.loc[df.experiment_type=='DMS_MaP']\n        \n        split = list(KFold(n_splits=nfolds, random_state=seed, \n                shuffle=True).split(df_2A3))[fold][0 if mode=='train' else 1]\n        df_2A3 = df_2A3.iloc[split].reset_index(drop=True)\n        df_DMS = df_DMS.iloc[split].reset_index(drop=True)\n        \n        m = (df_2A3['SN_filter'].values > 0) & (df_DMS['SN_filter'].values > 0)\n        df_2A3 = df_2A3.loc[m].reset_index(drop=True)\n        df_DMS = df_DMS.loc[m].reset_index(drop=True)\n        \n        self.seq = df_2A3['sequence'].values\n        self.L = df_2A3['L'].values\n        \n        self.react_2A3 = df_2A3[[c for c in df_2A3.columns if \\\n                                 'reactivity_0' in c]].values\n        self.react_DMS = df_DMS[[c for c in df_DMS.columns if \\\n                                 'reactivity_0' in c]].values\n        self.react_err_2A3 = df_2A3[[c for c in df_2A3.columns if \\\n                                 'reactivity_error_0' in c]].values\n        self.react_err_DMS = df_DMS[[c for c in df_DMS.columns if \\\n                                'reactivity_error_0' in c]].values\n        self.sn_2A3 = df_2A3['signal_to_noise'].values\n        self.sn_DMS = df_DMS['signal_to_noise'].values\n        self.mask_only = mask_only\n        \n    def __len__(self):\n        return len(self.seq)  \n    \n    def __getitem__(self, idx):\n        seq = self.seq[idx]\n        if self.mask_only:\n            mask = torch.zeros(self.Lmax, dtype=torch.bool)\n            mask[:len(seq)] = True\n            return {'mask':mask},{'mask':mask}\n        seq = [self.seq_map[s] for s in seq]\n        seq = np.array(seq)\n        mask = torch.zeros(self.Lmax, dtype=torch.bool)\n        mask[:len(seq)] = True\n        seq = np.pad(seq,(0,self.Lmax-len(seq)))\n        \n        react = torch.from_numpy(np.stack([self.react_2A3[idx],\n                                           self.react_DMS[idx]],-1))\n        react_err = torch.from_numpy(np.stack([self.react_err_2A3[idx],\n                                               self.react_err_DMS[idx]],-1))\n        sn = torch.FloatTensor([self.sn_2A3[idx],self.sn_DMS[idx]])\n        \n        return {'seq':torch.from_numpy(seq), 'mask':mask}, \\\n               {'react':react, 'react_err':react_err,\n                'sn':sn, 'mask':mask}\n    \nclass LenMatchBatchSampler(torch.utils.data.BatchSampler):\n    def __iter__(self):\n        buckets = [[]] * 100\n        yielded = 0\n\n        for idx in self.sampler:\n            s = self.sampler.data_source[idx]\n            if isinstance(s,tuple): L = s[0][\"mask\"].sum()\n            else: L = s[\"mask\"].sum()\n            L = max(1,L // 16) \n            if len(buckets[L]) == 0:  buckets[L] = []\n            buckets[L].append(idx)\n            \n            if len(buckets[L]) == self.batch_size:\n                batch = list(buckets[L])\n                yield batch\n                yielded += 1\n                buckets[L] = []\n                \n        batch = []\n        leftover = [idx for bucket in buckets for idx in bucket]\n\n        for idx in leftover:\n            batch.append(idx)\n            if len(batch) == self.batch_size:\n                yielded += 1\n                yield batch\n                batch = []\n\n        if len(batch) > 0 and not self.drop_last:\n            yielded += 1\n            yield batch\n            \ndef dict_to(x, device='cuda'):\n    return {k:x[k].to(device) for k in x}\n\ndef to_device(x, device='cuda'):\n    return tuple(dict_to(e,device) for e in x)\n\nclass DeviceDataLoader:\n    def __init__(self, dataloader, device='cuda'):\n        self.dataloader = dataloader\n        self.device = device\n    \n    def __len__(self):\n        return len(self.dataloader)\n    \n    def __iter__(self):\n        for batch in self.dataloader:\n            yield tuple(dict_to(x, self.device) for x in batch)","metadata":{"execution":{"iopub.status.busy":"2024-06-25T05:10:49.624967Z","iopub.execute_input":"2024-06-25T05:10:49.625236Z","iopub.status.idle":"2024-06-25T05:10:49.659232Z","shell.execute_reply.started":"2024-06-25T05:10:49.625213Z","shell.execute_reply":"2024-06-25T05:10:49.658385Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model\nTransformers are ideal for the considered task because they naturally capture long dependences in RNN, which define the secondary structure of the molecule and the corresponding chemical reactivity. For illustration purposes, below I provide a simple S-size transformer model.","metadata":{}},{"cell_type":"markdown","source":"モデル\nトランスフォーマーは、分子の二次構造とそれに対応する化学反応性を定義する RNN の長い依存関係を自然にキャプチャするため、検討中のタスクに最適です。説明のために、以下にシンプルな S サイズのトランスフォーマー モデルを示します。","metadata":{}},{"cell_type":"code","source":"class SinusoidalPosEmb(nn.Module):\n    def __init__(self, dim=16, M=10000):\n        super().__init__()\n        self.dim = dim\n        self.M = M\n\n    def forward(self, x):\n        device = x.device\n        half_dim = self.dim // 2\n        emb = math.log(self.M) / half_dim\n        emb = torch.exp(torch.arange(half_dim, device=device) * (-emb))\n        emb = x[...,None] * emb[None,...]\n        emb = torch.cat((emb.sin(), emb.cos()), dim=-1)\n        return emb\n\nclass RNA_Model(nn.Module):\n    def __init__(self, dim=192, depth=15, head_size=32, **kwargs):\n        super().__init__()\n        self.emb = nn.Embedding(4,dim)\n        self.pos_enc = SinusoidalPosEmb(dim)\n        self.transformer = nn.TransformerEncoder(\n            nn.TransformerEncoderLayer(d_model=dim, nhead=dim//head_size, dim_feedforward=4*dim,\n                dropout=0.1, activation=nn.GELU(), batch_first=True, norm_first=True), depth)\n        self.proj_out = nn.Linear(dim,2)\n    \n    def forward(self, x0):\n        mask = x0['mask']\n        Lmax = mask.sum(-1).max()\n        mask = mask[:,:Lmax]\n        x = x0['seq'][:,:Lmax]\n        \n        pos = torch.arange(Lmax, device=x.device).unsqueeze(0)\n        pos = self.pos_enc(pos)\n        x = self.emb(x)\n        x = x + pos\n        \n        x = self.transformer(x, src_key_padding_mask=~mask)\n        x = self.proj_out(x)\n        \n        return x","metadata":{"execution":{"iopub.status.busy":"2024-06-25T05:10:49.660460Z","iopub.execute_input":"2024-06-25T05:10:49.660953Z","iopub.status.idle":"2024-06-25T05:10:49.673922Z","shell.execute_reply.started":"2024-06-25T05:10:49.660923Z","shell.execute_reply":"2024-06-25T05:10:49.673187Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Loss & Metric\nThe metric accumulates all predictions and then performs average to be consistent with the competition metric. However, the difference with a simple batch-based average is negligible.","metadata":{}},{"cell_type":"markdown","source":"","metadata":{}},{"cell_type":"markdown","source":"損失とメトリック\nメトリックはすべての予測を蓄積し、競合メトリックと一致するように平均を実行します。ただし、単純なバッチベースの平均との差はごくわずかです。","metadata":{}},{"cell_type":"code","source":"def loss(pred,target):\n    p = pred[target['mask'][:,:pred.shape[1]]]\n    y = target['react'][target['mask']].clip(0,1)\n    loss = F.l1_loss(p, y, reduction='none')\n    loss = loss[~torch.isnan(loss)].mean()\n    \n    return loss\n\nclass MAE(Metric):\n    def __init__(self): \n        self.reset()\n        \n    def reset(self): \n        self.x,self.y = [],[]\n        \n    def accumulate(self, learn):\n        x = learn.pred[learn.y['mask'][:,:learn.pred.shape[1]]]\n        y = learn.y['react'][learn.y['mask']].clip(0,1)\n        self.x.append(x)\n        self.y.append(y)\n\n    @property\n    def value(self):\n        x,y = torch.cat(self.x,0),torch.cat(self.y,0)\n        loss = F.l1_loss(x, y, reduction='none')\n        loss = loss[~torch.isnan(loss)].mean()\n        return loss","metadata":{"execution":{"iopub.status.busy":"2024-06-25T05:10:49.675085Z","iopub.execute_input":"2024-06-25T05:10:49.675367Z","iopub.status.idle":"2024-06-25T05:10:49.687960Z","shell.execute_reply.started":"2024-06-25T05:10:49.675345Z","shell.execute_reply":"2024-06-25T05:10:49.687101Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Training","metadata":{}},{"cell_type":"code","source":"seed_everything(SEED)\nos.makedirs(OUT, exist_ok=True)\ndf = pd.read_parquet(os.path.join(PATH,'train_data.parquet'))\n\nfor fold in [0]: # running multiple folds at kaggle may cause OOM\n    ds_train = RNA_Dataset(df, mode='train', fold=fold, nfolds=nfolds)\n    ds_train_len = RNA_Dataset(df, mode='train', fold=fold, \n                nfolds=nfolds, mask_only=True)\n    sampler_train = torch.utils.data.RandomSampler(ds_train_len)\n    len_sampler_train = LenMatchBatchSampler(sampler_train, batch_size=bs,\n                drop_last=True)\n    dl_train = DeviceDataLoader(torch.utils.data.DataLoader(ds_train, \n                batch_sampler=len_sampler_train, num_workers=num_workers,\n                persistent_workers=True), device)\n\n    ds_val = RNA_Dataset(df, mode='eval', fold=fold, nfolds=nfolds)\n    ds_val_len = RNA_Dataset(df, mode='eval', fold=fold, nfolds=nfolds, \n               mask_only=True)\n    sampler_val = torch.utils.data.SequentialSampler(ds_val_len)\n    len_sampler_val = LenMatchBatchSampler(sampler_val, batch_size=bs, \n               drop_last=False)\n    dl_val= DeviceDataLoader(torch.utils.data.DataLoader(ds_val, \n               batch_sampler=len_sampler_val, num_workers=num_workers), device)\n    gc.collect()\n\n    data = DataLoaders(dl_train,dl_val)\n    model = RNA_Model()   \n    model = model.to(device)\n    learn = Learner(data, model, loss_func=loss,cbs=[GradientClip(3.0)],\n                metrics=[MAE()]).to_fp16() \n    #fp16 doesn't help at P100 but gives x1.6-1.8 speedup at modern hardware\n\n    learn.fit_one_cycle(32, lr_max=5e-4, wd=0.05, pct_start=0.02)\n    torch.save(learn.model.state_dict(),os.path.join(OUT,f'{fname}_{fold}.pth'))\n    gc.collect()","metadata":{"scrolled":true,"execution":{"iopub.status.busy":"2024-06-25T05:10:49.688972Z","iopub.execute_input":"2024-06-25T05:10:49.689213Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}