File size: 1,847 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 | """
Update this script and run it inside the worker container to run your retrosynthesis predictions via script
"""
print("Setting up retrosynthesis prediction...")
# 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)
# 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)
|