{"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"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"It is an example of inference with model trained in [RNA starter kernel](https://www.kaggle.com/code/iafoss/rna-starter)","metadata":{}},{"cell_type":"code","source":"import pandas as pd\nimport os, gc\nimport numpy as np\nfrom tqdm.notebook import tqdm\nimport math\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,"execution":{"iopub.status.busy":"2023-09-11T02:45:53.339696Z","iopub.execute_input":"2023-09-11T02:45:53.339953Z","iopub.status.idle":"2023-09-11T02:45:58.394952Z","shell.execute_reply.started":"2023-09-11T02:45:53.339928Z","shell.execute_reply":"2023-09-11T02:45:58.393964Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"MODELS = ['/kaggle/input/rna-starter/example0_0.pth']\nPATH = '/kaggle/input/stanford-ribonanza-rna-folding-converted/'\nbs = 256\nnum_workers = 2\ndevice = 'cuda' if torch.cuda.is_available() else 'cpu'","metadata":{"execution":{"iopub.status.busy":"2023-09-11T02:45:58.397109Z","iopub.execute_input":"2023-09-11T02:45:58.397627Z","iopub.status.idle":"2023-09-11T02:45:58.427500Z","shell.execute_reply.started":"2023-09-11T02:45:58.397586Z","shell.execute_reply":"2023-09-11T02:45:58.425474Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Data","metadata":{}},{"cell_type":"code","source":"class RNA_Dataset_Test(Dataset):\n    def __init__(self, df, mask_only=False, **kwargs):\n        self.seq_map = {'A':0,'C':1,'G':2,'U':3}\n        df['L'] = df.sequence.apply(len)\n        self.Lmax = df['L'].max()\n        self.df = df\n        self.mask_only = mask_only\n        \n    def __len__(self):\n        return len(self.df)  \n    \n    def __getitem__(self, idx):\n        id_min, id_max, seq = self.df.loc[idx, ['id_min','id_max','sequence']]\n        mask = torch.zeros(self.Lmax, dtype=torch.bool)\n        L = len(seq)\n        mask[:L] = True\n        if self.mask_only: return {'mask':mask},{}\n        ids = np.arange(id_min,id_max+1)\n        \n        seq = [self.seq_map[s] for s in seq]\n        seq = np.array(seq)\n        seq = np.pad(seq,(0,self.Lmax-L))\n        ids = np.pad(ids,(0,self.Lmax-L), constant_values=-1)\n        \n        return {'seq':torch.from_numpy(seq), 'mask':mask}, \\\n               {'ids':ids}\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":"2023-09-11T02:45:58.428972Z","iopub.execute_input":"2023-09-11T02:45:58.429790Z","iopub.status.idle":"2023-09-11T02:45:58.448952Z","shell.execute_reply.started":"2023-09-11T02:45:58.429756Z","shell.execute_reply":"2023-09-11T02:45:58.448144Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model","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=12, 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":"2023-09-11T02:45:58.452172Z","iopub.execute_input":"2023-09-11T02:45:58.452491Z","iopub.status.idle":"2023-09-11T02:45:58.466710Z","shell.execute_reply.started":"2023-09-11T02:45:58.452465Z","shell.execute_reply":"2023-09-11T02:45:58.465807Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Inference","metadata":{}},{"cell_type":"code","source":"df_test = pd.read_parquet(os.path.join(PATH,'test_sequences.parquet'))\nds = RNA_Dataset_Test(df_test)\ndl = DeviceDataLoader(torch.utils.data.DataLoader(ds, batch_size=bs, \n               shuffle=False, drop_last=False, num_workers=num_workers), device)\ndel df_test\ngc.collect()\n\nmodels = []\nfor m in MODELS:\n    model = RNA_Model()   \n    model = model.to(device)\n    model.load_state_dict(torch.load(m,map_location=torch.device('cpu')))\n    model.eval()\n    models.append(model)","metadata":{"execution":{"iopub.status.busy":"2023-09-11T02:45:58.469199Z","iopub.execute_input":"2023-09-11T02:45:58.469958Z","iopub.status.idle":"2023-09-11T02:46:06.167850Z","shell.execute_reply.started":"2023-09-11T02:45:58.469926Z","shell.execute_reply":"2023-09-11T02:46:06.166810Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"ids,preds = [],[]\nfor x,y in tqdm(dl):\n    with torch.no_grad(),torch.cuda.amp.autocast():\n        p = torch.stack([torch.nan_to_num(model(x)) for model in models]\n                        ,0).mean(0).clip(0,1)\n        \n    for idx, mask, pi in zip(y['ids'].cpu(), x['mask'].cpu(), p.cpu()):\n        ids.append(idx[mask])\n        preds.append(pi[mask[:pi.shape[0]]])\n\nids = torch.concat(ids)\npreds = torch.concat(preds)\n\ndf = pd.DataFrame({'id':ids.numpy(), 'reactivity_DMS_MaP':preds[:,1].numpy(), \n                   'reactivity_2A3_MaP':preds[:,0].numpy()})\ndf.to_csv('submission.csv', index=False, float_format='%.4f') # 6.5GB\ndf.head()","metadata":{"scrolled":true,"execution":{"iopub.status.busy":"2023-09-11T02:46:06.170245Z","iopub.execute_input":"2023-09-11T02:46:06.170945Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}],"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}}