{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":67356,"databundleVersionId":8006601,"sourceType":"competition"}],"dockerImageVersionId":30732,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"## Disclaimer !!!\n\nThis is illustrative code for a high performance transformer. I use my local machine for training and inference. I did not check if this code can be run on kaggle notbook or not.\n\ni achieve lb0.428 (new metric) for single-fold without tricks.","metadata":{}},{"cell_type":"code","source":"%%script false --no-raise-error\n\n#tokenization ====================================\n\n# https://www.ascii-code.com/\nMOLECULE_DICT = {\n    'l': 1, 'y': 2, '@': 3, '3': 4, 'H': 5, 'S': 6, 'F': 7, 'C': 8, 'r': 9, 's': 10, '/': 11, 'c': 12, 'o': 13,\n    '+': 14, 'I': 15, '5': 16, '(': 17, '2': 18, ')': 19, '9': 20, 'i': 21, '#': 22, '6': 23, '8': 24, '4': 25,\n    '=': 26, '1': 27, 'O': 28, '[': 29, 'D': 30, 'B': 31, ']': 32, 'N': 33, '7': 34, 'n': 35, '-': 36\n}\nMAX_MOLECULE_ID = np.max(list(MOLECULE_DICT.values()))\nVOCAB_SIZE = MAX_MOLECULE_ID + 10\nUNK = 255  # disallow: will cuase error\nBOS = MAX_MOLECULE_ID + 1\nEOS = MAX_MOLECULE_ID + 2\n# rest are reserved\nPAD = 0\nMAX_LENGTH = 160\n\nMOLECULE_LUT = np.full(256, fill_value=UNK, dtype=np.uint8)\nfor k, v in MOLECULE_DICT.items():\n    ascii = ord(k)\n    MOLECULE_LUT[ascii] = v\n\n        \n\ndef make_token(s):\n    \n    t = np.frombuffer(s, np.uint8)\n    t = MOLECULE_LUT[t]\n    t = t.tolist()\n\n    L = len(t) + 2\n    token_id = [BOS] + t + [EOS] + [PAD] * (MAX_LENGTH - L)\n    token_mask = [1] * L + [0] * (MAX_LENGTH - L)\n        \n    return token_id, token_mask\n\n\n#note this is byte-string!!!\n# string to byte-string: str.encode(s)\nsmiles = b'C#CCOc1ccc(CNc2nc(NCc3cccc(Br)n3)nc(N[C@@H](CC#C)CC(=O)N[Dy])n2)cc1'  \ntoken_id, token_mask = make_token(smiles)","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2024-06-07T03:08:51.921268Z","iopub.execute_input":"2024-06-07T03:08:51.921630Z","iopub.status.idle":"2024-06-07T03:08:51.975279Z","shell.execute_reply.started":"2024-06-07T03:08:51.921601Z","shell.execute_reply":"2024-06-07T03:08:51.974325Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"for modeling, you need flash attnetion + torch compile to make it run fast.\nI train on A6000 gpu with batch size= 2000.","metadata":{}},{"cell_type":"code","source":"%%script false --no-raise-error\n\n#model ====================================\n# https://gist.github.com/kklemon/98e491ff877c497668c715541f1bf478\n# refer to the link above to get fast flash attention wrapper\n\nclass FlashAttentionTransformerEncoder(nn.Module):\n    def __init__(\n            self,\n            dim_model,\n            num_layers,\n            num_heads=None,\n            dim_feedforward=None,\n            dropout=0.0,\n            norm_first=False,\n            activation=F.gelu,\n            rotary_emb_dim=0,\n    ):\n        super().__init__()\n        ... \n\n\nclass PositionalEncoding(nn.Module):\n    def __init__(self, d_model, max_len=256):\n        super(PositionalEncoding, self).__init__()\n\n        pe = torch.zeros(max_len, d_model)\n        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)\n        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))\n        pe[:, 0::2] = torch.sin(position * div_term)\n        pe[:, 1::2] = torch.cos(position * div_term)\n        pe = pe.unsqueeze(0)\n        self.register_buffer('pe', pe)\n\n    def forward(self, x):\n        x = x + self.pe[ :,:x.size(1)]\n\t\treturn x\n\nclass Net(nn.Module):\n    def __init__(self, ):\n        super().__init__()\n\n        embed_dim=512\n\n        self.output_type = ['infer', 'loss']\n        self.pe = PositionalEncoding(embed_dim,max_len=256)\n        self.embedding = nn.Embedding(VOCAB_SIZE, 64, padding_idx=PAD)\n        self.conv_embedding = nn.Sequential(\n            Conv1dBnRelu(64, embed_dim, kernel_size=3,stride=1,padding=1, is_bn=True),\n        )  #just a simple conv1d-bn-relu . for bn use: BN = partial(nn.BatchNorm1d, eps=5e-3,momentum=0.1)\n\n        self.tx_encoder = FlashAttentionTransformerEncoder(\n            dim_model=embed_dim,\n            num_heads=8,\n            dim_feedforward=embed_dim*4,\n            dropout=0.1,\n            norm_first=False,\n            activation=F.gelu,\n            rotary_emb_dim=0,\n            num_layers=7,\n        )\n\n        self.bind = nn.Sequential(\n            nn.Linear(embed_dim, 3),\n        )\n\n\n    def forward(self, batch):\n        smiles_token_id   = batch['smiles_token_id'].long()\n        smiles_token_mask = batch['smiles_token_mask'].long()\n        B, L = smiles_token_id.shape\n\n        x = self.embedding(smiles_token_id)\n        x = x.permute(0,2,1).float()\n        x = self.conv_embedding(x)\n        x = x.permute(0,2,1).contiguous()\n\n        x = self.pe(x)\n        z = self.tx_encoder(\n            x=x,\n            src_key_padding_mask=smiles_token_mask==0,\n        )\n\n        m = smiles_token_mask.unsqueeze(-1).float()\n        pool = (z*m).sum(1)/m.sum(1)\n        bind = self.bind(pool)\n\n        # --------------------------\n        output = {}\n        if 'loss' in self.output_type:\n            target = batch['bind']\n            output['bce_loss'] = F.binary_cross_entropy_with_logits(bind.float(), target.float())\n\n        if 'infer' in self.output_type:\n            output['bind'] = torch.sigmoid(bind)\n\n        return output\n\n    \n#-------------------------------------\n#dummy code to check net\ndef run_check_net():\n    max_length = MAX_LENGTH\n    batch_size = 500\n\n    batch = {\n        'smiles_token_id': torch.from_numpy(np.random.choice(VOCAB_SIZE, (batch_size, max_length))).byte().cuda(),\n        'smiles_token_mask': torch.from_numpy(np.random.choice(2, (batch_size, max_length))).byte().cuda(),\n        'bind': torch.from_numpy(np.random.choice(2, (batch_size, 3))).float().cuda(),\n    }\n     \n    net = Net().cuda()\n    #print(net)\n    #net.train()\n\n    with torch.no_grad():\n        with torch.cuda.amp.autocast(enabled=True): # dtype=torch.float16):\n            output = net(batch)\n\n    # ---\n    print('batch')\n    for k, v in batch.items():\n        if k=='idx':\n            print(f'{k:>32} : {len(v)} ')\n        else:\n            print(f'{k:>32} : {v.shape} ')\n\n    print('output')\n    for k, v in output.items():\n        if 'loss' not in k:\n            print(f'{k:>32} : {v.shape} ')\n    print('loss')\n    for k, v in output.items():\n        if 'loss' in k:\n            print(f'{k:>32} : {v.item()} ')","metadata":{"execution":{"iopub.status.busy":"2024-06-07T03:08:51.977250Z","iopub.execute_input":"2024-06-07T03:08:51.977563Z","iopub.status.idle":"2024-06-07T03:08:51.988189Z","shell.execute_reply.started":"2024-06-07T03:08:51.977537Z","shell.execute_reply":"2024-06-07T03:08:51.987025Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"I also tried other next generation seq models.  \ne.g. mamba, xlstm, Griffin (deepmind's attention RNN)  \n\nI am looking for a faster alternative. But so far transformer + flash attnetion2 is still the fastest  (for small dim and num of layers). for performance, i think these models are smiliar.","metadata":{}},{"cell_type":"code","source":"%%script false --no-raise-error\n\n#xlstm model\n# offical repo: https://github.com/NX-AI/xlstm\n\nclass Net(nn.Module):\n    def __init__(self, ):\n        super().__init__()\n\n        embed_dim = 256\n\n        self.output_type = ['infer', 'loss']\n        self.embedding = nn.Embedding(VOCAB_SIZE, 64, padding_idx=0)\n\n        self.conv_embedding = nn.Sequential(\n            Conv1dBnRelu(64, embed_dim, kernel_size=3, stride=1, padding=1, is_bn=False),\n        )\n\n        self.lstm_encoder = xLSTMBlockStack(\n            config = xLSTMBlockStackConfig(\n                mlstm_block=mLSTMBlockConfig(\n                    mlstm=mLSTMLayerConfig(\n                        conv1d_kernel_size=4, qkv_proj_blocksize=4, num_heads=4\n                    )\n                ),\n                slstm_block=sLSTMBlockConfig(\n                    slstm=sLSTMLayerConfig(\n                        backend='cuda',\n                        batch_size=64,\n                        num_heads=8,\n                        conv1d_kernel_size=4,\n                        bias_init='powerlaw_blockdependent',\n                    ),\n                    feedforward=FeedForwardConfig(proj_factor=1.3, act_fn='gelu'),\n                ),\n                context_length=MAX_LENGTH,\n                num_blocks=6,\n                embedding_dim=embed_dim,\n                slstm_at=[1],\n            )\n        )\n\n        self.bind = nn.Sequential(\n            nn.Linear(embed_dim, 3),\n        )\n\n\n    def forward(self, batch):\n        smiles_token_id = batch['smiles_token_id'].long()\n        B, L = smiles_token_id.shape\n\n        x = self.embedding(smiles_token_id)\n        x = x.permute(0,2,1).float()\n        x = self.conv_embedding(x)\n        x = x.permute(0,2,1).contiguous()\n\n        x = self.lstm_encoder(x)\n        last = x.mean(1)\n        bind = self.bind(last)\n\n        # --------------------------\n        output = {}\n        if 'loss' in self.output_type:\n            target = batch['bind']\n            output['bce_loss'] = F.binary_cross_entropy_with_logits(bind.float(), target.float())\n\n        if 'infer' in self.output_type:\n            output['bind'] = torch.sigmoid(bind)\n\n        return output\n","metadata":{"execution":{"iopub.status.busy":"2024-06-07T03:08:51.989893Z","iopub.execute_input":"2024-06-07T03:08:51.990281Z","iopub.status.idle":"2024-06-07T03:08:52.003495Z","shell.execute_reply.started":"2024-06-07T03:08:51.990250Z","shell.execute_reply":"2024-06-07T03:08:52.002567Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%script false --no-raise-error\n# official repo: https://github.com/state-spaces/mamba\n\n#https://github.com/state-spaces/mamba/issues/355\n# there is a bug? i cannot try mamba2\n\nclass Net(nn.Module):\n    def __init__(self, ):\n        super().__init__()\n\n        embed_dim=256\n        num_layer=6\n        self.output_type = ['infer', 'loss']\n        self.embedding = nn.Embedding(VOCAB_SIZE, 64, padding_idx=PAD)\n        self.pe = PositionalEncoding(embed_dim,max_len=256)\n\n        self.conv_embedding = nn.Sequential(\n            Conv1dBnRelu(64, embed_dim, kernel_size=3,stride=1,padding=1, is_bn=True),\n        )\n\n        self.mamba_encoder = nn.ModuleList(\n            [\n                create_block(\n                    embed_dim,\n                    d_intermediate=embed_dim//2,\n                    ssm_cfg={'layer': 'Mamba1'},\n                    attn_layer_idx=None,\n                    attn_cfg=None,\n                    norm_epsilon=1e-4,\n                    rms_norm=1e-4,\n                    residual_in_fp32=False,\n                    fused_add_norm=True,\n                    layer_idx=i,\n                )\n                for i in range(num_layer)\n            ])\n\n        self.norm_f = nn.LayerNorm ( #RMSNorm\n            embed_dim, eps=1e-4\n        )\n\n        self.bind = nn.Sequential(\n            nn.Linear(embed_dim, 3),\n        )\n\n        self.apply(\n            partial(\n                _init_weights,\n                n_layer=num_layer,\n                n_residuals_per_layer=2,\n            )\n        )\n\n    def forward(self, batch):\n        smiles_token_id = batch['smiles_token_id'].long()\n        smiles_token_mask = batch['smiles_token_mask'].long()\n        B, L  = smiles_token_id.shape\n\n        x = self.embedding(smiles_token_id)\n        x = x.permute(0,2,1).float()\n        x = self.conv_embedding(x)\n        x = x.permute(0,2,1).contiguous()\n      \n        hidden, residual = x, None\n        for mamba in self.mamba_encoder:\n            hidden, residual = mamba(\n                hidden, residual, inference_params=None\n            )\n            hidden = F.dropout(hidden,p=0.1, training=self.training)\n\n        #z=hidden\n        z = layer_norm_fn(\n            hidden,\n            self.norm_f.weight,\n            self.norm_f.bias,\n            eps=self.norm_f.eps,\n            residual=residual,\n            prenorm=False,\n            residual_in_fp32=False,\n            is_rms_norm=isinstance(self.norm_f, RMSNorm)\n        )\n\n        #pool = z.mean(1)\n        m = smiles_token_mask.unsqueeze(2).float()\n        pool = (z*m).sum(1)/m.sum(1)\n        bind = self.bind(pool)\n\n        # --------------------------\n        output = {}\n        if 'loss' in self.output_type:\n            target = batch['bind']\n            output['bce_loss'] = F.binary_cross_entropy_with_logits(bind.float(), target.float())\n\n        if 'infer' in self.output_type:\n            output['bind'] = torch.sigmoid(bind)\n\n        return output","metadata":{"execution":{"iopub.status.busy":"2024-06-07T03:08:52.005646Z","iopub.execute_input":"2024-06-07T03:08:52.006055Z","iopub.status.idle":"2024-06-07T03:08:52.020395Z","shell.execute_reply.started":"2024-06-07T03:08:52.006010Z","shell.execute_reply":"2024-06-07T03:08:52.019248Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# to be updated: hawk and griff\n# https://github.com/proger/hippogriff/blob/main/hippogriff.py","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"#### ","metadata":{}}]}