Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import torch | |
| import numpy as np | |
| import gradio as gr | |
| from rdkit import Chem | |
| from rdkit.Chem import Draw, Descriptors, rdMolDescriptors, QED | |
| from PIL import Image, ImageDraw, ImageFont | |
| from llama_model import ConditionalTransformer, ModelArgs | |
| from generate_no_cache import generate as conditional_generate | |
| from transformers import PreTrainedTokenizerFast | |
| def decode(out, is_cond=False): | |
| '''Decodes and returns generations form a batch ''' | |
| out = [tokenizer.decode(o) for o in out] | |
| out = [''.join(o.split(' ')) for o in out] | |
| if is_cond: | |
| out = [o.split('[SEP]')[0] for o in out] | |
| return out | |
| else: | |
| outs = [] | |
| for o in out: | |
| outs.extend(o.split('[SEP]')[1:-1]) | |
| np.random.shuffle(outs) | |
| return outs | |
| def verify_smiles(smiles: str) -> bool: | |
| """Returns True if the SMILES is valid, False otherwise.""" | |
| try: | |
| mol = Chem.MolFromSmiles(smiles) | |
| return mol is not None | |
| except: | |
| return False | |
| def generate_mol_background_color(mol) -> tuple: | |
| logp = Descriptors.MolLogP(mol) | |
| h_donors = Descriptors.NumHDonors(mol) | |
| mol_wt = Descriptors.MolWt(mol) | |
| red = min(int(logp / 6 * 255), 255) # logP: -3 to +3 → 0–255 | |
| green = min(h_donors * 30, 255) # H-donors: 0–8 | |
| blue = min(int(mol_wt / 500 * 255), 255) # MW: ~0–500 | |
| return (red, green, blue) | |
| def generate_baseball_card(smiles: str, img_size=(300, 300), font_path=None) -> Image.Image: | |
| mol = Chem.MolFromSmiles(smiles) | |
| if mol is None: | |
| raise ValueError("Invalid SMILES") | |
| # Get molecule-specific background color | |
| bg_color = generate_mol_background_color(mol) | |
| # Create transparent molecule image | |
| mol_img = Draw.MolToImage(mol, size=img_size, kekulize=True, bgColor=(255, 255, 255)) | |
| # Property summary | |
| props = { | |
| "Formula": rdMolDescriptors.CalcMolFormula(mol), | |
| "Mol Weight": f"{Descriptors.MolWt(mol):.2f}", | |
| "LogP": f"{Descriptors.MolLogP(mol):.2f}", | |
| "H-Donors": Descriptors.NumHDonors(mol), | |
| "H-Acceptors": Descriptors.NumHAcceptors(mol), | |
| "Rotatable Bonds": Descriptors.NumRotatableBonds(mol), | |
| "QED": round(QED.qed(mol), 2), | |
| } | |
| # Font | |
| font_size = 20 | |
| font = ImageFont.truetype(font_path, font_size) if font_path else ImageFont.load_default() | |
| # Layout dimensions | |
| padding = 20 | |
| box_margin = 6 | |
| stats_height = len(props) * (font_size + 6) + padding | |
| card_width = img_size[0] + 2 * padding | |
| card_height = img_size[1] + stats_height + 3 * padding | |
| # Create background card | |
| card = Image.new("RGB", (card_width, card_height), color=bg_color) | |
| draw = ImageDraw.Draw(card) | |
| # --- Draw Molecule Image Box --- | |
| img_x = padding | |
| img_y = padding | |
| mol_box_coords = [ | |
| img_x - box_margin, | |
| img_y - box_margin, | |
| img_x + img_size[0] + box_margin - 1, | |
| img_y + img_size[1] + box_margin - 1 | |
| ] | |
| draw.rectangle(mol_box_coords, fill="white", outline="black", width=2) | |
| card.paste(mol_img, (img_x, img_y)) | |
| # line_width = draw.textlength(smiles, font=font) | |
| # draw.text(((card_width - line_width) // 2, img_y), smiles, fill="black", font=font) | |
| # --- Draw Stats Box --- | |
| stats_y = img_y + img_size[1] + padding | |
| stats_box_coords = [ | |
| padding - box_margin, | |
| stats_y - box_margin, | |
| card_width - padding + box_margin - 1, | |
| stats_y + stats_height - padding + box_margin - 1 | |
| ] | |
| draw.rectangle(stats_box_coords, fill="white", outline="black", width=2) | |
| # Write centered stats | |
| text_y = stats_y + 6 | |
| for key, val in props.items(): | |
| line = f"{key}: {val}" | |
| line_width = draw.textlength(line, font=font) | |
| draw.text(((card_width - line_width) // 2, text_y), line, fill="black", font=font) | |
| text_y += font_size + 6 | |
| return card | |
| def generate_default_card(img_size=(300, 300)): | |
| font_size = 20 | |
| padding = 20 | |
| n_stats = 7 | |
| stats_height = n_stats * (font_size + 6) + padding | |
| card_width = img_size[0] + 2 * padding | |
| card_height = img_size[1] + stats_height + 3 * padding | |
| card = Image.new("RGB", (card_width, card_height), color='gray') | |
| return card | |
| # Load model config and checkpoint | |
| tokenizer = PreTrainedTokenizerFast(tokenizer_file='tokenizer_4096.json') | |
| checkpoint = torch.load('ckpt.pt', map_location=torch.device('cpu')) | |
| qed_bins = np.load('20_bins.npy') # maps [0, 1] -> [1, 20] by pdf and idx 0 is reserved for null embedding | |
| args = ModelArgs(**checkpoint['model_args']) | |
| model = ConditionalTransformer(args) | |
| model.load_state_dict(checkpoint['model']) | |
| model.eval() | |
| # O=C(C)Oc1ccccc1C(=O)O | |
| def generate_continuations(text): | |
| idx = [4096, 3] + tokenizer.encode(text) | |
| idx = torch.from_numpy(np.asarray(idx)).long().unsqueeze(0).expand(4, -1) | |
| out = model.generate(idx=idx, max_new_tokens=64, temperature=1.0).cpu().detach().numpy() | |
| # only the first generation in eac batch index is a continuation | |
| out = [tokenizer.decode(o) for o in out] | |
| out = [''.join(o.split(' ')) for o in out] | |
| out = [o.split('[SEP]')[1] for o in out] | |
| outs = [] | |
| for o in out: | |
| if verify_smiles(o): | |
| outs.append(generate_baseball_card(o)) | |
| if len(outs) < 3: | |
| outs.extend([generate_default_card() for _ in range(3 - len(outs))]) | |
| return outs[:3] | |
| def generate_similar(text): | |
| idx = [4096, 3] + tokenizer.encode(text) | |
| idx = torch.from_numpy(np.asarray(idx)).long().unsqueeze(0).expand(4, -1) | |
| out = model.generate(idx=idx, max_new_tokens=64, temperature=1.0).cpu().detach().numpy() | |
| out = decode(out, is_cond=False) | |
| outs = [] | |
| for o in out: | |
| if verify_smiles(o): | |
| outs.append(generate_baseball_card(o)) | |
| if len(outs) < 3: | |
| outs.extend([generate_default_card() for _ in range(3 - len(outs))]) | |
| return outs[:3] | |
| cond_map = { | |
| 'Low': 0, | |
| 'Medium': 0.5, | |
| 'High': 1 | |
| } | |
| def generate_cfg(qed): | |
| if qed == 'None' or qed is None: | |
| sep = torch.from_numpy(np.asarray([4096, 3])).long().unsqueeze(0).expand(2, -1) | |
| out = model.generate(idx=sep, max_new_tokens=64, temperature=1.0).cpu().detach().numpy() | |
| out = decode(out, is_cond=False) | |
| elif qed in cond_map: | |
| sep = torch.from_numpy(np.asarray([4096 + np.digitize(cond_map[qed], qed_bins), 3])).long().unsqueeze(0).expand(4, -1) | |
| out = conditional_generate(model, sep, max_new_tokens=64, cfg_scale=3).cpu().detach().numpy() | |
| out = decode(out[:, 2:], is_cond=True) | |
| else: | |
| raise NotImplementedError(f'QED = {qed} not implemented in generate') | |
| outs = [] | |
| for o in out: | |
| if verify_smiles(o): | |
| outs.append(generate_baseball_card(o)) | |
| if len(outs) < 3: | |
| outs.extend([generate_default_card() for _ in range(3 - len(outs))]) | |
| return outs[:3] | |
| def generate(qed, text, text_type): | |
| # figure out how to generate | |
| if qed == 'None' or qed is None: | |
| if text != '': | |
| if text_type == 'Continuation': | |
| return generate_continuations(text) | |
| elif text_type == 'Similar': | |
| return generate_similar(text) | |
| elif text_type == 'None' or text_type is None: | |
| sep = torch.from_numpy(np.asarray([4096, 3])).long().unsqueeze(0).expand(2, -1) | |
| out = model.generate(idx=sep, max_new_tokens=64, temperature=1.0).cpu().detach().numpy() | |
| out = decode(out, is_cond=False) | |
| else: | |
| raise NotImplementedError(f'Text Type {text_type} is not implemented in generate') | |
| else: | |
| sep = torch.from_numpy(np.asarray([4096, 3])).long().unsqueeze(0).expand(2, -1) | |
| out = model.generate(idx=sep, max_new_tokens=64, temperature=1.0).cpu().detach().numpy() | |
| out = decode(out, is_cond=False) | |
| elif qed in cond_map: | |
| sep = torch.from_numpy(np.asarray([4096 + np.digitize(cond_map[qed], qed_bins), 3])).long().unsqueeze(0).expand(4, -1) | |
| out = conditional_generate(model, sep, max_new_tokens=64, cfg_scale=3).cpu().detach().numpy() | |
| out = decode(out[:, 2:], is_cond=True) | |
| else: | |
| raise NotImplementedError(f'QED = {qed} not implemented in generate') | |
| outs = [] | |
| for o in out: | |
| if verify_smiles(o): | |
| outs.append(generate_baseball_card(o)) | |
| if len(outs) < 3: | |
| outs.extend([generate_default_card() for _ in range(3 - len(outs))]) | |
| return outs[:3] | |
| with gr.Blocks() as demo: | |
| gr.Markdown("# A Lightweight, Conditional, and Steerable Molecule Generator") | |
| with gr.Row(): | |
| btn = gr.Button("Generate Molecules") | |
| with gr.Row(): | |
| img1 = gr.Image() | |
| img2 = gr.Image() | |
| img3 = gr.Image() | |
| with gr.Accordion("Set optional conditions for molecular generation. For now, only one conditional option is allowed at a time!", open=True): | |
| qed = gr.Radio(['None', 'Low', 'Medium', 'High'], value='None', label="Quantitative Estimate of Druglikeness (QED)", show_label=True) | |
| text = gr.Textbox('O=C(C)Oc1ccccc1C', label='Type SMILES to condition generations on', show_label=True) | |
| text_button = gr.Radio(['None', 'Continuation', 'Similar'], value='None', label='How should the above SMILES be used?', show_label=True) | |
| btn.click( | |
| fn=generate, | |
| inputs=[qed, text, text_button], | |
| outputs=[img1, img2, img3] | |
| ) | |
| demo.launch() | |