notebooks/pretrained_molformer/frozen_embeddings_classification.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Pre-trained MolFormer embeddings demo\n",
    "\n",
    "This notebook demonstrates how to load a MolFormer checkpoint file and use it to output pre-trained embeddings (without fine-tuning) given a list of SMILES strings. These can then be used with a simple model to perform a classification task."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load checkpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Namespace(accelerator='ddp', batch_size=64, beam_size=0, checkpoint_every=5000, clip_grad=50, config_load=None, config_save=None, d_dropout=0.2, data_path='', data_root='/dccstor/medscan7/smallmolecule/runs/ba-predictor/small-data/affinity', dataset_length=None, dataset_name='sol', debug=True, device='cuda', dropout=0.1, eval_every=1000, fast_dev_run=False, fc_h=512, finetune_path='', freeze_model=False, from_scratch=False, gen_save=None, gpus=8, grad_acc=1, log_file=None, lr=0.001, lr_end=0.00030000000000000003, lr_multiplier=8, lr_start=3e-05, max_epochs=4, max_len=202, measure_name='measure', min_len=1, mode='cls', model_arch='BERT_16GPU_Both_10percent_rotate_no_masking', model_load=None, model_save='model.pt', model_save_dir='./models_dump/', n_batch=1800, n_embd=768, n_head=12, n_jobs=1, n_last=1000, n_layer=12, n_samples=None, n_workers=8, nucleus_thresh=0.9, num_epoch=1, num_feats=32, num_nodes=1, num_seq_returned=0, num_workers=0, pretext_size=0, q_dropout=0.5, restart_path='', resume=False, root_dir='.', rotate=True, seed=12345, smiles_embedding='/dccstor/medscan7/smallmolecule/runs/ba-predictor/small-data/embeddings/protein/ba_embeddings_tanh_512_2986138_2.pt', tb_loc=None, train_load='both', unlike_alpha=1.0, unlikelihood=False, use_linear=0, val_load=None, vocab_load=None, vocab_save=None)"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from argparse import Namespace\n",
    "import yaml\n",
    "\n",
    "with open('../../data/Pretrained MoLFormer/hparams.yaml', 'r') as f:\n",
    "    config = Namespace(**yaml.safe_load(f))\n",
    "config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "OrderedDict([('<bos>', 0),\n",
       "             ('<eos>', 1),\n",
       "             ('<pad>', 2),\n",
       "             ('<mask>', 3),\n",
       "             ('C', 4),\n",
       "             ('c', 5),\n",
       "             ('(', 6),\n",
       "             (')', 7),\n",
       "             ('1', 8),\n",
       "             ('O', 9),\n",
       "             ('N', 10),\n",
       "             ('2', 11),\n",
       "             ('=', 12),\n",
       "             ('n', 13),\n",
       "             ('3', 14),\n",
       "             ('[C@H]', 15),\n",
       "             ('[C@@H]', 16),\n",
       "             ('F', 17),\n",
       "             ('S', 18),\n",
       "             ('4', 19),\n",
       "             ('Cl', 20),\n",
       "             ('-', 21),\n",
       "             ('o', 22),\n",
       "             ('s', 23),\n",
       "             ('[nH]', 24),\n",
       "             ('#', 25),\n",
       "             ('/', 26),\n",
       "             ('Br', 27),\n",
       "             ('[C@]', 28),\n",
       "             ('[C@@]', 29),\n",
       "             ('[N+]', 30),\n",
       "             ('[O-]', 31),\n",
       "             ('5', 32),\n",
       "             ('\\\\', 33),\n",
       "             ('.', 34),\n",
       "             ('I', 35),\n",
       "             ('6', 36),\n",
       "             ('[S@]', 37),\n",
       "             ('[S@@]', 38),\n",
       "             ('P', 39),\n",
       "             ('[N-]', 40),\n",
       "             ('[Si]', 41),\n",
       "             ('7', 42),\n",
       "             ('[n+]', 43),\n",
       "             ('[2H]', 44),\n",
       "             ('8', 45),\n",
       "             ('[NH+]', 46),\n",
       "             ('B', 47),\n",
       "             ('9', 48),\n",
       "             ('[C-]', 49),\n",
       "             ('[Na+]', 50),\n",
       "             ('[Cl-]', 51),\n",
       "             ('[c-]', 52),\n",
       "             ('[CH]', 53),\n",
       "             ('%10', 54),\n",
       "             ('[NH2+]', 55),\n",
       "             ('[P+]', 56),\n",
       "             ('[B]', 57),\n",
       "             ('[I-]', 58),\n",
       "             ('%11', 59),\n",
       "             ('[CH2-]', 60),\n",
       "             ('[O+]', 61),\n",
       "             ('[NH3+]', 62),\n",
       "             ('[C]', 63),\n",
       "             ('[Br-]', 64),\n",
       "             ('[IH2]', 65),\n",
       "             ('[S-]', 66),\n",
       "             ('[cH-]', 67),\n",
       "             ('%12', 68),\n",
       "             ('[nH+]', 69),\n",
       "             ('[B-]', 70),\n",
       "             ('[K+]', 71),\n",
       "             ('[Sn]', 72),\n",
       "             ('[Se]', 73),\n",
       "             ('[CH-]', 74),\n",
       "             ('[HH]', 75),\n",
       "             ('[Y]', 76),\n",
       "             ('[n-]', 77),\n",
       "             ('[CH3-]', 78),\n",
       "             ('[SiH]', 79),\n",
       "             ('[S+]', 80),\n",
       "             ('%13', 81),\n",
       "             ('[SiH2]', 82),\n",
       "             ('[Li+]', 83),\n",
       "             ('[NH-]', 84),\n",
       "             ('%14', 85),\n",
       "             ('[Na]', 86),\n",
       "             ('[CH2]', 87),\n",
       "             ('[O-2]', 88),\n",
       "             ('[U+2]', 89),\n",
       "             ('[W]', 90),\n",
       "             ('[Al]', 91),\n",
       "             ('[P@]', 92),\n",
       "             ('[Fe+2]', 93),\n",
       "             ('[PH+]', 94),\n",
       "             ('%15', 95),\n",
       "             ('[Cl+3]', 96),\n",
       "             ('[Zn+2]', 97),\n",
       "             ('[Ir]', 98),\n",
       "             ('[Mg+2]', 99),\n",
       "             ('[Pt+2]', 100),\n",
       "             ('[OH2+]', 101),\n",
       "             ('[As]', 102),\n",
       "             ('[Fe]', 103),\n",
       "             ('[OH+]', 104),\n",
       "             ('[Zr+2]', 105),\n",
       "             ('[3H]', 106),\n",
       "             ('[Ge]', 107),\n",
       "             ('[SiH3]', 108),\n",
       "             ('[OH-]', 109),\n",
       "             ('[NH4+]', 110),\n",
       "             ('[Cu+2]', 111),\n",
       "             ('[P@@]', 112),\n",
       "             ('p', 113),\n",
       "             ('[Pt]', 114),\n",
       "             ('%16', 115),\n",
       "             ('[Ca+2]', 116),\n",
       "             ('[Zr]', 117),\n",
       "             ('[F-]', 118),\n",
       "             ('[C+]', 119),\n",
       "             ('[Ti]', 120),\n",
       "             ('[P-]', 121),\n",
       "             ('[V]', 122),\n",
       "             ('[se]', 123),\n",
       "             ('[U]', 124),\n",
       "             ('[O]', 125),\n",
       "             ('[Ni+2]', 126),\n",
       "             ('[Zn]', 127),\n",
       "             ('[Co]', 128),\n",
       "             ('[Ni]', 129),\n",
       "             ('[Pd+2]', 130),\n",
       "             ('[Cu]', 131),\n",
       "             ('%17', 132),\n",
       "             ('[Cu+]', 133),\n",
       "             ('[Te]', 134),\n",
       "             ('[H+]', 135),\n",
       "             ('[CH+]', 136),\n",
       "             ('[Li]', 137),\n",
       "             ('[Pd]', 138),\n",
       "             ('[Mo]', 139),\n",
       "             ('[Ru+2]', 140),\n",
       "             ('[o+]', 141),\n",
       "             ('[Re]', 142),\n",
       "             ('[SH+]', 143),\n",
       "             ('%18', 144),\n",
       "             ('[Ac]', 145),\n",
       "             ('[Cr]', 146),\n",
       "             ('[NH2-]', 147),\n",
       "             ('[K]', 148),\n",
       "             ('[13CH2]', 149),\n",
       "             ('[c]', 150),\n",
       "             ('[Zr+4]', 151),\n",
       "             ('[Tl]', 152),\n",
       "             ('[13C]', 153),\n",
       "             ('[Mn]', 154),\n",
       "             ('[N@+]', 155),\n",
       "             ('[Hg]', 156),\n",
       "             ('[Rh]', 157),\n",
       "             ('[Ti+4]', 158),\n",
       "             ('[Sb]', 159),\n",
       "             ('[Co+2]', 160),\n",
       "             ('[Ag+]', 161),\n",
       "             ('[Ru]', 162),\n",
       "             ('%19', 163),\n",
       "             ('[N@@+]', 164),\n",
       "             ('[Ti+2]', 165),\n",
       "             ('[Al+3]', 166),\n",
       "             ('[Pb]', 167),\n",
       "             ('[I+]', 168),\n",
       "             ('[18F]', 169),\n",
       "             ('[s+]', 170),\n",
       "             ('[Rb+]', 171),\n",
       "             ('[Ba+2]', 172),\n",
       "             ('[H-]', 173),\n",
       "             ('[Fe+3]', 174),\n",
       "             ('[Ir+3]', 175),\n",
       "             ('[13cH]', 176),\n",
       "             ('%20', 177),\n",
       "             ('[AlH2]', 178),\n",
       "             ('[Au+]', 179),\n",
       "             ('[13c]', 180),\n",
       "             ('[SH2+]', 181),\n",
       "             ('[Sn+2]', 182),\n",
       "             ('[Mn+2]', 183),\n",
       "             ('[Si-]', 184),\n",
       "             ('[Ag]', 185),\n",
       "             ('[N]', 186),\n",
       "             ('[Bi]', 187),\n",
       "             ('%21', 188),\n",
       "             ('[In]', 189),\n",
       "             ('[CH2+]', 190),\n",
       "             ('[Y+3]', 191),\n",
       "             ('[Ga]', 192),\n",
       "             ('%22', 193),\n",
       "             ('[Co+3]', 194),\n",
       "             ('[Au]', 195),\n",
       "             ('[13CH3]', 196),\n",
       "             ('[Mg]', 197),\n",
       "             ('[Cs+]', 198),\n",
       "             ('[W+2]', 199),\n",
       "             ('[Hf]', 200),\n",
       "             ('[Zn+]', 201),\n",
       "             ('[Se-]', 202),\n",
       "             ('[S-2]', 203),\n",
       "             ('[Ca]', 204),\n",
       "             ('[pH]', 205),\n",
       "             ('[ClH+]', 206),\n",
       "             ('[Ti+3]', 207),\n",
       "             ('%23', 208),\n",
       "             ('[Ru+]', 209),\n",
       "             ('[SH-]', 210),\n",
       "             ('[13CH]', 211),\n",
       "             ('[IH+]', 212),\n",
       "             ('[Hf+4]', 213),\n",
       "             ('[Rf]', 214),\n",
       "             ('[OH3+]', 215),\n",
       "             ('%24', 216),\n",
       "             ('[Pt+4]', 217),\n",
       "             ('[Zr+3]', 218),\n",
       "             ('[PH3+]', 219),\n",
       "             ('[Sr+2]', 220),\n",
       "             ('[Cd+2]', 221),\n",
       "             ('[Cd]', 222),\n",
       "             ('%25', 223),\n",
       "             ('[Os]', 224),\n",
       "             ('[BH-]', 225),\n",
       "             ('[Sn+4]', 226),\n",
       "             ('[Cr+3]', 227),\n",
       "             ('[Ru+3]', 228),\n",
       "             ('[PH2+]', 229),\n",
       "             ('[Rh+2]', 230),\n",
       "             ('[V+2]', 231),\n",
       "             ('%26', 232),\n",
       "             ('[Gd+3]', 233),\n",
       "             ('[Pb+2]', 234),\n",
       "             ('[PH]', 235),\n",
       "             ('[Hg+]', 236),\n",
       "             ('[Mo+2]', 237),\n",
       "             ('[AlH]', 238),\n",
       "             ('[Sn+]', 239),\n",
       "             ('%27', 240),\n",
       "             ('[Pd+]', 241),\n",
       "             ('b', 242),\n",
       "             ('[Rh+3]', 243),\n",
       "             ('[Hg+2]', 244),\n",
       "             ('[15NH]', 245),\n",
       "             ('[14C]', 246),\n",
       "             ('%28', 247),\n",
       "             ('[Mn+3]', 248),\n",
       "             ('[Si+]', 249),\n",
       "             ('[SeH]', 250),\n",
       "             ('[13C@H]', 251),\n",
       "             ('[NH]', 252),\n",
       "             ('[Ga+3]', 253),\n",
       "             ('[SiH-]', 254),\n",
       "             ('[13C@@H]', 255),\n",
       "             ('[Ce]', 256),\n",
       "             ('[Au+3]', 257),\n",
       "             ('[Bi+3]', 258),\n",
       "             ('[15N]', 259),\n",
       "             ('%29', 260),\n",
       "             ('[BH3-]', 261),\n",
       "             ('[14cH]', 262),\n",
       "             ('[Ti+]', 263),\n",
       "             ('[Gd]', 264),\n",
       "             ('[cH+]', 265),\n",
       "             ('[Cr+2]', 266),\n",
       "             ('[Sb-]', 267),\n",
       "             ('%30', 268),\n",
       "             ('[Be+2]', 269),\n",
       "             ('[Al+]', 270),\n",
       "             ('[te]', 271),\n",
       "             ('[11CH3]', 272),\n",
       "             ('[Sm]', 273),\n",
       "             ('[Pr]', 274),\n",
       "             ('[La]', 275),\n",
       "             ('%31', 276),\n",
       "             ('[Al-]', 277),\n",
       "             ('[Ta]', 278),\n",
       "             ('[125I]', 279),\n",
       "             ('[BH2-]', 280),\n",
       "             ('[Nb]', 281),\n",
       "             ('[Si@]', 282),\n",
       "             ('%32', 283),\n",
       "             ('[14c]', 284),\n",
       "             ('[Sb+3]', 285),\n",
       "             ('[Ba]', 286),\n",
       "             ('%33', 287),\n",
       "             ('[Os+2]', 288),\n",
       "             ('[Si@@]', 289),\n",
       "             ('[La+3]', 290),\n",
       "             ('[15n]', 291),\n",
       "             ('[15NH2]', 292),\n",
       "             ('[Nd+3]', 293),\n",
       "             ('%34', 294),\n",
       "             ('[14CH2]', 295),\n",
       "             ('[18O]', 296),\n",
       "             ('[Nd]', 297),\n",
       "             ('[GeH]', 298),\n",
       "             ('[Ni+3]', 299),\n",
       "             ('[Eu]', 300),\n",
       "             ('[Dy+3]', 301),\n",
       "             ('[Sc]', 302),\n",
       "             ('%36', 303),\n",
       "             ('[Se-2]', 304),\n",
       "             ('[As+]', 305),\n",
       "             ('%35', 306),\n",
       "             ('[AsH]', 307),\n",
       "             ('[Tb]', 308),\n",
       "             ('[Sb+5]', 309),\n",
       "             ('[Se+]', 310),\n",
       "             ('[Ce+3]', 311),\n",
       "             ('[c+]', 312),\n",
       "             ('[In+3]', 313),\n",
       "             ('[SnH]', 314),\n",
       "             ('[Mo+4]', 315),\n",
       "             ('%37', 316),\n",
       "             ('[V+4]', 317),\n",
       "             ('[Eu+3]', 318),\n",
       "             ('[Hf+2]', 319),\n",
       "             ('%38', 320),\n",
       "             ('[Pt+]', 321),\n",
       "             ('[p+]', 322),\n",
       "             ('[123I]', 323),\n",
       "             ('[Tl+]', 324),\n",
       "             ('[Sm+3]', 325),\n",
       "             ('%39', 326),\n",
       "             ('[Yb+3]', 327),\n",
       "             ('%40', 328),\n",
       "             ('[Yb]', 329),\n",
       "             ('[Os+]', 330),\n",
       "             ('%41', 331),\n",
       "             ('[10B]', 332),\n",
       "             ('[Sc+3]', 333),\n",
       "             ('[Al+2]', 334),\n",
       "             ('%42', 335),\n",
       "             ('[Sr]', 336),\n",
       "             ('[Tb+3]', 337),\n",
       "             ('[Po]', 338),\n",
       "             ('[Tc]', 339),\n",
       "             ('[PH-]', 340),\n",
       "             ('[AlH3]', 341),\n",
       "             ('[Ar]', 342),\n",
       "             ('[U+4]', 343),\n",
       "             ('[SnH2]', 344),\n",
       "             ('[Cl+2]', 345),\n",
       "             ('[si]', 346),\n",
       "             ('[Fe+]', 347),\n",
       "             ('[14CH3]', 348),\n",
       "             ('[U+3]', 349),\n",
       "             ('[Cl+]', 350),\n",
       "             ('%43', 351),\n",
       "             ('[GeH2]', 352),\n",
       "             ('%44', 353),\n",
       "             ('[Er+3]', 354),\n",
       "             ('[Mo+3]', 355),\n",
       "             ('[I+2]', 356),\n",
       "             ('[Fe+4]', 357),\n",
       "             ('[99Tc]', 358),\n",
       "             ('%45', 359),\n",
       "             ('[11C]', 360),\n",
       "             ('%46', 361),\n",
       "             ('[SnH3]', 362),\n",
       "             ('[S]', 363),\n",
       "             ('[Te+]', 364),\n",
       "             ('[Er]', 365),\n",
       "             ('[Lu+3]', 366),\n",
       "             ('[11B]', 367),\n",
       "             ('%47', 368),\n",
       "             ('%48', 369),\n",
       "             ('[P]', 370),\n",
       "             ('[Tm]', 371),\n",
       "             ('[Th]', 372),\n",
       "             ('[Dy]', 373),\n",
       "             ('[Pr+3]', 374),\n",
       "             ('[Ta+5]', 375),\n",
       "             ('[Nb+5]', 376),\n",
       "             ('[Rb]', 377),\n",
       "             ('[GeH3]', 378),\n",
       "             ('[Br+2]', 379),\n",
       "             ('%49', 380),\n",
       "             ('[131I]', 381),\n",
       "             ('[Fm]', 382),\n",
       "             ('[Cs]', 383),\n",
       "             ('[BH4-]', 384),\n",
       "             ('[Lu]', 385),\n",
       "             ('[15nH]', 386),\n",
       "             ('%50', 387),\n",
       "             ('[Ru+6]', 388),\n",
       "             ('[b-]', 389),\n",
       "             ('[Ho]', 390),\n",
       "             ('[Th+4]', 391),\n",
       "             ('[Ru+4]', 392),\n",
       "             ('%52', 393),\n",
       "             ('[14CH]', 394),\n",
       "             ('%51', 395),\n",
       "             ('[Cr+6]', 396),\n",
       "             ('[18OH]', 397),\n",
       "             ('[Ho+3]', 398),\n",
       "             ('[Ce+4]', 399),\n",
       "             ('[Bi+2]', 400),\n",
       "             ('[Co+]', 401),\n",
       "             ('%53', 402),\n",
       "             ('[Yb+2]', 403),\n",
       "             ('[Fe+6]', 404),\n",
       "             ('[Be]', 405),\n",
       "             ('%54', 406),\n",
       "             ('[SH3+]', 407),\n",
       "             ('[Np]', 408),\n",
       "             ('[As-]', 409),\n",
       "             ('%55', 410),\n",
       "             ('[14C@@H]', 411),\n",
       "             ('[Ir+2]', 412),\n",
       "             ('[GaH3]', 413),\n",
       "             ('[p-]', 414),\n",
       "             ('[GeH4]', 415),\n",
       "             ('[Sn+3]', 416),\n",
       "             ('[Os+4]', 417),\n",
       "             ('%56', 418),\n",
       "             ('[14C@H]', 419),\n",
       "             ('[sH+]', 420),\n",
       "             ('[19F]', 421),\n",
       "             ('[Eu+2]', 422),\n",
       "             ('[TlH]', 423),\n",
       "             ('%57', 424),\n",
       "             ('[Cr+4]', 425),\n",
       "             ('%58', 426),\n",
       "             ('[B@@-]', 427),\n",
       "             ('[SiH+]', 428),\n",
       "             ('[At]', 429),\n",
       "             ('[Am]', 430),\n",
       "             ('[Fe+5]', 431),\n",
       "             ('[AsH2]', 432),\n",
       "             ('[Si+4]', 433),\n",
       "             ('[B@-]', 434),\n",
       "             ('[Pu]', 435),\n",
       "             ('[SbH]', 436),\n",
       "             ('[P-2]', 437),\n",
       "             ('[Tm+3]', 438),\n",
       "             ('*', 439),\n",
       "             ('%59', 440),\n",
       "             ('[se+]', 441),\n",
       "             ('%60', 442),\n",
       "             ('[oH+]', 443),\n",
       "             ('[1H]', 444),\n",
       "             ('[15N+]', 445),\n",
       "             ('[124I]', 446),\n",
       "             ('[S@@+]', 447),\n",
       "             ('[P-3]', 448),\n",
       "             ('[H]', 449),\n",
       "             ('[IH2+]', 450),\n",
       "             ('[TeH]', 451),\n",
       "             ('[Xe]', 452),\n",
       "             ('[PH4+]', 453),\n",
       "             ('[Cr+]', 454),\n",
       "             ('[Cm]', 455),\n",
       "             ('[I+3]', 456),\n",
       "             ('%61', 457),\n",
       "             ('[Nb+2]', 458),\n",
       "             ('[Ru+5]', 459),\n",
       "             ('%62', 460),\n",
       "             ('[Ta+2]', 461),\n",
       "             ('[Tc+4]', 462),\n",
       "             ('[CH3+]', 463),\n",
       "             ('[Pm]', 464),\n",
       "             ('[Si@H]', 465),\n",
       "             ('[No]', 466),\n",
       "             ('%63', 467),\n",
       "             ('[Cr+5]', 468),\n",
       "             ('[Th+2]', 469),\n",
       "             ('[Zn-2]', 470),\n",
       "             ('[13C@]', 471),\n",
       "             ('[Lr]', 472),\n",
       "             ('%64', 473),\n",
       "             ('[99Tc+3]', 474),\n",
       "             ('%65', 475),\n",
       "             ('[13C@@]', 476),\n",
       "             ('%66', 477),\n",
       "             ('[Fe-]', 478),\n",
       "             ('[17O]', 479),\n",
       "             ('[siH]', 480),\n",
       "             ('[Sb+]', 481),\n",
       "             ('[OH]', 482),\n",
       "             ('[IH]', 483),\n",
       "             ('[11CH2]', 484),\n",
       "             ('[Cf]', 485),\n",
       "             ('[SiH2+]', 486),\n",
       "             ('[Gd+2]', 487),\n",
       "             ('[In+]', 488),\n",
       "             ('[Si@@H]', 489),\n",
       "             ('[Mn+]', 490),\n",
       "             ('[99Tc+4]', 491),\n",
       "             ('[Ga-]', 492),\n",
       "             ('%67', 493),\n",
       "             ('[S@+]', 494),\n",
       "             ('[Ge+4]', 495),\n",
       "             ('[Tl+3]', 496),\n",
       "             ('[16OH]', 497),\n",
       "             ('%68', 498),\n",
       "             ('[2H-]', 499),\n",
       "             ('[Ra]', 500),\n",
       "             ('[si-]', 501),\n",
       "             ('[NiH2]', 502),\n",
       "             ('[P@@H]', 503),\n",
       "             ('[Rh+]', 504),\n",
       "             ('[12C]', 505),\n",
       "             ('[35S]', 506),\n",
       "             ('[32P]', 507),\n",
       "             ('[SiH2-]', 508),\n",
       "             ('[AlH2+]', 509),\n",
       "             ('[16O]', 510),\n",
       "             ('%69', 511),\n",
       "             ('[BiH]', 512),\n",
       "             ('[BiH2]', 513),\n",
       "             ('[Zn-]', 514),\n",
       "             ('[BH]', 515),\n",
       "             ('[Tc+3]', 516),\n",
       "             ('[Ir+]', 517),\n",
       "             ('[Ni+]', 518),\n",
       "             ('%70', 519),\n",
       "             ('[InH2]', 520),\n",
       "             ('[InH]', 521),\n",
       "             ('[Nb+3]', 522),\n",
       "             ('[PbH]', 523),\n",
       "             ('[Bi+]', 524),\n",
       "             ('%71', 525),\n",
       "             ('[As+3]', 526),\n",
       "             ('%72', 527),\n",
       "             ('[18O-]', 528),\n",
       "             ('[68Ga+3]', 529),\n",
       "             ('%73', 530),\n",
       "             ('[Pa]', 531),\n",
       "             ('[76Br]', 532),\n",
       "             ('[Tc+5]', 533),\n",
       "             ('[pH+]', 534),\n",
       "             ('[64Cu+2]', 535),\n",
       "             ('[Ru+8]', 536),\n",
       "             ('%74', 537),\n",
       "             ('[PH2-]', 538),\n",
       "             ('[Si+2]', 539),\n",
       "             ('[17OH]', 540),\n",
       "             ('[RuH]', 541),\n",
       "             ('[111In+3]', 542),\n",
       "             ('[AlH+]', 543),\n",
       "             ('%75', 544),\n",
       "             ('%76', 545),\n",
       "             ('[W+]', 546),\n",
       "             ('[SbH2]', 547),\n",
       "             ('[PoH]', 548),\n",
       "             ('[Ru-]', 549),\n",
       "             ('[XeH]', 550),\n",
       "             ('[Tc+2]', 551),\n",
       "             ('[13C-]', 552),\n",
       "             ('[Br+]', 553),\n",
       "             ('[Pt-2]', 554),\n",
       "             ('[Es]', 555),\n",
       "             ('[Cu-]', 556),\n",
       "             ('[Mg+]', 557),\n",
       "             ('[3HH]', 558),\n",
       "             ('[P@H]', 559),\n",
       "             ('[ClH2+]', 560),\n",
       "             ('%77', 561),\n",
       "             ('[SH]', 562),\n",
       "             ('[Au-]', 563),\n",
       "             ('[2HH]', 564),\n",
       "             ('%78', 565),\n",
       "             ('[Sn-]', 566),\n",
       "             ('[11CH]', 567),\n",
       "             ('[PdH2]', 568),\n",
       "             ('0', 569),\n",
       "             ('[Os+6]', 570),\n",
       "             ('%79', 571),\n",
       "             ('[Mo+]', 572),\n",
       "             ('%80', 573),\n",
       "             ('[al]', 574),\n",
       "             ('[PbH2]', 575),\n",
       "             ('[64Cu]', 576),\n",
       "             ('[Cl]', 577),\n",
       "             ('[12CH3]', 578),\n",
       "             ('%81', 579),\n",
       "             ('[Tc+7]', 580),\n",
       "             ('[11c]', 581),\n",
       "             ('%82', 582),\n",
       "             ('[Li-]', 583),\n",
       "             ('[99Tc+5]', 584),\n",
       "             ('[He]', 585),\n",
       "             ('[12c]', 586),\n",
       "             ('[Kr]', 587),\n",
       "             ('[RuH+2]', 588),\n",
       "             ('[35Cl]', 589),\n",
       "             ('[Pd-2]', 590),\n",
       "             ('[GaH2]', 591),\n",
       "             ('[4H]', 592),\n",
       "             ('[Sg]', 593),\n",
       "             ('[Cu-2]', 594),\n",
       "             ('[Br+3]', 595),\n",
       "             ('%83', 596),\n",
       "             ('[37Cl]', 597),\n",
       "             ('[211At]', 598),\n",
       "             ('[IrH+2]', 599),\n",
       "             ('[Mt]', 600),\n",
       "             ('[Ir-2]', 601),\n",
       "             ('[In-]', 602),\n",
       "             ('[12cH]', 603),\n",
       "             ('[12CH2]', 604),\n",
       "             ('[RuH2]', 605),\n",
       "             ('[99Tc+7]', 606),\n",
       "             ('%84', 607),\n",
       "             ('[15n+]', 608),\n",
       "             ('[ClH2+2]', 609),\n",
       "             ('[16N]', 610),\n",
       "             ('[111In]', 611),\n",
       "             ('[Tc+]', 612),\n",
       "             ('[Ru-2]', 613),\n",
       "             ('[12CH]', 614),\n",
       "             ('[si+]', 615),\n",
       "             ('[Tc+6]', 616),\n",
       "             ('%85', 617),\n",
       "             ('%86', 618),\n",
       "             ('[90Y]', 619),\n",
       "             ('[Pd-]', 620),\n",
       "             ('[188Re]', 621),\n",
       "             ('[RuH+]', 622),\n",
       "             ('[NiH]', 623),\n",
       "             ('[SiH3-]', 624),\n",
       "             ('[14n]', 625),\n",
       "             ('[CH3]', 626),\n",
       "             ('[14N]', 627),\n",
       "             ('[10BH2]', 628),\n",
       "             ('%88', 629),\n",
       "             ('%89', 630),\n",
       "             ('%90', 631),\n",
       "             ('[34S]', 632),\n",
       "             ('[77Br]', 633),\n",
       "             ('[GaH]', 634),\n",
       "             ('[Br]', 635),\n",
       "             ('[Ge@]', 636),\n",
       "             ('[B@@H-]', 637),\n",
       "             ('[CuH]', 638),\n",
       "             ('[SiH4]', 639),\n",
       "             ('[3H-]', 640),\n",
       "             ('%87', 641),\n",
       "             ('%91', 642),\n",
       "             ('%92', 643),\n",
       "             ('[67Cu]', 644),\n",
       "             ('[I]', 645),\n",
       "             ('[177Lu]', 646),\n",
       "             ('[ReH]', 647),\n",
       "             ('[67Ga+3]', 648),\n",
       "             ('[Db]', 649),\n",
       "             ('[177Lu+3]', 650),\n",
       "             ('[AlH2-]', 651),\n",
       "             ('[Si+3]', 652),\n",
       "             ('[Ti-2]', 653),\n",
       "             ('[RuH+3]', 654),\n",
       "             ('[al+]', 655),\n",
       "             ('[68Ga]', 656),\n",
       "             ('[2H+]', 657),\n",
       "             ('[B@H-]', 658),\n",
       "             ('[WH2]', 659),\n",
       "             ('[OsH]', 660),\n",
       "             ('[Ir-3]', 661),\n",
       "             ('[AlH-]', 662),\n",
       "             ('[Bk]', 663),\n",
       "             ('[75Se]', 664),\n",
       "             ('[14C@]', 665),\n",
       "             ('[Pt-]', 666),\n",
       "             ('[N@@H+]', 667),\n",
       "             ('[Nb-]', 668),\n",
       "             ('[13NH2]', 669),\n",
       "             ('%93', 670),\n",
       "             ('[186Re]', 671),\n",
       "             ('[Tb+4]', 672),\n",
       "             ('[PtH]', 673),\n",
       "             ('[IrH2]', 674),\n",
       "             ('[Hg-2]', 675),\n",
       "             ('[AlH3-]', 676),\n",
       "             ('[PdH+]', 677),\n",
       "             ('[Md]', 678),\n",
       "             ('[RhH+2]', 679),\n",
       "             ('[11cH]', 680),\n",
       "             ('[Co-2]', 681),\n",
       "             ('[15N-]', 682),\n",
       "             ('[ZrH2]', 683),\n",
       "             ('%94', 684),\n",
       "             ('[Hg-]', 685),\n",
       "             ('[127I]', 686),\n",
       "             ('[AsH2+]', 687),\n",
       "             ('[MoH2]', 688),\n",
       "             ('[Te+4]', 689),\n",
       "             ('[14C@@]', 690),\n",
       "             ('[As+5]', 691),\n",
       "             ('[SnH+3]', 692),\n",
       "             ('[Ge@@]', 693),\n",
       "             ('[6Li+]', 694),\n",
       "             ('[WH]', 695),\n",
       "             ('[Ne]', 696),\n",
       "             ('[14NH2]', 697),\n",
       "             ('[14NH]', 698),\n",
       "             ('[12C@@H]', 699),\n",
       "             ('[Os+7]', 700),\n",
       "             ('[RhH]', 701),\n",
       "             ('[Al-3]', 702),\n",
       "             ('[SnH+]', 703),\n",
       "             ('[15NH3+]', 704),\n",
       "             ('[Zr+]', 705),\n",
       "             ('[197Hg+]', 706),\n",
       "             ('%95', 707),\n",
       "             ('%96', 708),\n",
       "             ('[90Y+3]', 709),\n",
       "             ('[Os-2]', 710),\n",
       "             ('[98Tc+5]', 711),\n",
       "             ('[15NH3]', 712),\n",
       "             ('[bH-]', 713),\n",
       "             ('[33P]', 714),\n",
       "             ('[Zr-2]', 715),\n",
       "             ('[15O]', 716),\n",
       "             ('[Rh-]', 717),\n",
       "             ('[PbH3]', 718),\n",
       "             ('[PH2]', 719),\n",
       "             ('[Ni-]', 720),\n",
       "             ('[CuH+]', 721),\n",
       "             ('%97', 722),\n",
       "             ('%98', 723),\n",
       "             ('%99', 724),\n",
       "             ('[Os+5]', 725),\n",
       "             ('[PtH+]', 726),\n",
       "             ('[ReH4]', 727),\n",
       "             ('[16NH]', 728),\n",
       "             ('[82Br]', 729),\n",
       "             ('[W-]', 730),\n",
       "             ('[18F-]', 731),\n",
       "             ('[15NH4+]', 732),\n",
       "             ('[Se+4]', 733),\n",
       "             ('[SeH-]', 734),\n",
       "             ('[67Cu+2]', 735),\n",
       "             ('[12C@H]', 736),\n",
       "             ('[AsH3]', 737),\n",
       "             ('[HgH]', 738),\n",
       "             ('[10B-]', 739),\n",
       "             ('[99Tc+6]', 740),\n",
       "             ('[117Sn+4]', 741),\n",
       "             ('[Te@]', 742),\n",
       "             ('[P@+]', 743),\n",
       "             ('[35SH]', 744),\n",
       "             ('[SeH+]', 745),\n",
       "             ('[Ni-2]', 746),\n",
       "             ('[Al-2]', 747),\n",
       "             ('[TeH2]', 748),\n",
       "             ('[Bh]', 749),\n",
       "             ('[99Tc+2]', 750),\n",
       "             ('[Os+8]', 751),\n",
       "             ('[PH-2]', 752),\n",
       "             ('[7Li+]', 753),\n",
       "             ('[14nH]', 754),\n",
       "             ('[AlH+2]', 755),\n",
       "             ('[18FH]', 756),\n",
       "             ('[SnH4]', 757),\n",
       "             ('[18O-2]', 758),\n",
       "             ('[IrH]', 759),\n",
       "             ('[13N]', 760),\n",
       "             ('[Te@@]', 761),\n",
       "             ('[Rh-3]', 762),\n",
       "             ('[15NH+]', 763),\n",
       "             ('[AsH3+]', 764),\n",
       "             ('[SeH2]', 765),\n",
       "             ('[AsH+]', 766),\n",
       "             ('[CoH2]', 767),\n",
       "             ('[16NH2]', 768),\n",
       "             ('[AsH-]', 769),\n",
       "             ('[203Hg+]', 770),\n",
       "             ('[P@@+]', 771),\n",
       "             ('[166Ho+3]', 772),\n",
       "             ('[60Co+3]', 773),\n",
       "             ('[13CH2-]', 774),\n",
       "             ('[SeH2+]', 775),\n",
       "             ('[75Br]', 776),\n",
       "             ('[TlH2]', 777),\n",
       "             ('[80Br]', 778),\n",
       "             ('[siH+]', 779),\n",
       "             ('[Ca+]', 780),\n",
       "             ('[153Sm+3]', 781),\n",
       "             ('[PdH]', 782),\n",
       "             ('[225Ac]', 783),\n",
       "             ('[13CH3-]', 784),\n",
       "             ('[AlH4-]', 785),\n",
       "             ('[FeH]', 786),\n",
       "             ('[13CH-]', 787),\n",
       "             ('[14C-]', 788),\n",
       "             ('[11C-]', 789),\n",
       "             ('[153Sm]', 790),\n",
       "             ('[Re-]', 791),\n",
       "             ('[te+]', 792),\n",
       "             ('[13CH4]', 793),\n",
       "             ('[ClH+2]', 794),\n",
       "             ('[8CH2]', 795),\n",
       "             ('[99Mo]', 796),\n",
       "             ('[ClH3+3]', 797),\n",
       "             ('[SbH3]', 798),\n",
       "             ('[25Mg+2]', 799),\n",
       "             ('[16N+]', 800),\n",
       "             ('[SnH2+]', 801),\n",
       "             ('[11C@H]', 802),\n",
       "             ('[122I]', 803),\n",
       "             ('[Re-2]', 804),\n",
       "             ('[RuH2+2]', 805),\n",
       "             ('[ZrH]', 806),\n",
       "             ('[Bi-]', 807),\n",
       "             ('[Pr+]', 808),\n",
       "             ('[Rn]', 809),\n",
       "             ('[Fr]', 810),\n",
       "             ('[36Cl]', 811),\n",
       "             ('[18o]', 812),\n",
       "             ('[YH]', 813),\n",
       "             ('[79Br]', 814),\n",
       "             ('[121I]', 815),\n",
       "             ('[113In+3]', 816),\n",
       "             ('[TaH]', 817),\n",
       "             ('[RhH2]', 818),\n",
       "             ('[Ta-]', 819),\n",
       "             ('[67Ga]', 820),\n",
       "             ('[ZnH+]', 821),\n",
       "             ('[SnH2-]', 822),\n",
       "             ('[OsH2]', 823),\n",
       "             ('[16F]', 824),\n",
       "             ('[FeH2]', 825),\n",
       "             ('[14O]', 826),\n",
       "             ('[PbH2+2]', 827),\n",
       "             ('[BH2]', 828),\n",
       "             ('[6H]', 829),\n",
       "             ('[125Te]', 830),\n",
       "             ('[197Hg]', 831),\n",
       "             ('[TaH2]', 832),\n",
       "             ('[TaH3]', 833),\n",
       "             ('[76As]', 834),\n",
       "             ('[Nb-2]', 835),\n",
       "             ('[14N+]', 836),\n",
       "             ('[125I-]', 837),\n",
       "             ('[33S]', 838),\n",
       "             ('[IH2+2]', 839),\n",
       "             ('[NH2]', 840),\n",
       "             ('[PtH2]', 841),\n",
       "             ('[MnH]', 842),\n",
       "             ('[19C]', 843),\n",
       "             ('[17F]', 844),\n",
       "             ('[1H-]', 845),\n",
       "             ('[SnH4+2]', 846),\n",
       "             ('[Mn-2]', 847),\n",
       "             ('[15NH2+]', 848),\n",
       "             ('[TiH2]', 849),\n",
       "             ('[ReH7]', 850),\n",
       "             ('[Cd-2]', 851),\n",
       "             ('[Fe-3]', 852),\n",
       "             ('[SH2]', 853),\n",
       "             ('[17O-]', 854),\n",
       "             ('[siH-]', 855),\n",
       "             ('[CoH+]', 856),\n",
       "             ('[VH]', 857),\n",
       "             ('[10BH]', 858),\n",
       "             ('[Ru-3]', 859),\n",
       "             ('[13O]', 860),\n",
       "             ('[5H]', 861),\n",
       "             ('[15n-]', 862),\n",
       "             ('[153Gd]', 863),\n",
       "             ('[12C@]', 864),\n",
       "             ('[11CH3-]', 865),\n",
       "             ('[IrH3]', 866),\n",
       "             ('[RuH3]', 867),\n",
       "             ('[74Se]', 868),\n",
       "             ('[Se@]', 869),\n",
       "             ('[Hf+]', 870),\n",
       "             ('[77Se]', 871),\n",
       "             ('[166Ho]', 872),\n",
       "             ('[59Fe+2]', 873),\n",
       "             ('[203Hg]', 874),\n",
       "             ('[18OH-]', 875),\n",
       "             ('[8CH]', 876),\n",
       "             ('[12C@@]', 877),\n",
       "             ('[11CH4]', 878),\n",
       "             ('[15C]', 879),\n",
       "             ('[249Cf]', 880),\n",
       "             ('[PbH4]', 881),\n",
       "             ('[64Zn]', 882),\n",
       "             ('[99Tc+]', 883),\n",
       "             ('[14c-]', 884),\n",
       "             ('[149Pm]', 885),\n",
       "             ('[IrH4]', 886),\n",
       "             ('[Se@@]', 887),\n",
       "             ('[13OH]', 888),\n",
       "             ('[14CH3-]', 889),\n",
       "             ('[28Si]', 890),\n",
       "             ('[Rh-2]', 891),\n",
       "             ('[Fe-2]', 892),\n",
       "             ('[131I-]', 893),\n",
       "             ('[51Cr]', 894),\n",
       "             ('[62Cu+2]', 895),\n",
       "             ('[81Br]', 896),\n",
       "             ('[121Sb]', 897),\n",
       "             ('[7Li]', 898),\n",
       "             ('[89Zr+4]', 899),\n",
       "             ('[SbH3+]', 900),\n",
       "             ('[11C@@H]', 901),\n",
       "             ('[98Tc]', 902),\n",
       "             ('[59Fe+3]', 903),\n",
       "             ('[BiH2+]', 904),\n",
       "             ('[SbH+]', 905),\n",
       "             ('[TiH]', 906),\n",
       "             ('[14NH3]', 907),\n",
       "             ('[15OH]', 908),\n",
       "             ('[119Sn]', 909),\n",
       "             ('[201Hg]', 910),\n",
       "             ('[MnH+]', 911),\n",
       "             ('[201Tl]', 912),\n",
       "             ('[51Cr+3]', 913),\n",
       "             ('[123I-]', 914),\n",
       "             ('[MoH]', 915),\n",
       "             ('[AlH6-3]', 916),\n",
       "             ('[MnH2]', 917),\n",
       "             ('[WH3]', 918),\n",
       "             ('[213Bi+3]', 919),\n",
       "             ('[SnH2+2]', 920),\n",
       "             ('[123IH]', 921),\n",
       "             ('[13CH+]', 922),\n",
       "             ('[Zr-]', 923),\n",
       "             ('[74As]', 924),\n",
       "             ('[13C+]', 925),\n",
       "             ('[32P+]', 926),\n",
       "             ('[KrH]', 927),\n",
       "             ('[SiH+2]', 928),\n",
       "             ('[ClH3+2]', 929),\n",
       "             ('[13NH]', 930),\n",
       "             ('[9CH2]', 931),\n",
       "             ('[ZrH2+2]', 932),\n",
       "             ('[87Sr+2]', 933),\n",
       "             ('[35s]', 934),\n",
       "             ('[239Pu]', 935),\n",
       "             ('[198Au]', 936),\n",
       "             ('[241Am]', 937),\n",
       "             ('[203Hg+2]', 938),\n",
       "             ('[V+]', 939),\n",
       "             ('[YH2]', 940),\n",
       "             ('[195Pt]', 941),\n",
       "             ('[203Pb]', 942),\n",
       "             ('[RuH4]', 943),\n",
       "             ('[ThH2]', 944),\n",
       "             ('[AuH]', 945),\n",
       "             ('[66Ga+3]', 946),\n",
       "             ('[11B-]', 947),\n",
       "             ('[F]', 948),\n",
       "             ('[24Na+]', 949),\n",
       "             ('[85Sr+2]', 950),\n",
       "             ('[201Tl+]', 951),\n",
       "             ('[14CH4]', 952),\n",
       "             ('[32S]', 953),\n",
       "             ('[TeH2+]', 954),\n",
       "             ('[ClH2+3]', 955),\n",
       "             ('[AgH]', 956),\n",
       "             ('[Ge@H]', 957),\n",
       "             ('[44Ca+2]', 958),\n",
       "             ('[Os-]', 959),\n",
       "             ('[31P]', 960),\n",
       "             ('[15nH+]', 961),\n",
       "             ('[SbH4]', 962),\n",
       "             ('[TiH+]', 963),\n",
       "             ('[Ba+]', 964),\n",
       "             ('[57Co+2]', 965),\n",
       "             ('[Ta+]', 966),\n",
       "             ('[125IH]', 967),\n",
       "             ('[77As]', 968),\n",
       "             ('[129I]', 969),\n",
       "             ('[Fe-4]', 970),\n",
       "             ('[Ta-2]', 971),\n",
       "             ('[19O]', 972),\n",
       "             ('[12O]', 973),\n",
       "             ('[BiH3]', 974),\n",
       "             ('[237Np]', 975),\n",
       "             ('[252Cf]', 976),\n",
       "             ('[86Y]', 977),\n",
       "             ('[Cr-2]', 978),\n",
       "             ('[89Y]', 979),\n",
       "             ('[195Pt+2]', 980),\n",
       "             ('[si+2]', 981),\n",
       "             ('[58Fe+2]', 982),\n",
       "             ('[Hs]', 983),\n",
       "             ('[S@@H]', 984),\n",
       "             ('[8CH4]', 985),\n",
       "             ('[164Dy+3]', 986),\n",
       "             ('[47Ca+2]', 987),\n",
       "             ('[57Co]', 988),\n",
       "             ('[NbH2]', 989),\n",
       "             ('[ReH2]', 990),\n",
       "             ('[ZnH2]', 991),\n",
       "             ('[CrH2]', 992),\n",
       "             ('[17NH]', 993),\n",
       "             ('[ZrH3]', 994),\n",
       "             ('[RhH3]', 995),\n",
       "             ('[12C-]', 996),\n",
       "             ('[18O+]', 997),\n",
       "             ('[Bi-2]', 998),\n",
       "             ('[ClH4+3]', 999),\n",
       "             ...])"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from tokenizer.tokenizer import MolTranBertTokenizer\n",
    "\n",
    "tokenizer = MolTranBertTokenizer('bert_vocab.txt')\n",
    "tokenizer.vocab"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "`fused_weight_gradient_mlp_cuda` module not found. gradient accumulation fusion with weight gradient computation disabled.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using Rotation Embedding\n",
      "Using Rotation Embedding\n",
      "Using Rotation Embedding\n",
      "Using Rotation Embedding\n",
      "Using Rotation Embedding\n",
      "Using Rotation Embedding\n",
      "Using Rotation Embedding\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Global seed set to 12345\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using Rotation Embedding\n",
      "Using Rotation Embedding\n",
      "Using Rotation Embedding\n",
      "Using Rotation Embedding\n",
      "Using Rotation Embedding\n",
      "Using Rotation Embedding\n",
      "Using Rotation Embedding\n",
      "Using Rotation Embedding\n",
      "Using Rotation Embedding\n",
      "Using Rotation Embedding\n",
      "Using Rotation Embedding\n",
      "Using Rotation Embedding\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Global seed set to 12345\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using Rotation Embedding\n",
      "Using Rotation Embedding\n",
      "Using Rotation Embedding\n",
      "Using Rotation Embedding\n",
      "Using Rotation Embedding\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "LightningModule(\n",
       "  (tok_emb): Embedding(2362, 768)\n",
       "  (drop): Dropout(p=0.2, inplace=False)\n",
       "  (blocks): TransformerEncoder(\n",
       "    (layers): ModuleList(\n",
       "      (0): TransformerEncoderLayer(\n",
       "        (attention): RotateAttentionLayer(\n",
       "          (inner_attention): LinearAttention(\n",
       "            (feature_map): GeneralizedRandomFeatures()\n",
       "          )\n",
       "          (query_projection): Linear(in_features=768, out_features=768, bias=True)\n",
       "          (key_projection): Linear(in_features=768, out_features=768, bias=True)\n",
       "          (value_projection): Linear(in_features=768, out_features=768, bias=True)\n",
       "          (out_projection): Linear(in_features=768, out_features=768, bias=True)\n",
       "          (rotaryemb): RotaryEmbedding()\n",
       "        )\n",
       "        (linear1): Linear(in_features=768, out_features=768, bias=True)\n",
       "        (linear2): Linear(in_features=768, out_features=768, bias=True)\n",
       "        (norm1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        (norm2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        (dropout): Dropout(p=0.1, inplace=False)\n",
       "      )\n",
       "      (1): TransformerEncoderLayer(\n",
       "        (attention): RotateAttentionLayer(\n",
       "          (inner_attention): LinearAttention(\n",
       "            (feature_map): GeneralizedRandomFeatures()\n",
       "          )\n",
       "          (query_projection): Linear(in_features=768, out_features=768, bias=True)\n",
       "          (key_projection): Linear(in_features=768, out_features=768, bias=True)\n",
       "          (value_projection): Linear(in_features=768, out_features=768, bias=True)\n",
       "          (out_projection): Linear(in_features=768, out_features=768, bias=True)\n",
       "          (rotaryemb): RotaryEmbedding()\n",
       "        )\n",
       "        (linear1): Linear(in_features=768, out_features=768, bias=True)\n",
       "        (linear2): Linear(in_features=768, out_features=768, bias=True)\n",
       "        (norm1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        (norm2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        (dropout): Dropout(p=0.1, inplace=False)\n",
       "      )\n",
       "      (2): TransformerEncoderLayer(\n",
       "        (attention): RotateAttentionLayer(\n",
       "          (inner_attention): LinearAttention(\n",
       "            (feature_map): GeneralizedRandomFeatures()\n",
       "          )\n",
       "          (query_projection): Linear(in_features=768, out_features=768, bias=True)\n",
       "          (key_projection): Linear(in_features=768, out_features=768, bias=True)\n",
       "          (value_projection): Linear(in_features=768, out_features=768, bias=True)\n",
       "          (out_projection): Linear(in_features=768, out_features=768, bias=True)\n",
       "          (rotaryemb): RotaryEmbedding()\n",
       "        )\n",
       "        (linear1): Linear(in_features=768, out_features=768, bias=True)\n",
       "        (linear2): Linear(in_features=768, out_features=768, bias=True)\n",
       "        (norm1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        (norm2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        (dropout): Dropout(p=0.1, inplace=False)\n",
       "      )\n",
       "      (3): TransformerEncoderLayer(\n",
       "        (attention): RotateAttentionLayer(\n",
       "          (inner_attention): LinearAttention(\n",
       "            (feature_map): GeneralizedRandomFeatures()\n",
       "          )\n",
       "          (query_projection): Linear(in_features=768, out_features=768, bias=True)\n",
       "          (key_projection): Linear(in_features=768, out_features=768, bias=True)\n",
       "          (value_projection): Linear(in_features=768, out_features=768, bias=True)\n",
       "          (out_projection): Linear(in_features=768, out_features=768, bias=True)\n",
       "          (rotaryemb): RotaryEmbedding()\n",
       "        )\n",
       "        (linear1): Linear(in_features=768, out_features=768, bias=True)\n",
       "        (linear2): Linear(in_features=768, out_features=768, bias=True)\n",
       "        (norm1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        (norm2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        (dropout): Dropout(p=0.1, inplace=False)\n",
       "      )\n",
       "      (4): TransformerEncoderLayer(\n",
       "        (attention): RotateAttentionLayer(\n",
       "          (inner_attention): LinearAttention(\n",
       "            (feature_map): GeneralizedRandomFeatures()\n",
       "          )\n",
       "          (query_projection): Linear(in_features=768, out_features=768, bias=True)\n",
       "          (key_projection): Linear(in_features=768, out_features=768, bias=True)\n",
       "          (value_projection): Linear(in_features=768, out_features=768, bias=True)\n",
       "          (out_projection): Linear(in_features=768, out_features=768, bias=True)\n",
       "          (rotaryemb): RotaryEmbedding()\n",
       "        )\n",
       "        (linear1): Linear(in_features=768, out_features=768, bias=True)\n",
       "        (linear2): Linear(in_features=768, out_features=768, bias=True)\n",
       "        (norm1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        (norm2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        (dropout): Dropout(p=0.1, inplace=False)\n",
       "      )\n",
       "      (5): TransformerEncoderLayer(\n",
       "        (attention): RotateAttentionLayer(\n",
       "          (inner_attention): LinearAttention(\n",
       "            (feature_map): GeneralizedRandomFeatures()\n",
       "          )\n",
       "          (query_projection): Linear(in_features=768, out_features=768, bias=True)\n",
       "          (key_projection): Linear(in_features=768, out_features=768, bias=True)\n",
       "          (value_projection): Linear(in_features=768, out_features=768, bias=True)\n",
       "          (out_projection): Linear(in_features=768, out_features=768, bias=True)\n",
       "          (rotaryemb): RotaryEmbedding()\n",
       "        )\n",
       "        (linear1): Linear(in_features=768, out_features=768, bias=True)\n",
       "        (linear2): Linear(in_features=768, out_features=768, bias=True)\n",
       "        (norm1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        (norm2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        (dropout): Dropout(p=0.1, inplace=False)\n",
       "      )\n",
       "      (6): TransformerEncoderLayer(\n",
       "        (attention): RotateAttentionLayer(\n",
       "          (inner_attention): LinearAttention(\n",
       "            (feature_map): GeneralizedRandomFeatures()\n",
       "          )\n",
       "          (query_projection): Linear(in_features=768, out_features=768, bias=True)\n",
       "          (key_projection): Linear(in_features=768, out_features=768, bias=True)\n",
       "          (value_projection): Linear(in_features=768, out_features=768, bias=True)\n",
       "          (out_projection): Linear(in_features=768, out_features=768, bias=True)\n",
       "          (rotaryemb): RotaryEmbedding()\n",
       "        )\n",
       "        (linear1): Linear(in_features=768, out_features=768, bias=True)\n",
       "        (linear2): Linear(in_features=768, out_features=768, bias=True)\n",
       "        (norm1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        (norm2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        (dropout): Dropout(p=0.1, inplace=False)\n",
       "      )\n",
       "      (7): TransformerEncoderLayer(\n",
       "        (attention): RotateAttentionLayer(\n",
       "          (inner_attention): LinearAttention(\n",
       "            (feature_map): GeneralizedRandomFeatures()\n",
       "          )\n",
       "          (query_projection): Linear(in_features=768, out_features=768, bias=True)\n",
       "          (key_projection): Linear(in_features=768, out_features=768, bias=True)\n",
       "          (value_projection): Linear(in_features=768, out_features=768, bias=True)\n",
       "          (out_projection): Linear(in_features=768, out_features=768, bias=True)\n",
       "          (rotaryemb): RotaryEmbedding()\n",
       "        )\n",
       "        (linear1): Linear(in_features=768, out_features=768, bias=True)\n",
       "        (linear2): Linear(in_features=768, out_features=768, bias=True)\n",
       "        (norm1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        (norm2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        (dropout): Dropout(p=0.1, inplace=False)\n",
       "      )\n",
       "      (8): TransformerEncoderLayer(\n",
       "        (attention): RotateAttentionLayer(\n",
       "          (inner_attention): LinearAttention(\n",
       "            (feature_map): GeneralizedRandomFeatures()\n",
       "          )\n",
       "          (query_projection): Linear(in_features=768, out_features=768, bias=True)\n",
       "          (key_projection): Linear(in_features=768, out_features=768, bias=True)\n",
       "          (value_projection): Linear(in_features=768, out_features=768, bias=True)\n",
       "          (out_projection): Linear(in_features=768, out_features=768, bias=True)\n",
       "          (rotaryemb): RotaryEmbedding()\n",
       "        )\n",
       "        (linear1): Linear(in_features=768, out_features=768, bias=True)\n",
       "        (linear2): Linear(in_features=768, out_features=768, bias=True)\n",
       "        (norm1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        (norm2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        (dropout): Dropout(p=0.1, inplace=False)\n",
       "      )\n",
       "      (9): TransformerEncoderLayer(\n",
       "        (attention): RotateAttentionLayer(\n",
       "          (inner_attention): LinearAttention(\n",
       "            (feature_map): GeneralizedRandomFeatures()\n",
       "          )\n",
       "          (query_projection): Linear(in_features=768, out_features=768, bias=True)\n",
       "          (key_projection): Linear(in_features=768, out_features=768, bias=True)\n",
       "          (value_projection): Linear(in_features=768, out_features=768, bias=True)\n",
       "          (out_projection): Linear(in_features=768, out_features=768, bias=True)\n",
       "          (rotaryemb): RotaryEmbedding()\n",
       "        )\n",
       "        (linear1): Linear(in_features=768, out_features=768, bias=True)\n",
       "        (linear2): Linear(in_features=768, out_features=768, bias=True)\n",
       "        (norm1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        (norm2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        (dropout): Dropout(p=0.1, inplace=False)\n",
       "      )\n",
       "      (10): TransformerEncoderLayer(\n",
       "        (attention): RotateAttentionLayer(\n",
       "          (inner_attention): LinearAttention(\n",
       "            (feature_map): GeneralizedRandomFeatures()\n",
       "          )\n",
       "          (query_projection): Linear(in_features=768, out_features=768, bias=True)\n",
       "          (key_projection): Linear(in_features=768, out_features=768, bias=True)\n",
       "          (value_projection): Linear(in_features=768, out_features=768, bias=True)\n",
       "          (out_projection): Linear(in_features=768, out_features=768, bias=True)\n",
       "          (rotaryemb): RotaryEmbedding()\n",
       "        )\n",
       "        (linear1): Linear(in_features=768, out_features=768, bias=True)\n",
       "        (linear2): Linear(in_features=768, out_features=768, bias=True)\n",
       "        (norm1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        (norm2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        (dropout): Dropout(p=0.1, inplace=False)\n",
       "      )\n",
       "      (11): TransformerEncoderLayer(\n",
       "        (attention): RotateAttentionLayer(\n",
       "          (inner_attention): LinearAttention(\n",
       "            (feature_map): GeneralizedRandomFeatures()\n",
       "          )\n",
       "          (query_projection): Linear(in_features=768, out_features=768, bias=True)\n",
       "          (key_projection): Linear(in_features=768, out_features=768, bias=True)\n",
       "          (value_projection): Linear(in_features=768, out_features=768, bias=True)\n",
       "          (out_projection): Linear(in_features=768, out_features=768, bias=True)\n",
       "          (rotaryemb): RotaryEmbedding()\n",
       "        )\n",
       "        (linear1): Linear(in_features=768, out_features=768, bias=True)\n",
       "        (linear2): Linear(in_features=768, out_features=768, bias=True)\n",
       "        (norm1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        (norm2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        (dropout): Dropout(p=0.1, inplace=False)\n",
       "      )\n",
       "    )\n",
       "    (norm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "  )\n",
       "  (lang_model): lm_layer(\n",
       "    (embed): Linear(in_features=768, out_features=768, bias=True)\n",
       "    (ln_f): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "    (head): Linear(in_features=768, out_features=2362, bias=False)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from train_pubchem_light import LightningModule\n",
    "\n",
    "ckpt = '../../data/Pretrained MoLFormer/checkpoints/N-Step-Checkpoint_3_30000.ckpt'\n",
    "lm = LightningModule(config, tokenizer.vocab).load_from_checkpoint(ckpt, config=config, vocab=tokenizer.vocab)\n",
    "lm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run inference (get embeddings)\n",
    "\n",
    "note: this runs on CPU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from fast_transformers.masking import LengthMask as LM\n",
    "\n",
    "def batch_split(data, batch_size=64):\n",
    "    i = 0\n",
    "    while i < len(data):\n",
    "        yield data[i:min(i+batch_size, len(data))]\n",
    "        i += batch_size\n",
    "\n",
    "def embed(model, smiles, tokenizer, batch_size=64):\n",
    "    model.eval()\n",
    "    embeddings = []\n",
    "    for batch in batch_split(smiles, batch_size=batch_size):\n",
    "        batch_enc = tokenizer.batch_encode_plus(batch, padding=True, add_special_tokens=True)\n",
    "        idx, mask = torch.tensor(batch_enc['input_ids']), torch.tensor(batch_enc['attention_mask'])\n",
    "        with torch.no_grad():\n",
    "            token_embeddings = model.blocks(model.tok_emb(idx), length_mask=LM(mask.sum(-1)))\n",
    "        # average pooling over tokens\n",
    "        input_mask_expanded = mask.unsqueeze(-1).expand(token_embeddings.size()).float()\n",
    "        sum_embeddings = torch.sum(token_embeddings * input_mask_expanded, 1)\n",
    "        sum_mask = torch.clamp(input_mask_expanded.sum(1), min=1e-9)\n",
    "        embedding = sum_embeddings / sum_mask\n",
    "        embeddings.append(embedding.detach().cpu())\n",
    "    return torch.cat(embeddings)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use linear head for classification task\n",
    "\n",
    "There are many ways to use the embeddings for downstream tasks, this is clearly a toy example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>smiles</th>\n",
       "      <th>CID</th>\n",
       "      <th>Class</th>\n",
       "      <th>Unnamed: 3</th>\n",
       "      <th>pIC50</th>\n",
       "      <th>MW</th>\n",
       "      <th>AlogP</th>\n",
       "      <th>HBA</th>\n",
       "      <th>HBD</th>\n",
       "      <th>RB</th>\n",
       "      <th>...</th>\n",
       "      <th>PEOE6 (PEOE6)</th>\n",
       "      <th>PEOE7 (PEOE7)</th>\n",
       "      <th>PEOE8 (PEOE8)</th>\n",
       "      <th>PEOE9 (PEOE9)</th>\n",
       "      <th>PEOE10 (PEOE10)</th>\n",
       "      <th>PEOE11 (PEOE11)</th>\n",
       "      <th>PEOE12 (PEOE12)</th>\n",
       "      <th>PEOE13 (PEOE13)</th>\n",
       "      <th>PEOE14 (PEOE14)</th>\n",
       "      <th>canvasUID</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>O1CC[C@@H](NC(=O)[C@@H](Cc2cc3cc(ccc3nc2N)-c2c...</td>\n",
       "      <td>BACE_1</td>\n",
       "      <td>1</td>\n",
       "      <td>NaN</td>\n",
       "      <td>9.154901</td>\n",
       "      <td>431.56979</td>\n",
       "      <td>4.4014</td>\n",
       "      <td>3</td>\n",
       "      <td>2</td>\n",
       "      <td>5</td>\n",
       "      <td>...</td>\n",
       "      <td>53.205711</td>\n",
       "      <td>78.640335</td>\n",
       "      <td>226.855410</td>\n",
       "      <td>107.434910</td>\n",
       "      <td>37.133846</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>7.980170</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>S1(=O)(=O)N(c2cc(cc3c2n(cc3CC)CC1)C(=O)N[C@H](...</td>\n",
       "      <td>BACE_3</td>\n",
       "      <td>1</td>\n",
       "      <td>NaN</td>\n",
       "      <td>8.698970</td>\n",
       "      <td>591.74091</td>\n",
       "      <td>2.5499</td>\n",
       "      <td>4</td>\n",
       "      <td>3</td>\n",
       "      <td>11</td>\n",
       "      <td>...</td>\n",
       "      <td>70.365707</td>\n",
       "      <td>47.941147</td>\n",
       "      <td>192.406520</td>\n",
       "      <td>255.752550</td>\n",
       "      <td>23.654478</td>\n",
       "      <td>0.230159</td>\n",
       "      <td>15.879790</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>24.663788</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>S1(=O)(=O)N(c2cc(cc3c2n(cc3CC)CC1)C(=O)N[C@H](...</td>\n",
       "      <td>BACE_5</td>\n",
       "      <td>1</td>\n",
       "      <td>NaN</td>\n",
       "      <td>8.698970</td>\n",
       "      <td>629.71283</td>\n",
       "      <td>3.5086</td>\n",
       "      <td>3</td>\n",
       "      <td>3</td>\n",
       "      <td>11</td>\n",
       "      <td>...</td>\n",
       "      <td>78.945702</td>\n",
       "      <td>39.361153</td>\n",
       "      <td>179.712880</td>\n",
       "      <td>220.461300</td>\n",
       "      <td>23.654478</td>\n",
       "      <td>0.230159</td>\n",
       "      <td>15.879790</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>26.100143</td>\n",
       "      <td>5</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>S(=O)(=O)(CCCCC)C[C@@H](NC(=O)c1cccnc1)C(=O)N[...</td>\n",
       "      <td>BACE_7</td>\n",
       "      <td>1</td>\n",
       "      <td>NaN</td>\n",
       "      <td>8.698970</td>\n",
       "      <td>645.78009</td>\n",
       "      <td>3.1973</td>\n",
       "      <td>5</td>\n",
       "      <td>4</td>\n",
       "      <td>18</td>\n",
       "      <td>...</td>\n",
       "      <td>63.830162</td>\n",
       "      <td>52.390511</td>\n",
       "      <td>263.781340</td>\n",
       "      <td>190.542130</td>\n",
       "      <td>45.370659</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>23.859961</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>24.663788</td>\n",
       "      <td>7</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>O1c2c(cc(cc2)CC)[C@@H]([NH2+]C[C@@H](O)[C@H]2N...</td>\n",
       "      <td>BACE_9</td>\n",
       "      <td>1</td>\n",
       "      <td>NaN</td>\n",
       "      <td>8.602060</td>\n",
       "      <td>556.71503</td>\n",
       "      <td>4.7010</td>\n",
       "      <td>4</td>\n",
       "      <td>3</td>\n",
       "      <td>5</td>\n",
       "      <td>...</td>\n",
       "      <td>53.205711</td>\n",
       "      <td>68.418541</td>\n",
       "      <td>299.000030</td>\n",
       "      <td>140.683620</td>\n",
       "      <td>28.755558</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>15.879790</td>\n",
       "      <td>6.904104</td>\n",
       "      <td>24.663788</td>\n",
       "      <td>9</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1204</th>\n",
       "      <td>Clc1cc2nc([nH]c2cc1)N</td>\n",
       "      <td>BACE_1541</td>\n",
       "      <td>0</td>\n",
       "      <td>NaN</td>\n",
       "      <td>3.136083</td>\n",
       "      <td>167.59560</td>\n",
       "      <td>1.4715</td>\n",
       "      <td>1</td>\n",
       "      <td>2</td>\n",
       "      <td>0</td>\n",
       "      <td>...</td>\n",
       "      <td>26.907076</td>\n",
       "      <td>25.739992</td>\n",
       "      <td>1.916970</td>\n",
       "      <td>40.882919</td>\n",
       "      <td>20.071724</td>\n",
       "      <td>19.404020</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>1541</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1205</th>\n",
       "      <td>Clc1cc2nc(n(c2cc1)C(CC(=O)NCC1CCOCC1)CC)N</td>\n",
       "      <td>BACE_1543</td>\n",
       "      <td>0</td>\n",
       "      <td>NaN</td>\n",
       "      <td>3.000000</td>\n",
       "      <td>364.86969</td>\n",
       "      <td>2.5942</td>\n",
       "      <td>3</td>\n",
       "      <td>2</td>\n",
       "      <td>6</td>\n",
       "      <td>...</td>\n",
       "      <td>37.212799</td>\n",
       "      <td>37.681076</td>\n",
       "      <td>180.226410</td>\n",
       "      <td>95.670128</td>\n",
       "      <td>30.107586</td>\n",
       "      <td>9.368159</td>\n",
       "      <td>7.980170</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>1543</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1206</th>\n",
       "      <td>Clc1cc2nc(n(c2cc1)C(CC(=O)NCc1ncccc1)CC)N</td>\n",
       "      <td>BACE_1544</td>\n",
       "      <td>0</td>\n",
       "      <td>NaN</td>\n",
       "      <td>3.000000</td>\n",
       "      <td>357.83731</td>\n",
       "      <td>2.8229</td>\n",
       "      <td>3</td>\n",
       "      <td>2</td>\n",
       "      <td>6</td>\n",
       "      <td>...</td>\n",
       "      <td>45.792797</td>\n",
       "      <td>47.349350</td>\n",
       "      <td>122.401500</td>\n",
       "      <td>99.877144</td>\n",
       "      <td>30.107586</td>\n",
       "      <td>9.368159</td>\n",
       "      <td>7.980170</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>1544</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1207</th>\n",
       "      <td>Brc1cc(ccc1)C1CC1C=1N=C(N)N(C)C(=O)C=1</td>\n",
       "      <td>BACE_1545</td>\n",
       "      <td>0</td>\n",
       "      <td>NaN</td>\n",
       "      <td>2.953115</td>\n",
       "      <td>320.18451</td>\n",
       "      <td>3.0895</td>\n",
       "      <td>2</td>\n",
       "      <td>1</td>\n",
       "      <td>2</td>\n",
       "      <td>...</td>\n",
       "      <td>47.790600</td>\n",
       "      <td>22.563574</td>\n",
       "      <td>96.290794</td>\n",
       "      <td>58.798935</td>\n",
       "      <td>20.071724</td>\n",
       "      <td>9.368159</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>6.904104</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>1545</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1208</th>\n",
       "      <td>O=C1N(C)C(=NC(=C1)C1CC1c1cc(ccc1)-c1ccccc1)N</td>\n",
       "      <td>BACE_1546</td>\n",
       "      <td>0</td>\n",
       "      <td>NaN</td>\n",
       "      <td>2.733298</td>\n",
       "      <td>317.38440</td>\n",
       "      <td>3.8595</td>\n",
       "      <td>2</td>\n",
       "      <td>1</td>\n",
       "      <td>3</td>\n",
       "      <td>...</td>\n",
       "      <td>77.219978</td>\n",
       "      <td>9.316234</td>\n",
       "      <td>95.907784</td>\n",
       "      <td>112.609720</td>\n",
       "      <td>20.071724</td>\n",
       "      <td>9.368159</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>6.904104</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>1546</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>1209 rows × 595 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                 smiles        CID  Class  \\\n",
       "0     O1CC[C@@H](NC(=O)[C@@H](Cc2cc3cc(ccc3nc2N)-c2c...     BACE_1      1   \n",
       "1     S1(=O)(=O)N(c2cc(cc3c2n(cc3CC)CC1)C(=O)N[C@H](...     BACE_3      1   \n",
       "2     S1(=O)(=O)N(c2cc(cc3c2n(cc3CC)CC1)C(=O)N[C@H](...     BACE_5      1   \n",
       "3     S(=O)(=O)(CCCCC)C[C@@H](NC(=O)c1cccnc1)C(=O)N[...     BACE_7      1   \n",
       "4     O1c2c(cc(cc2)CC)[C@@H]([NH2+]C[C@@H](O)[C@H]2N...     BACE_9      1   \n",
       "...                                                 ...        ...    ...   \n",
       "1204                              Clc1cc2nc([nH]c2cc1)N  BACE_1541      0   \n",
       "1205          Clc1cc2nc(n(c2cc1)C(CC(=O)NCC1CCOCC1)CC)N  BACE_1543      0   \n",
       "1206          Clc1cc2nc(n(c2cc1)C(CC(=O)NCc1ncccc1)CC)N  BACE_1544      0   \n",
       "1207             Brc1cc(ccc1)C1CC1C=1N=C(N)N(C)C(=O)C=1  BACE_1545      0   \n",
       "1208       O=C1N(C)C(=NC(=C1)C1CC1c1cc(ccc1)-c1ccccc1)N  BACE_1546      0   \n",
       "\n",
       "      Unnamed: 3     pIC50         MW   AlogP  HBA  HBD  RB  ...  \\\n",
       "0            NaN  9.154901  431.56979  4.4014    3    2   5  ...   \n",
       "1            NaN  8.698970  591.74091  2.5499    4    3  11  ...   \n",
       "2            NaN  8.698970  629.71283  3.5086    3    3  11  ...   \n",
       "3            NaN  8.698970  645.78009  3.1973    5    4  18  ...   \n",
       "4            NaN  8.602060  556.71503  4.7010    4    3   5  ...   \n",
       "...          ...       ...        ...     ...  ...  ...  ..  ...   \n",
       "1204         NaN  3.136083  167.59560  1.4715    1    2   0  ...   \n",
       "1205         NaN  3.000000  364.86969  2.5942    3    2   6  ...   \n",
       "1206         NaN  3.000000  357.83731  2.8229    3    2   6  ...   \n",
       "1207         NaN  2.953115  320.18451  3.0895    2    1   2  ...   \n",
       "1208         NaN  2.733298  317.38440  3.8595    2    1   3  ...   \n",
       "\n",
       "      PEOE6 (PEOE6)  PEOE7 (PEOE7)  PEOE8 (PEOE8)  PEOE9 (PEOE9)  \\\n",
       "0         53.205711      78.640335     226.855410     107.434910   \n",
       "1         70.365707      47.941147     192.406520     255.752550   \n",
       "2         78.945702      39.361153     179.712880     220.461300   \n",
       "3         63.830162      52.390511     263.781340     190.542130   \n",
       "4         53.205711      68.418541     299.000030     140.683620   \n",
       "...             ...            ...            ...            ...   \n",
       "1204      26.907076      25.739992       1.916970      40.882919   \n",
       "1205      37.212799      37.681076     180.226410      95.670128   \n",
       "1206      45.792797      47.349350     122.401500      99.877144   \n",
       "1207      47.790600      22.563574      96.290794      58.798935   \n",
       "1208      77.219978       9.316234      95.907784     112.609720   \n",
       "\n",
       "      PEOE10 (PEOE10)  PEOE11 (PEOE11)  PEOE12 (PEOE12)  PEOE13 (PEOE13)  \\\n",
       "0           37.133846         0.000000         7.980170         0.000000   \n",
       "1           23.654478         0.230159        15.879790         0.000000   \n",
       "2           23.654478         0.230159        15.879790         0.000000   \n",
       "3           45.370659         0.000000        23.859961         0.000000   \n",
       "4           28.755558         0.000000        15.879790         6.904104   \n",
       "...               ...              ...              ...              ...   \n",
       "1204        20.071724        19.404020         0.000000         0.000000   \n",
       "1205        30.107586         9.368159         7.980170         0.000000   \n",
       "1206        30.107586         9.368159         7.980170         0.000000   \n",
       "1207        20.071724         9.368159         0.000000         6.904104   \n",
       "1208        20.071724         9.368159         0.000000         6.904104   \n",
       "\n",
       "      PEOE14 (PEOE14)  canvasUID  \n",
       "0            0.000000          1  \n",
       "1           24.663788          3  \n",
       "2           26.100143          5  \n",
       "3           24.663788          7  \n",
       "4           24.663788          9  \n",
       "...               ...        ...  \n",
       "1204         0.000000       1541  \n",
       "1205         0.000000       1543  \n",
       "1206         0.000000       1544  \n",
       "1207         0.000000       1545  \n",
       "1208         0.000000       1546  \n",
       "\n",
       "[1209 rows x 595 columns]"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pandas as pd\n",
    "\n",
    "df = pd.read_csv('../../data/bace/train.csv')#.sample(frac=0.1)  # speed things up...\n",
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/dccstor/shoffman/conda-envs/MolTran_CUDA11/lib/python3.8/site-packages/sklearn/linear_model/_logistic.py:763: ConvergenceWarning: lbfgs failed to converge (status=1):\n",
      "STOP: TOTAL NO. of ITERATIONS REACHED LIMIT.\n",
      "\n",
      "Increase the number of iterations (max_iter) or scale the data as shown in:\n",
      "    https://scikit-learn.org/stable/modules/preprocessing.html\n",
      "Please also refer to the documentation for alternative solver options:\n",
      "    https://scikit-learn.org/stable/modules/linear_model.html#logistic-regression\n",
      "  n_iter_i = _check_optimize_result(\n"
     ]
    }
   ],
   "source": [
    "from rdkit import Chem\n",
    "from sklearn.linear_model import LogisticRegression\n",
    "\n",
    "def canonicalize(s):\n",
    "    return Chem.MolToSmiles(Chem.MolFromSmiles(s), canonical=True, isomericSmiles=False)\n",
    "\n",
    "smiles = df.smiles.apply(canonicalize)\n",
    "X = embed(lm, smiles, tokenizer).numpy()\n",
    "y = df.Class\n",
    "head = LogisticRegression().fit(X, y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.8518518518518519"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from sklearn.metrics import roc_auc_score\n",
    "\n",
    "test_df = pd.read_csv('../../data/bace/test.csv').sample(frac=0.1)\n",
    "X_test = embed(lm, test_df.smiles.apply(canonicalize), tokenizer).numpy()\n",
    "roc_auc_score(test_df.Class, head.predict_proba(X_test)[:, 1])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}