<jupyter_start><jupyter_text>Example Workflows - Model Optimization>Basic model optimization workflow  Model Optimization WorkflowsThis notebook shows a basic workflow for optimizing a generative model. The focus here is on showing how to set up the code, rather than maximizing performance. For this reason, we will use a simple template and a simple reward function.  Performance NotesThe workflow in this notebook is more CPU-constrained than GPU-constrained due to the need to evaluate samples on CPU. If you have a multi-core machine, it is recommended that you uncomment and run the `set_global_pool` cells in the notebook. This will trigger the use of multiprocessing, which will result in 2-4x speedups.This notebook may run slow on Collab due to CPU limitations.If running on Collab, remember to change the runtime to GPU<jupyter_code>import sys
sys.path.append('..')

from mrl.imports import *
from mrl.core import *
from mrl.chem import *
from mrl.templates.all import *

from mrl.torch_imports import *
from mrl.torch_core import *
from mrl.layers import *
from mrl.dataloaders import *
from mrl.g_models.all import *
from mrl.vocab import *
from mrl.policy_gradient import *
from mrl.train.all import *
from mrl.model_zoo import *<jupyter_output>/home/dmai/miniconda3/envs/mrl/lib/python3.7/importlib/_bootstrap.py:219: RuntimeWarning: to-Python converter for boost::shared_ptr<RDKit::FilterCatalogEntry const> already registered; second conversion method ignored.
  return f(*args, **kwds)<jupyter_text>AgentHere we create the model we want to optimize. We will use the `LSTM_LM_Small_ZINC` - a LSTM-based language model pretrained on part of the ZINC database<jupyter_code>agent = LSTM_LM_Small_ZINC(drop_scale=0.5,opt_kwargs={'lr':5e-5})<jupyter_output><empty_output><jupyter_text>Optional: Fine TuningWe could optionally fine-tune the pretrained model on a dataset of interest. This is an example of how to add a dataset to a pretrained `Agent` and finetune<jupyter_code># if in the repo
df = pd.read_csv('../files/smiles.csv')

# if in Collab
# download_files()
# df = pd.read_csv('files/smiles.csv')
agent.update_dataset_from_inputs(df.smiles.values)
agent.train_supervised(64, 1, 1e-5)<jupyter_output><empty_output><jupyter_text>TemplateWe will set up a very basic template that will only check compounds for structural validity<jupyter_code>template = Template([ValidityFilter(), 
                     SingleCompoundFilter()],
                    [])

template_cb = TemplateCallback(template, prefilter=True)<jupyter_output><empty_output><jupyter_text>RewardFor the reward, we will load a scikit-learn linear regression model trained to predict affinity against erbB1 using molecular fingerprints.This score function is extremely simple and won't translate well to affinity. It is used as a lightweight example<jupyter_code>class FP_Regression_Score():
    def __init__(self, fname):
        self.model = torch.load(fname)
        self.fp_function = partial(failsafe_fp, fp_function=ECFP6)
        
    def __call__(self, samples):
        mols = to_mols(samples)
        fps = maybe_parallel(self.fp_function, mols)
        fps = [fp_to_array(i) for i in fps]
        x_vals = np.stack(fps)
        preds = self.model.predict(x_vals)
        return preds
    
    
# if in the repo
reward_function = FP_Regression_Score('../files/erbB1_regression.sklearn')

# if in Collab
# download_files()
# reward_function = FP_Regression_Score('files/erbB1_regression.sklearn')

reward = Reward(reward_function, weight=1.)

aff_reward = RewardCallback(reward, 'aff')<jupyter_output><empty_output><jupyter_text>Loss FunctionWe will use the `PolicyGradient` class, the simplest policy gradient algorithm<jupyter_code>pg = PolicyGradient(discount=True, gamma=0.97)

loss = PolicyLoss(pg, 'PG')<jupyter_output><empty_output><jupyter_text>Optional: PPOA more sophisticated policy gradient algorithm<jupyter_code># pg = PPO(0.99,
#         0.5,
#         lam=0.95,
#         v_coef=0.5,
#         cliprange=0.3,
#         v_cliprange=0.3,
#         ent_coef=0.01,
#         kl_target=0.03,
#         kl_horizon=3000,
#         scale_rewards=True)

# loss = PolicyLoss(pg, 'PPO', 
#                    value_head=ValueHead(256), 
#                    v_update_iter=2, 
#                    vopt_kwargs={'lr':1e-3})<jupyter_output><empty_output><jupyter_text>SamplersWe will use a basic sampler setup which will build the buffer from the main model and train only on samples from the buffer<jupyter_code>gen_bs = 1500
sampler1 = ModelSampler(agent.vocab, agent.model, 'live', 400, 0., gen_bs)
samplers = [sampler1]<jupyter_output><empty_output><jupyter_text>Optional: More SamplersSamplers that draw from the base model and the log<jupyter_code># sampler2 = ModelSampler(agent.vocab, agent.base_model, 'base', 400, 0., gen_bs)
# sampler3 = LogSampler('samples', 'rewards', 10, 97, 400)

# samplers += [sampler2, sampler3]<jupyter_output><empty_output><jupyter_text>Other CallbacksWe will add two callbacks for printing metrics<jupyter_code>live_max = MaxCallback('rewards', 'live')
live_p90 = PercentileCallback('rewards', 'live', 90)

cbs = [live_p90, live_max]<jupyter_output><empty_output><jupyter_text>EnvironmentWe create our environment with the objects assembled so far<jupyter_code>env = Environment(agent, template_cb, samplers=samplers, rewards=[aff_reward], losses=[loss],
                 cbs=cbs)<jupyter_output><empty_output><jupyter_text>Train<jupyter_code># set_global_pool(cpus=min(10, os.cpu_count()))
env.fit(128, 90, 500, 25)
env.log.plot_metrics()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.09446428571428571, 'fraction_numerical': 0.0175, 'mean_word_length': 4.834375, 'pattern_counts': {'":': 0, '<': 49, '<?xml version=': 0, '>': 50, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '12', 'dir': 'jupyter-structured-clean-dedup', 'id': '317458', 'license': 'MIT', 'n_tokens_mistral': 1902, 'n_tokens_neox': 1805, 'n_words': 492, 'path': 'nbs/tutorials/tutorials.workflows.model_optimization.ipynb', 'repo_name': 'DarkMatterAI/mrl', 'size': '5600'}