{"metadata":{"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"}},"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-10-03T05:20:26.599136Z","iopub.execute_input":"2023-10-03T05:20:26.600113Z","iopub.status.idle":"2023-10-03T05:20:28.729487Z","shell.execute_reply.started":"2023-10-03T05:20:26.600067Z","shell.execute_reply":"2023-10-03T05:20:28.728532Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"MODELS = ['/kaggle/input/roformer-base-pretrain/roformer_base_pretrain_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-10-03T05:20:28.731207Z","iopub.execute_input":"2023-10-03T05:20:28.732225Z","iopub.status.idle":"2023-10-03T05:20:28.760884Z","shell.execute_reply.started":"2023-10-03T05:20:28.732191Z","shell.execute_reply":"2023-10-03T05:20:28.759946Z"},"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':3,'C':4,'G':5,'U':6}\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-10-03T05:20:28.762579Z","iopub.execute_input":"2023-10-03T05:20:28.762902Z","iopub.status.idle":"2023-10-03T05:20:28.778421Z","shell.execute_reply.started":"2023-10-03T05:20:28.762867Z","shell.execute_reply":"2023-10-03T05:20:28.777381Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model","metadata":{}},{"cell_type":"code","source":"from transformers import AutoConfig, AutoModel\n\nclass RNA_Model(nn.Module):\n    def __init__(self, dim=768, depth=12, heads=12, **kwargs):\n        super().__init__()\n        config = AutoConfig.from_pretrained('/kaggle/input/roformer-base-pretrain/config.json')\n        self.transformer = AutoModel.from_config(config)\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        x = self.transformer(x, attention_mask=mask)[0]\n        x = self.proj_out(x)\n        \n        return x","metadata":{"execution":{"iopub.status.busy":"2023-10-03T05:20:28.780943Z","iopub.execute_input":"2023-10-03T05:20:28.781681Z","iopub.status.idle":"2023-10-03T05:20:28.792714Z","shell.execute_reply.started":"2023-10-03T05:20:28.781645Z","shell.execute_reply":"2023-10-03T05:20:28.791824Z"},"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-10-03T05:20:28.794058Z","iopub.execute_input":"2023-10-03T05:20:28.794398Z","iopub.status.idle":"2023-10-03T05:20:34.979355Z","shell.execute_reply.started":"2023-10-03T05:20:28.794367Z","shell.execute_reply":"2023-10-03T05:20:34.978344Z"},"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\n\n# import polars as pl\n# schema={k:pl.Float32 for k in ['id', 'reactivity_DMS_MaP', 'reactivity_2A3_MaP']}\n# schema['id'] = pl.Int64\n\n# df = pl.DataFrame(\n#     data=[ids.numpy(), preds[:,1].numpy(), preds[:,0].numpy()],\n#     schema=schema\n# )\n# df.write_csv('submission.csv', float_precision=4) \ndf.head()","metadata":{"scrolled":true,"execution":{"iopub.status.busy":"2023-10-03T05:20:34.981041Z","iopub.execute_input":"2023-10-03T05:20:34.981625Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}