{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"**This notebook contains weights for a decent vae and a small dalle that needs way more training**\n\n - trained on competition data using the code from https://github.com/lucidrains/DALLE-pytorch \n\n*lucid is all you need*\n\nThis is mostly just for fun, but if you figure out how to use this to improve the score, please let everyone know, thanks.","metadata":{}},{"cell_type":"code","source":"# !conda install -y -c rdkit rdkit\n!pip install rdkit-pypi dalle-pytorch youtokentome einops","metadata":{"execution":{"iopub.status.busy":"2021-05-22T03:06:25.291780Z","iopub.execute_input":"2021-05-22T03:06:25.292183Z","iopub.status.idle":"2021-05-22T03:06:36.264368Z","shell.execute_reply.started":"2021-05-22T03:06:25.292097Z","shell.execute_reply":"2021-05-22T03:06:36.263420Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import numpy as np, os\nimport pandas as pd\nimport re\nfrom tqdm.auto import tqdm\ntqdm.pandas()\n\nimport torch\nfrom torch.utils.data import Dataset\nimport cv2\n\nfrom albumentations import (\n    Compose, OneOf, Normalize, Resize, HorizontalFlip, VerticalFlip, Rotate, RandomRotate90, CenterCrop\n    )\n\nfrom albumentations.pytorch import ToTensorV2\n\nimport rdkit.Chem as Chem\n\nfrom pathlib import Path\nfrom tqdm import tqdm\n\n# torch\n\nimport torch\n\nfrom einops import repeat\n\n# vision imports\n\nfrom PIL import Image\nfrom torchvision.utils import make_grid, save_image\n\n# dalle related classes and utils\n\nfrom dalle_pytorch import DiscreteVAE, OpenAIDiscreteVAE, VQGanVAE1024, DALLE\nfrom dalle_pytorch.tokenizer import tokenizer, HugTokenizer, YttmTokenizer, ChineseTokenizer\n\nfrom matplotlib import pyplot as plt\n\n## functions to generate images via rdkit from this excellent notebook: https://www.kaggle.com/tuckerarrants/inchi-allowed-external-data\n\ndef sp_noise(image):\n    #https://gist.github.com/lucaswiman/1e877a164a69f78694f845eab45c381a\n    output = image.copy()\n    if len(image.shape) == 2:\n        black = 0\n        white = 255            \n    else:\n        colorspace = image.shape[2]\n        if colorspace == 3:  # RGB\n            black = np.array([0, 0, 0], dtype='uint8')\n            white = np.array([255, 255, 255], dtype='uint8')\n        else:  # RGBA\n            black = np.array([0, 0, 0, 255], dtype='uint8')\n            white = np.array([255, 255, 255, 255], dtype='uint8')\n    probs = np.random.random(image.shape[:2])\n    image[probs < .00015] = black\n    image[probs > .85] = white\n    return image\n\ndef noisy_inchi(inchi, inchi_path='tmp1.png', add_noise=True, crop_and_pad=True):\n    mol = Chem.MolFromInchi(inchi)\n    d = Chem.Draw.rdMolDraw2D.MolDraw2DCairo(512, 512)\n    # https://www.kaggle.com/stainsby/improved-synthetic-data-for-bms-competition-v3\n    Chem.rdDepictor.SetPreferCoordGen(True)\n    d.drawOptions().maxFontSize=16\n    d.drawOptions().multipleBondOffset=np.random.uniform(0.05, 0.2)\n    d.drawOptions().useBWAtomPalette()\n    d.drawOptions().bondLineWidth=2\n    d.drawOptions().additionalAtomLabelPadding=np.random.uniform(0, .2)\n    d.DrawMolecule(mol)\n    d.FinishDrawing()\n    d.WriteDrawingText(inchi_path)  \n    if crop_and_pad:\n        img = cv2.imread(inchi_path, cv2.IMREAD_GRAYSCALE)\n        crop_rows = img[~np.all(img==255, axis=1), :]\n        img = crop_rows[:, ~np.all(crop_rows==255, axis=0)]\n        img = cv2.copyMakeBorder(img, 30, 20, 30, 20, cv2.BORDER_CONSTANT, value=255)\n        img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)\n    else:\n        img = cv2.imread(inchi_path)\n    if add_noise:\n        img = sp_noise(img)\n#         cv2.imwrite(inchi_path, img)\n    return img","metadata":{"execution":{"iopub.status.busy":"2021-05-22T04:20:54.371127Z","iopub.execute_input":"2021-05-22T04:20:54.371449Z","iopub.status.idle":"2021-05-22T04:20:54.387503Z","shell.execute_reply.started":"2021-05-22T04:20:54.371418Z","shell.execute_reply":"2021-05-22T04:20:54.386658Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df = pd.read_pickle('../input/extra-inchis/allowed_inchi_processed.pkl')","metadata":{"execution":{"iopub.status.busy":"2021-05-22T03:08:24.465549Z","iopub.execute_input":"2021-05-22T03:08:24.465870Z","iopub.status.idle":"2021-05-22T03:09:09.463827Z","shell.execute_reply.started":"2021-05-22T03:08:24.465840Z","shell.execute_reply":"2021-05-22T03:09:09.462808Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"rows = df.sample(n=8)[['InChI', 'InChI_text']].values","metadata":{"execution":{"iopub.status.busy":"2021-05-22T05:16:22.909582Z","iopub.execute_input":"2021-05-22T05:16:22.909932Z","iopub.status.idle":"2021-05-22T05:16:23.252353Z","shell.execute_reply.started":"2021-05-22T05:16:22.909899Z","shell.execute_reply":"2021-05-22T05:16:23.251383Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"inchis, texts = rows[:,0], rows[:, 1]","metadata":{"execution":{"iopub.status.busy":"2021-05-22T05:16:25.040602Z","iopub.execute_input":"2021-05-22T05:16:25.040928Z","iopub.status.idle":"2021-05-22T05:16:25.047185Z","shell.execute_reply.started":"2021-05-22T05:16:25.040898Z","shell.execute_reply":"2021-05-22T05:16:25.046394Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def display_multiple_img(images, rows = 1, cols=1):\n    figure, ax = plt.subplots(nrows=rows,ncols=cols )\n    for ind,title in enumerate(images):\n        ax.ravel()[ind].imshow(images[title])\n        ax.ravel()[ind].set_title(title)\n        ax.ravel()[ind].set_axis_off()\n    plt.tight_layout()\n    plt.gcf().set_size_inches(12,8)\n    plt.show()\n    \ndef show_generated(output, raw_text):\n\n    imgs = [h.cpu().permute(1,2,0).numpy()*2+1 for h in output.unbind(0)]+[noisy_inchi(raw_text), noisy_inchi(raw_text)]\n    ttls = ['dalle_'+str(i) for i in range(len(output))] + ['rdkit_1', 'rdkit_2']\n\n    images = {ttls[i]: imgs[i] for i in range(len(ttls))}\n\n    display_multiple_img(images, 2, 3)","metadata":{"execution":{"iopub.status.busy":"2021-05-22T04:28:56.081801Z","iopub.execute_input":"2021-05-22T04:28:56.082159Z","iopub.status.idle":"2021-05-22T04:28:56.091255Z","shell.execute_reply.started":"2021-05-22T04:28:56.082130Z","shell.execute_reply":"2021-05-22T04:28:56.090306Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Args:\n    dalle_path = '../input/moldal/dalle_gen2.pth'\n    text = ''\n    num_images = 4\n    batch_size = 4\n    top_k = 0.97\n    outputs_dir = './'\n    bpe_path = '../input/moldal/bms_bpe.model'\n    hug = False\n    chinese = False\n    taming = False\n\nargs=Args()\n\ndef exists(val):\n    return val is not None\n\nif exists(args.bpe_path):\n    klass = HugTokenizer if args.hug else YttmTokenizer\n    tokenizer = klass(args.bpe_path)\nelif args.chinese:\n    tokenizer = ChineseTokenizer()\n\n# load DALL-E\n\ndalle_path = Path(args.dalle_path)\n\nassert dalle_path.exists(), 'trained DALL-E must exist'\n\nload_obj = torch.load(str(dalle_path), map_location='cpu')\n\ndalle_params, vae_params, weights = load_obj.pop('hparams'), load_obj.pop('vae_params'), load_obj.pop('weights')\n\ndalle_params.pop('vae', None) # cleanup later\n\nimport gc\ngc.collect()\n\nvae = DiscreteVAE(**vae_params)\n\ndalle = DALLE(vae = vae, **dalle_params).cuda()\n\ndalle.load_state_dict(weights)\n\nimage_size = vae.image_size    ","metadata":{"execution":{"iopub.status.busy":"2021-05-22T05:16:50.378821Z","iopub.execute_input":"2021-05-22T05:16:50.379150Z","iopub.status.idle":"2021-05-22T05:16:50.869540Z","shell.execute_reply.started":"2021-05-22T05:16:50.379116Z","shell.execute_reply":"2021-05-22T05:16:50.868755Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"for inchi, text in zip(inchis, texts):\n    text = tokenizer.tokenize(text,\n                              dalle.text_seq_len,\n                              truncate_text=True,\n                             ).cuda()\n\n    text = repeat(text, '() n -> b n', b = args.num_images)\n    for text_chunk in tqdm(text.split(args.batch_size), desc = f'generating images for - {inchi}'):\n        output = dalle.generate_images(text_chunk, filter_thres = args.top_k)\n        show_generated(output, inchi)","metadata":{"execution":{"iopub.status.busy":"2021-05-22T05:17:08.822354Z","iopub.execute_input":"2021-05-22T05:17:08.822720Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# for inchi, text in zip(inchis, texts):\n#     text = tokenizer.tokenize(text,\n#                               dalle.text_seq_len,\n#                               truncate_text=True,\n#                              ).cuda()\n\n#     text = repeat(text, '() n -> b n', b = args.num_images)\n#     for text_chunk in tqdm(text.split(args.batch_size), desc = f'generating images for - {inchi}'):\n#         output = dalle.generate_images(text_chunk, filter_thres = args.top_k)\n#         show_generated(output, inchi)\n#         print(text_chunk)\n#         break\n#     break","metadata":{"execution":{"iopub.status.busy":"2021-05-22T04:30:14.797477Z","iopub.execute_input":"2021-05-22T04:30:14.797836Z","iopub.status.idle":"2021-05-22T05:13:56.489202Z","shell.execute_reply.started":"2021-05-22T04:30:14.797805Z","shell.execute_reply":"2021-05-22T05:13:56.488296Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}