File size: 6,054 Bytes
25f9bfc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 | print("Setup underway...")
# All necessary imports
from celery import Celery
from utils import wait_for_result
# Initialize Celery
celery_app = Celery()
print("Broker:", celery_app.conf.broker_url)
print("Backend:", celery_app.conf.result_backend)
print("")
print("============================")
print("1a. Product prediction - batch")
# Set up a list of reactants to make predictions
reactants_list = ["CCI.O=Cc1ccc([N+](=O)[O-])c(O)c1", "CCOc1cc(O)c(C=O)cc1OCC.OCCCBr", "C=CCc1cc(OCc2ccccc2)ccc1O.CCBr"]
# Setup task kwargs
kwargs = {
"topn": 1, # Number of results per reactant
"num_beams": 3, # Number of beams used for prediction. Must be >= topn
"device": None, # Device used for predicting, either "cuda" or "cpu", None defaults to cuda if available
"ckpt_forward": "Pistachio2025Q2-Forward", # Default forward model
"vocab": "Pistachio2025Q2", # Vocab for default forward model
# "ckpt_forward_path": "models/forward/Pistachio2025Q2-Forward.ckpt", # Can be used instead of ckpt_forward
# "vocab_path": "vocab/Pistachio2025Q2.txt", # Can be used instead of vocab
}
# Send the product_prediction task with the reaction list and kwargs
task = celery_app.send_task(
"tasks.product_prediction",
[reactants_list],
kwargs=kwargs,
queue="product_prediction",
)
print("Task sent. Assigned task_id: {}".format(task.id))
# Use the task id to get the result. Increase timeout if needed.
wait_for_result(celery_app, task.id, timeout=180)
print("")
print("============================")
print("1b. Product prediction - top 3")
# Set up a list of reactants to make predictions
reactants_list = ["CCI.O=Cc1ccc([N+](=O)[O-])c(O)c1"]
# Setup task kwargs
kwargs = {
"topn": 3, # Number of results per reactant
"num_beams": 5, # Number of beams used for prediction. Must be >= topn
"device": None, # Device used for predicting, either "cuda" or "cpu", None defaults to cuda if available
"ckpt_forward": "Pistachio2025Q2-Forward", # Default forward model
"vocab": "Pistachio2025Q2", # Vocab for default forward model
# "ckpt_forward_path": "models/forward/Pistachio2025Q2-Forward.ckpt", # Can be used instead of ckpt_forward
# "vocab_path": "vocab/Pistachio2025Q2.txt", # Can be used instead of vocab
}
# Send the product_prediction task with the reaction list and kwargs
task = celery_app.send_task(
"tasks.product_prediction",
[reactants_list],
kwargs=kwargs,
queue="product_prediction",
)
print("Task sent. Assigned task_id: {}".format(task.id))
# Use the task id to get the result. Increase timeout if needed.
wait_for_result(celery_app, task.id, timeout=180)
print("")
print("============================")
print("2. Retrosynthesis prediction")
# Choose product for retrosynthesis prediction
product = "C=CC(=C)C[Si](C)(C)C"
# Setup task kwargs
kwargs = {
"topn": 15, # Number of results per reactant
"num_beams": 15, # Number of beams used for prediction. Must be >= topn
"fap": 0.6, # Forward likelihood acceptance probability (not length averaged)
"fld": 0.2, # Forward likelihood delta required between the top2 forward prediction results
"device": None, # Device used for predicting, either "cuda" or "cpu", None defaults to cuda if available
"ckpt_forward": "Pistachio2025Q2-Forward", # Default forward model
"ckpt_retro": "Pistachio2025Q2-Retro", # Default retrosynthesis model
"vocab": "Pistachio2025Q2", # Vocab for default forward and retrosynthesis models
# "ckpt_forward_path": "models/forward/Pistachio2025Q2-Forward.ckpt", # Can be used instead of ckpt_forward
# "ckpt_retro_path": "models/retrosynthesis/Pistachio2025Q2-Retro.ckpt", # Can be used instead of ckpt_retro
# "vocab_path": "vocab/Pistachio2025Q2.txt", # Can be used instead of vocab
}
# Send the retro_prediction task with the product and kwargs
task = celery_app.send_task(
"tasks.retro_prediction",
[product],
kwargs=kwargs,
queue="retro_prediction",
)
print("Task sent. Assigned task_id: {}".format(task.id))
# Use the task id to get the result. Increase timeout if needed.
wait_for_result(celery_app, task.id, timeout=300)
print("")
print("============================")
print("3. Retro tree prediction")
# Choose product for retrosynthesis tree prediction
product = "C1C(C[Si](C)(C)C)=CCC2C(=O)OC(=O)C12"
#product = "Cc1cc2scnc2cc1N"
#product = "CC(C)(C(=O)O)C1C=CC=C(C2CC2)C1=O"
#product = "Nc1ccc2scnc2c1Br"
# Setup task kwargs
kwargs = {
"topn": 15, # Number of results per reactant
"num_beams": 15, # Number of beams used for prediction. Must be >= topn
"fap": 0.6, # Forward likelihood acceptance probability (not length averaged)
"fld": 0.2, # Forward likelihood delta required between the top2 forward prediction results
"max_depth": 4, # Max depth of the retrosynthesis tree
"beam_width": 6, # Max amount of nodes being expanded in each step
"device": None, # Device used for predicting, either "cuda" or "cpu", None defaults to cuda if available
"ckpt_forward": "Pistachio2025Q2-Forward", # Default forward model
"ckpt_retro": "Pistachio2025Q2-Retro", # Default retrosynthesis model
"vocab": "Pistachio2025Q2", # Vocab for default forward and retrosynthesis models
# "ckpt_forward_path": "models/forward/Pistachio2025Q2-Forward.ckpt", # Can be used instead of ckpt_forward
# "ckpt_retro_path": "models/retrosynthesis/Pistachio2025Q2-Retro.ckpt", # Can be used instead of ckpt_retro
# "vocab_path": "vocab/Pistachio2025Q2.txt", # Can be used instead of vocab
}
# Send the retro_prediction_tree task with the product and kwargs
task = celery_app.send_task(
"tasks.retro_prediction_tree",
[product],
kwargs=kwargs,
queue="retro_prediction",
)
print("Task sent. Assigned task_id: {}".format(task.id))
# Use the task id to get the result. Increase timeout if needed.
wait_for_result(celery_app, task.id, timeout=3000)
|