decimer 2.7.1__tar.gz → 2.8.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {decimer-2.7.1 → decimer-2.8.0}/DECIMER/Efficient_Net_encoder.py +0 -2
- decimer-2.8.0/DECIMER/__init__.py +31 -0
- decimer-2.8.0/DECIMER/config.py +148 -0
- {decimer-2.7.1 → decimer-2.8.0}/DECIMER/decimer.py +44 -4
- decimer-2.7.1/DECIMER/config.py → decimer-2.8.0/DECIMER/pre_process.py +8 -189
- decimer-2.8.0/DECIMER/utils.py +143 -0
- {decimer-2.7.1 → decimer-2.8.0}/PKG-INFO +85 -59
- {decimer-2.7.1 → decimer-2.8.0}/README.md +59 -50
- {decimer-2.7.1 → decimer-2.8.0}/decimer.egg-info/PKG-INFO +85 -59
- {decimer-2.7.1 → decimer-2.8.0}/decimer.egg-info/SOURCES.txt +4 -1
- {decimer-2.7.1 → decimer-2.8.0}/decimer.egg-info/requires.txt +1 -1
- decimer-2.8.0/pyproject.toml +3 -0
- {decimer-2.7.1 → decimer-2.8.0}/setup.py +7 -10
- decimer-2.8.0/tests/test_functions.py +30 -0
- decimer-2.7.1/DECIMER/__init__.py +0 -29
- decimer-2.7.1/DECIMER/utils.py +0 -97
- {decimer-2.7.1 → decimer-2.8.0}/DECIMER/DECIMER_EfficinetNetV2_Transfomer_Trainer.py +0 -0
- {decimer-2.7.1 → decimer-2.8.0}/DECIMER/Predictor_usingCheckpoints.py +0 -0
- {decimer-2.7.1 → decimer-2.8.0}/DECIMER/Repack_model.py +0 -0
- {decimer-2.7.1 → decimer-2.8.0}/DECIMER/Transformer_decoder.py +0 -0
- {decimer-2.7.1 → decimer-2.8.0}/DECIMER/Utils/Create_TFrecord_From_pickles.py +0 -0
- {decimer-2.7.1 → decimer-2.8.0}/DECIMER/Utils/Create_image_tokenizer.py +0 -0
- {decimer-2.7.1 → decimer-2.8.0}/DECIMER/efficientnetv2/effnetv2_configs.py +0 -0
- {decimer-2.7.1 → decimer-2.8.0}/DECIMER/efficientnetv2/effnetv2_model.py +0 -0
- {decimer-2.7.1 → decimer-2.8.0}/DECIMER/efficientnetv2/hparams.py +0 -0
- {decimer-2.7.1 → decimer-2.8.0}/DECIMER/efficientnetv2/utils.py +0 -0
- {decimer-2.7.1 → decimer-2.8.0}/LICENSE +0 -0
- {decimer-2.7.1 → decimer-2.8.0}/decimer.egg-info/dependency_links.txt +0 -0
- {decimer-2.7.1 → decimer-2.8.0}/decimer.egg-info/entry_points.txt +0 -0
- {decimer-2.7.1 → decimer-2.8.0}/decimer.egg-info/top_level.txt +0 -0
- {decimer-2.7.1 → decimer-2.8.0}/setup.cfg +0 -0
|
@@ -0,0 +1,31 @@
|
|
|
1
|
+
# -*- coding: utf-8 -*-
|
|
2
|
+
"""
|
|
3
|
+
DECIMER V2.7.2 Python Package
|
|
4
|
+
=============================
|
|
5
|
+
|
|
6
|
+
Deep lEarning for Chemical ImagE Recognition (DECIMER) project
|
|
7
|
+
was launched to address the Optical Chemical Structure Recognition (OCSR)
|
|
8
|
+
problem using deep learning based methods,
|
|
9
|
+
providing an automated open-source software solution.
|
|
10
|
+
|
|
11
|
+
Typical usage example::
|
|
12
|
+
|
|
13
|
+
from decimer import predict_SMILES
|
|
14
|
+
|
|
15
|
+
# Chemical depiction to SMILES translation
|
|
16
|
+
image_path = "path/to/imagefile"
|
|
17
|
+
SMILES = predict_SMILES(image_path)
|
|
18
|
+
print(SMILES)
|
|
19
|
+
|
|
20
|
+
For comments, bug reports, or feature ideas,
|
|
21
|
+
please raise an issue on the GitHub repository.
|
|
22
|
+
"""
|
|
23
|
+
|
|
24
|
+
__version__ = "2.7.2"
|
|
25
|
+
|
|
26
|
+
__all__ = [
|
|
27
|
+
"DECIMER",
|
|
28
|
+
]
|
|
29
|
+
|
|
30
|
+
|
|
31
|
+
from .decimer import predict_SMILES
|
|
@@ -0,0 +1,148 @@
|
|
|
1
|
+
import tensorflow as tf
|
|
2
|
+
|
|
3
|
+
import DECIMER.Efficient_Net_encoder as Efficient_Net_encoder
|
|
4
|
+
import DECIMER.Transformer_decoder as Transformer_decoder
|
|
5
|
+
|
|
6
|
+
TARGET_DTYPE = tf.float32
|
|
7
|
+
|
|
8
|
+
|
|
9
|
+
class Config:
|
|
10
|
+
"""Configuration class."""
|
|
11
|
+
|
|
12
|
+
def __init__(
|
|
13
|
+
self,
|
|
14
|
+
):
|
|
15
|
+
self.encoder_config = {}
|
|
16
|
+
self.transformer_config = {}
|
|
17
|
+
self.lr_config = {}
|
|
18
|
+
|
|
19
|
+
def initialize_encoder_config(
|
|
20
|
+
self,
|
|
21
|
+
image_embedding_dim,
|
|
22
|
+
preprocessing_fn,
|
|
23
|
+
backbone_fn,
|
|
24
|
+
image_shape,
|
|
25
|
+
do_permute=False,
|
|
26
|
+
pretrained_weights=None,
|
|
27
|
+
):
|
|
28
|
+
"""This functions initializes the Efficient-Net V2 encoder with user
|
|
29
|
+
defined configurations.
|
|
30
|
+
|
|
31
|
+
Args:
|
|
32
|
+
image_embedding_dim (int): Embedding dimension of the input image
|
|
33
|
+
preprocessing_fn (method): Efficient Net preprocessing function for input image
|
|
34
|
+
backbone_fn (method): Calls Efficient-Net V2 as backbone for encoder
|
|
35
|
+
image_shape (int): Shape of the input image
|
|
36
|
+
do_permute (bool, optional): . Defaults to False.
|
|
37
|
+
pretrained_weights (keras weights, optional): Use pretrained efficient net weights or not. Defaults to None.
|
|
38
|
+
"""
|
|
39
|
+
self.encoder_config = dict(
|
|
40
|
+
image_embedding_dim=image_embedding_dim,
|
|
41
|
+
preprocessing_fn=preprocessing_fn,
|
|
42
|
+
backbone_fn=backbone_fn,
|
|
43
|
+
image_shape=image_shape,
|
|
44
|
+
do_permute=do_permute,
|
|
45
|
+
pretrained_weights=pretrained_weights,
|
|
46
|
+
)
|
|
47
|
+
|
|
48
|
+
def initialize_transformer_config(
|
|
49
|
+
self,
|
|
50
|
+
vocab_len,
|
|
51
|
+
max_len,
|
|
52
|
+
n_transformer_layers,
|
|
53
|
+
transformer_d_dff,
|
|
54
|
+
transformer_n_heads,
|
|
55
|
+
image_embedding_dim,
|
|
56
|
+
rate=0.1,
|
|
57
|
+
):
|
|
58
|
+
"""This functions initializes the Transformer model as decoder with
|
|
59
|
+
user defined configurations.
|
|
60
|
+
|
|
61
|
+
Args:
|
|
62
|
+
vocab_len (int): Total number of words in the input vocabulary
|
|
63
|
+
max_len (int): Maximum length of the string found on the training dataset
|
|
64
|
+
n_transformer_layers (int): Number of layers present in the transformer model
|
|
65
|
+
transformer_d_dff (int): Transformer feed forward upwards projection size
|
|
66
|
+
transformer_n_heads (int): Number of heads present in the transformer model
|
|
67
|
+
image_embedding_dim (int): Total number of dimension the image gets embedded
|
|
68
|
+
dropout_rate (float, optional): Fraction of the input units to drop. Defaults to 0.1.
|
|
69
|
+
"""
|
|
70
|
+
self.transformer_config = dict(
|
|
71
|
+
num_layers=n_transformer_layers,
|
|
72
|
+
d_model=image_embedding_dim,
|
|
73
|
+
num_heads=transformer_n_heads,
|
|
74
|
+
dff=transformer_d_dff,
|
|
75
|
+
target_vocab_size=vocab_len,
|
|
76
|
+
max_len=max_len,
|
|
77
|
+
rate=0.1,
|
|
78
|
+
)
|
|
79
|
+
|
|
80
|
+
def initialize_lr_config(self, warm_steps, n_epochs):
|
|
81
|
+
"""This function sets the configuration to initialize learning rate.
|
|
82
|
+
|
|
83
|
+
Args:
|
|
84
|
+
warm_steps (int): Number of steps The learning rate is increased
|
|
85
|
+
n_epochs (int): Number of epochs
|
|
86
|
+
"""
|
|
87
|
+
self.lr_config = dict(
|
|
88
|
+
warm_steps=warm_steps,
|
|
89
|
+
n_epochs=n_epochs,
|
|
90
|
+
)
|
|
91
|
+
|
|
92
|
+
|
|
93
|
+
class CustomSchedule(tf.keras.optimizers.schedules.LearningRateSchedule):
|
|
94
|
+
"""Custom schedule for learning rate used during training.
|
|
95
|
+
|
|
96
|
+
Args:
|
|
97
|
+
tf (_type_): keras learning rate schedule
|
|
98
|
+
"""
|
|
99
|
+
|
|
100
|
+
def __init__(self, d_model, warmup_steps=4000):
|
|
101
|
+
super(CustomSchedule, self).__init__()
|
|
102
|
+
|
|
103
|
+
self.d_model = d_model
|
|
104
|
+
self.d_model = tf.cast(self.d_model, tf.float32)
|
|
105
|
+
|
|
106
|
+
self.warmup_steps = warmup_steps
|
|
107
|
+
|
|
108
|
+
def __call__(self, step):
|
|
109
|
+
arg1 = tf.math.rsqrt(step)
|
|
110
|
+
arg2 = step * (self.warmup_steps**-1.5)
|
|
111
|
+
|
|
112
|
+
return tf.math.rsqrt(self.d_model) * tf.math.minimum(arg1, arg2)
|
|
113
|
+
|
|
114
|
+
|
|
115
|
+
def prepare_models(encoder_config, transformer_config, replica_batch_size, verbose=0):
|
|
116
|
+
"""This function is used to initiate the Encoder and the Transformer with
|
|
117
|
+
appropriate configs set by the user. After initiating the models this
|
|
118
|
+
function returns the Encoder,Transformer and the optimizer.
|
|
119
|
+
|
|
120
|
+
Args:
|
|
121
|
+
encoder_config ([type]): Encoder configuration set by user in the config class.
|
|
122
|
+
transformer_config ([type]): Transformer configuration set by user in the config class.
|
|
123
|
+
replica_batch_size ([type]): Per replica batch size set by user(during distributed training).
|
|
124
|
+
verbose (int, optional): Defaults to 0.
|
|
125
|
+
|
|
126
|
+
Returns:
|
|
127
|
+
[type]: Optimizer, Encoder model and the Transformer
|
|
128
|
+
"""
|
|
129
|
+
|
|
130
|
+
# Instantiate an optimizer
|
|
131
|
+
optimizer = tf.keras.optimizers.Adam(learning_rate=0.00051)
|
|
132
|
+
|
|
133
|
+
# Instantiate the encoder model
|
|
134
|
+
|
|
135
|
+
encoder = Efficient_Net_encoder.Encoder(**encoder_config)
|
|
136
|
+
|
|
137
|
+
# Instantiate the decoder model
|
|
138
|
+
transformer = Transformer_decoder.Decoder(**transformer_config)
|
|
139
|
+
|
|
140
|
+
# Show the model architectures and plot the learning rate
|
|
141
|
+
if verbose:
|
|
142
|
+
print("\nEncoder model\n")
|
|
143
|
+
print(encoder.summary())
|
|
144
|
+
|
|
145
|
+
print("\nTransformer model\n")
|
|
146
|
+
print(transformer.summary())
|
|
147
|
+
|
|
148
|
+
return optimizer, encoder, transformer
|
|
@@ -9,7 +9,7 @@ import numpy as np
|
|
|
9
9
|
import pystow
|
|
10
10
|
import tensorflow as tf
|
|
11
11
|
|
|
12
|
-
import DECIMER.
|
|
12
|
+
import DECIMER.pre_process as pre_process
|
|
13
13
|
import DECIMER.utils as utils
|
|
14
14
|
|
|
15
15
|
# Silence tensorflow model loading warnings.
|
|
@@ -18,6 +18,46 @@ logging.getLogger("absl").setLevel("ERROR")
|
|
|
18
18
|
# Silence tensorflow errors - not recommended if your model is not working properly.
|
|
19
19
|
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
|
|
20
20
|
|
|
21
|
+
|
|
22
|
+
# ============================================================================
|
|
23
|
+
# Keras 2 to Keras 3 Compatibility Layer
|
|
24
|
+
# ============================================================================
|
|
25
|
+
class KerasCompatibilityUnpickler(pickle.Unpickler):
|
|
26
|
+
"""Custom unpickler to handle Keras 2.x tokenizers in Keras 3 environment.
|
|
27
|
+
|
|
28
|
+
This handles the module path changes between Keras 2 and Keras 3:
|
|
29
|
+
- keras.preprocessing.text -> tensorflow.keras.preprocessing.text
|
|
30
|
+
"""
|
|
31
|
+
|
|
32
|
+
def find_class(self, module, name):
|
|
33
|
+
# Redirect old Keras 2.x imports to TensorFlow Keras
|
|
34
|
+
if module.startswith("keras."):
|
|
35
|
+
module = module.replace("keras.", "tensorflow.keras.")
|
|
36
|
+
return super().find_class(module, name)
|
|
37
|
+
|
|
38
|
+
|
|
39
|
+
def load_tokenizer(tokenizer_path: str) -> object:
|
|
40
|
+
"""Load tokenizer with Keras 2/3 compatibility.
|
|
41
|
+
|
|
42
|
+
Args:
|
|
43
|
+
tokenizer_path (str): Path to the pickled tokenizer file
|
|
44
|
+
|
|
45
|
+
Returns:
|
|
46
|
+
tokenizer: Loaded tokenizer object
|
|
47
|
+
"""
|
|
48
|
+
try:
|
|
49
|
+
# Try standard pickle load first (for Keras 3 tokenizers)
|
|
50
|
+
with open(tokenizer_path, "rb") as f:
|
|
51
|
+
return pickle.load(f)
|
|
52
|
+
except ModuleNotFoundError as e:
|
|
53
|
+
if "keras.preprocessing" in str(e):
|
|
54
|
+
# Use compatibility unpickler for Keras 2.x tokenizers
|
|
55
|
+
with open(tokenizer_path, "rb") as f:
|
|
56
|
+
return KerasCompatibilityUnpickler(f).load()
|
|
57
|
+
else:
|
|
58
|
+
raise
|
|
59
|
+
|
|
60
|
+
|
|
21
61
|
# Set the absolute path
|
|
22
62
|
HERE = os.path.dirname(os.path.abspath(__file__))
|
|
23
63
|
|
|
@@ -59,7 +99,7 @@ def get_models(model_urls: dict):
|
|
|
59
99
|
tokenizer_path = os.path.join(
|
|
60
100
|
model_paths["DECIMER"], "assets", "tokenizer_SMILES.pkl"
|
|
61
101
|
)
|
|
62
|
-
tokenizer =
|
|
102
|
+
tokenizer = load_tokenizer(tokenizer_path)
|
|
63
103
|
|
|
64
104
|
# Load DECIMER models
|
|
65
105
|
DECIMER_V2 = tf.saved_model.load(model_paths["DECIMER"])
|
|
@@ -123,7 +163,7 @@ def detokenize_output_add_confidence(
|
|
|
123
163
|
|
|
124
164
|
|
|
125
165
|
def predict_SMILES(
|
|
126
|
-
|
|
166
|
+
image_input: [str, np.ndarray], confidence: bool = False, hand_drawn: bool = False
|
|
127
167
|
) -> str:
|
|
128
168
|
"""Predicts SMILES representation of a molecule depicted in the given image.
|
|
129
169
|
|
|
@@ -135,7 +175,7 @@ def predict_SMILES(
|
|
|
135
175
|
Returns:
|
|
136
176
|
str: SMILES representation of the molecule in the input image, optionally with confidence values.
|
|
137
177
|
"""
|
|
138
|
-
chemical_structure =
|
|
178
|
+
chemical_structure = pre_process.decode_image(image_input)
|
|
139
179
|
|
|
140
180
|
model = DECIMER_Hand_drawn if hand_drawn else DECIMER_V2
|
|
141
181
|
predicted_tokens, confidence_values = model(tf.constant(chemical_structure))
|
|
@@ -1,24 +1,14 @@
|
|
|
1
|
-
# Network configuration file
|
|
2
1
|
import io
|
|
3
|
-
import zipfile
|
|
4
|
-
from pathlib import Path
|
|
5
|
-
|
|
6
2
|
import cv2
|
|
7
3
|
import efficientnet.tfkeras as efn
|
|
8
|
-
import numpy as np
|
|
9
|
-
import pystow
|
|
10
4
|
import tensorflow as tf
|
|
5
|
+
import numpy as np
|
|
6
|
+
|
|
11
7
|
from PIL import Image
|
|
12
8
|
from PIL import ImageEnhance
|
|
13
9
|
from pillow_heif import register_heif_opener
|
|
14
10
|
from typing import Union
|
|
15
11
|
|
|
16
|
-
import DECIMER.Efficient_Net_encoder as Efficient_Net_encoder
|
|
17
|
-
import DECIMER.Transformer_decoder as Transformer_decoder
|
|
18
|
-
|
|
19
|
-
|
|
20
|
-
TARGET_DTYPE = tf.float32
|
|
21
|
-
|
|
22
12
|
|
|
23
13
|
def resize_byratio(image):
|
|
24
14
|
"""This function takes a Pillow Image object and will resize the image by
|
|
@@ -106,6 +96,7 @@ def remove_transparent(image: Union[str, np.ndarray]) -> Image.Image:
|
|
|
106
96
|
Returns:
|
|
107
97
|
PIL.Image.Image: The image with transparency removed.
|
|
108
98
|
"""
|
|
99
|
+
|
|
109
100
|
def process_image(png: Image.Image) -> Image.Image:
|
|
110
101
|
"""
|
|
111
102
|
Helper function to remove transparency from a single image.
|
|
@@ -183,14 +174,14 @@ def increase_contrast(image):
|
|
|
183
174
|
Returns: PIL.Image
|
|
184
175
|
"""
|
|
185
176
|
|
|
186
|
-
# Get brightness range
|
|
187
|
-
|
|
188
|
-
|
|
177
|
+
# Get brightness range - cast to int to avoid uint8 overflow
|
|
178
|
+
min_val = int(np.min(image))
|
|
179
|
+
max_val = int(np.max(image))
|
|
189
180
|
|
|
190
181
|
# Use LUT to convert image values
|
|
191
182
|
LUT = np.zeros(256, dtype=np.uint8)
|
|
192
|
-
LUT[
|
|
193
|
-
start=0, stop=255, num=(
|
|
183
|
+
LUT[min_val : max_val + 1] = np.linspace(
|
|
184
|
+
start=0, stop=255, num=(max_val - min_val) + 1, endpoint=True, dtype=np.uint8
|
|
194
185
|
)
|
|
195
186
|
|
|
196
187
|
# Apply LUT and return image
|
|
@@ -250,175 +241,3 @@ def decode_image(image_path: Union[str, np.ndarray]):
|
|
|
250
241
|
img = tf.image.resize(img, (512, 512), method="gaussian", antialias=True)
|
|
251
242
|
img = efn.preprocess_input(img)
|
|
252
243
|
return img
|
|
253
|
-
|
|
254
|
-
|
|
255
|
-
class Config:
|
|
256
|
-
"""Configuration class."""
|
|
257
|
-
|
|
258
|
-
def __init__(
|
|
259
|
-
self,
|
|
260
|
-
):
|
|
261
|
-
self.encoder_config = {}
|
|
262
|
-
self.transformer_config = {}
|
|
263
|
-
self.lr_config = {}
|
|
264
|
-
|
|
265
|
-
def initialize_encoder_config(
|
|
266
|
-
self,
|
|
267
|
-
image_embedding_dim,
|
|
268
|
-
preprocessing_fn,
|
|
269
|
-
backbone_fn,
|
|
270
|
-
image_shape,
|
|
271
|
-
do_permute=False,
|
|
272
|
-
pretrained_weights=None,
|
|
273
|
-
):
|
|
274
|
-
"""This functions initializes the Efficient-Net V2 encoder with user
|
|
275
|
-
defined configurations.
|
|
276
|
-
|
|
277
|
-
Args:
|
|
278
|
-
image_embedding_dim (int): Embedding dimension of the input image
|
|
279
|
-
preprocessing_fn (method): Efficient Net preprocessing function for input image
|
|
280
|
-
backbone_fn (method): Calls Efficient-Net V2 as backbone for encoder
|
|
281
|
-
image_shape (int): Shape of the input image
|
|
282
|
-
do_permute (bool, optional): . Defaults to False.
|
|
283
|
-
pretrained_weights (keras weights, optional): Use pretrained efficient net weights or not. Defaults to None.
|
|
284
|
-
"""
|
|
285
|
-
self.encoder_config = dict(
|
|
286
|
-
image_embedding_dim=image_embedding_dim,
|
|
287
|
-
preprocessing_fn=preprocessing_fn,
|
|
288
|
-
backbone_fn=backbone_fn,
|
|
289
|
-
image_shape=image_shape,
|
|
290
|
-
do_permute=do_permute,
|
|
291
|
-
pretrained_weights=pretrained_weights,
|
|
292
|
-
)
|
|
293
|
-
|
|
294
|
-
def initialize_transformer_config(
|
|
295
|
-
self,
|
|
296
|
-
vocab_len,
|
|
297
|
-
max_len,
|
|
298
|
-
n_transformer_layers,
|
|
299
|
-
transformer_d_dff,
|
|
300
|
-
transformer_n_heads,
|
|
301
|
-
image_embedding_dim,
|
|
302
|
-
rate=0.1,
|
|
303
|
-
):
|
|
304
|
-
"""This functions initializes the Transformer model as decoder with
|
|
305
|
-
user defined configurations.
|
|
306
|
-
|
|
307
|
-
Args:
|
|
308
|
-
vocab_len (int): Total number of words in the input vocabulary
|
|
309
|
-
max_len (int): Maximum length of the string found on the training dataset
|
|
310
|
-
n_transformer_layers (int): Number of layers present in the transformer model
|
|
311
|
-
transformer_d_dff (int): Transformer feed forward upwards projection size
|
|
312
|
-
transformer_n_heads (int): Number of heads present in the transformer model
|
|
313
|
-
image_embedding_dim (int): Total number of dimension the image gets embedded
|
|
314
|
-
dropout_rate (float, optional): Fraction of the input units to drop. Defaults to 0.1.
|
|
315
|
-
"""
|
|
316
|
-
self.transformer_config = dict(
|
|
317
|
-
num_layers=n_transformer_layers,
|
|
318
|
-
d_model=image_embedding_dim,
|
|
319
|
-
num_heads=transformer_n_heads,
|
|
320
|
-
dff=transformer_d_dff,
|
|
321
|
-
target_vocab_size=vocab_len,
|
|
322
|
-
max_len=max_len,
|
|
323
|
-
rate=0.1,
|
|
324
|
-
)
|
|
325
|
-
|
|
326
|
-
def initialize_lr_config(self, warm_steps, n_epochs):
|
|
327
|
-
"""This function sets the configuration to initialize learning rate.
|
|
328
|
-
|
|
329
|
-
Args:
|
|
330
|
-
warm_steps (int): Number of steps The learning rate is increased
|
|
331
|
-
n_epochs (int): Number of epochs
|
|
332
|
-
"""
|
|
333
|
-
self.lr_config = dict(
|
|
334
|
-
warm_steps=warm_steps,
|
|
335
|
-
n_epochs=n_epochs,
|
|
336
|
-
)
|
|
337
|
-
|
|
338
|
-
|
|
339
|
-
class CustomSchedule(tf.keras.optimizers.schedules.LearningRateSchedule):
|
|
340
|
-
"""Custom schedule for learning rate used during training.
|
|
341
|
-
|
|
342
|
-
Args:
|
|
343
|
-
tf (_type_): keras learning rate schedule
|
|
344
|
-
"""
|
|
345
|
-
|
|
346
|
-
def __init__(self, d_model, warmup_steps=4000):
|
|
347
|
-
super(CustomSchedule, self).__init__()
|
|
348
|
-
|
|
349
|
-
self.d_model = d_model
|
|
350
|
-
self.d_model = tf.cast(self.d_model, tf.float32)
|
|
351
|
-
|
|
352
|
-
self.warmup_steps = warmup_steps
|
|
353
|
-
|
|
354
|
-
def __call__(self, step):
|
|
355
|
-
arg1 = tf.math.rsqrt(step)
|
|
356
|
-
arg2 = step * (self.warmup_steps**-1.5)
|
|
357
|
-
|
|
358
|
-
return tf.math.rsqrt(self.d_model) * tf.math.minimum(arg1, arg2)
|
|
359
|
-
|
|
360
|
-
|
|
361
|
-
def prepare_models(encoder_config, transformer_config, replica_batch_size, verbose=0):
|
|
362
|
-
"""This function is used to initiate the Encoder and the Transformer with
|
|
363
|
-
appropriate configs set by the user. After initiating the models this
|
|
364
|
-
function returns the Encoder,Transformer and the optimizer.
|
|
365
|
-
|
|
366
|
-
Args:
|
|
367
|
-
encoder_config ([type]): Encoder configuration set by user in the config class.
|
|
368
|
-
transformer_config ([type]): Transformer configuration set by user in the config class.
|
|
369
|
-
replica_batch_size ([type]): Per replica batch size set by user(during distributed training).
|
|
370
|
-
verbose (int, optional): Defaults to 0.
|
|
371
|
-
|
|
372
|
-
Returns:
|
|
373
|
-
[type]: Optimizer, Encoder model and the Transformer
|
|
374
|
-
"""
|
|
375
|
-
|
|
376
|
-
# Instantiate an optimizer
|
|
377
|
-
optimizer = tf.keras.optimizers.Adam(learning_rate=0.00051)
|
|
378
|
-
|
|
379
|
-
# Instantiate the encoder model
|
|
380
|
-
|
|
381
|
-
encoder = Efficient_Net_encoder.Encoder(**encoder_config)
|
|
382
|
-
|
|
383
|
-
# Instantiate the decoder model
|
|
384
|
-
transformer = Transformer_decoder.Decoder(**transformer_config)
|
|
385
|
-
|
|
386
|
-
# Show the model architectures and plot the learning rate
|
|
387
|
-
if verbose:
|
|
388
|
-
print("\nEncoder model\n")
|
|
389
|
-
print(encoder.summary())
|
|
390
|
-
|
|
391
|
-
print("\nTransformer model\n")
|
|
392
|
-
print(transformer.summary())
|
|
393
|
-
|
|
394
|
-
return optimizer, encoder, transformer
|
|
395
|
-
|
|
396
|
-
|
|
397
|
-
# Downloads the model and unzips the file downloaded, if the model is not present on the working directory.
|
|
398
|
-
def download_trained_weights(model_url: str, model_path: str, verbose=1):
|
|
399
|
-
"""This function downloads the trained models and tokenizers to a default
|
|
400
|
-
location. After downloading the zipped file the function unzips the file
|
|
401
|
-
automatically. If the model exists on the default location this function
|
|
402
|
-
will not work.
|
|
403
|
-
|
|
404
|
-
Args:
|
|
405
|
-
model_url (str): trained model url for downloading.
|
|
406
|
-
model_path (str): model default path to download.
|
|
407
|
-
|
|
408
|
-
Returns:
|
|
409
|
-
path (str): downloaded model.
|
|
410
|
-
"""
|
|
411
|
-
# Download trained models
|
|
412
|
-
if verbose > 0:
|
|
413
|
-
print("Downloading trained model to " + str(model_path))
|
|
414
|
-
model_path = pystow.ensure("DECIMER-V2", url=model_url)
|
|
415
|
-
if verbose > 0:
|
|
416
|
-
print(model_path)
|
|
417
|
-
print("... done downloading trained model!")
|
|
418
|
-
|
|
419
|
-
with zipfile.ZipFile(model_path.as_posix(), "r") as zip_ref:
|
|
420
|
-
zip_ref.extractall(model_path.parent.as_posix())
|
|
421
|
-
|
|
422
|
-
# Delete zipfile after downloading
|
|
423
|
-
if Path(model_path).exists():
|
|
424
|
-
Path(model_path).unlink()
|
|
@@ -0,0 +1,143 @@
|
|
|
1
|
+
import os
|
|
2
|
+
import re
|
|
3
|
+
import shutil
|
|
4
|
+
|
|
5
|
+
import pystow
|
|
6
|
+
import zipfile
|
|
7
|
+
from pathlib import Path
|
|
8
|
+
|
|
9
|
+
pattern = "R([0-9]*)|X([0-9]*)|Y([0-9]*)|Z([0-9]*)"
|
|
10
|
+
add_space_re = r"^(\W+)|(\W+)$"
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
def split_and_modify_atoms(SMILES):
|
|
14
|
+
splitted_SMILES = list(SMILES)
|
|
15
|
+
modified_SMILES = re.sub(r"\s+(?=[a-z])", "", " ".join(map(str, splitted_SMILES)))
|
|
16
|
+
return modified_SMILES
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
def replacer(match):
|
|
20
|
+
part = match.group(0)
|
|
21
|
+
text = (
|
|
22
|
+
part.replace("1", "!")
|
|
23
|
+
.replace("2", "$")
|
|
24
|
+
.replace("3", "^")
|
|
25
|
+
.replace("4", "<")
|
|
26
|
+
.replace("5", ">")
|
|
27
|
+
.replace("6", "?")
|
|
28
|
+
.replace("7", "£")
|
|
29
|
+
.replace("8", "¢")
|
|
30
|
+
.replace("9", "€")
|
|
31
|
+
.replace("0", "§")
|
|
32
|
+
)
|
|
33
|
+
return text
|
|
34
|
+
|
|
35
|
+
|
|
36
|
+
def add_space(match_):
|
|
37
|
+
if match_.group(1) is not None:
|
|
38
|
+
return "{} ".format(match_.group(1))
|
|
39
|
+
else:
|
|
40
|
+
return " {}".format(match_.group(2))
|
|
41
|
+
|
|
42
|
+
|
|
43
|
+
def encoder(SMILES):
|
|
44
|
+
replaced_SMILES = re.sub(pattern, replacer, SMILES)
|
|
45
|
+
splitted_SMILES = split_and_modify_atoms(replaced_SMILES)
|
|
46
|
+
modified_SMILES = " ".join(
|
|
47
|
+
[re.sub(add_space_re, add_space, word) for word in splitted_SMILES.split()]
|
|
48
|
+
)
|
|
49
|
+
return modified_SMILES
|
|
50
|
+
|
|
51
|
+
|
|
52
|
+
def decoder(predictions):
|
|
53
|
+
modified = (
|
|
54
|
+
predictions.replace("!", "1")
|
|
55
|
+
.replace("$", "2")
|
|
56
|
+
.replace("^", "3")
|
|
57
|
+
.replace("<", "4")
|
|
58
|
+
.replace(">", "5")
|
|
59
|
+
.replace("?", "6")
|
|
60
|
+
.replace("£", "7")
|
|
61
|
+
.replace("¢", "8")
|
|
62
|
+
.replace("€", "9")
|
|
63
|
+
.replace("§", "0")
|
|
64
|
+
)
|
|
65
|
+
return modified
|
|
66
|
+
|
|
67
|
+
|
|
68
|
+
# Downloads the model and unzips the file downloaded, if the model is not present on the working directory.
|
|
69
|
+
def download_trained_weights(model_url: str, model_path: str, verbose=1):
|
|
70
|
+
"""This function downloads the trained models and tokenizers to a default
|
|
71
|
+
location. After downloading the zipped file the function unzips the file
|
|
72
|
+
automatically. If the model exists on the default location this function
|
|
73
|
+
will not work.
|
|
74
|
+
|
|
75
|
+
Args:
|
|
76
|
+
model_url (str): trained model url for downloading.
|
|
77
|
+
model_path (str): model default path to download.
|
|
78
|
+
|
|
79
|
+
Returns:
|
|
80
|
+
path (str): downloaded model.
|
|
81
|
+
"""
|
|
82
|
+
# Download trained models
|
|
83
|
+
if verbose > 0:
|
|
84
|
+
print("Downloading trained model to " + str(model_path))
|
|
85
|
+
model_path = pystow.ensure("DECIMER-V2", url=model_url)
|
|
86
|
+
if verbose > 0:
|
|
87
|
+
print(model_path)
|
|
88
|
+
print("... done downloading trained model!")
|
|
89
|
+
|
|
90
|
+
with zipfile.ZipFile(model_path.as_posix(), "r") as zip_ref:
|
|
91
|
+
zip_ref.extractall(model_path.parent.as_posix())
|
|
92
|
+
|
|
93
|
+
# Delete zipfile after downloading
|
|
94
|
+
if Path(model_path).exists():
|
|
95
|
+
Path(model_path).unlink()
|
|
96
|
+
|
|
97
|
+
|
|
98
|
+
def ensure_models(default_path: str, model_urls: dict) -> dict:
|
|
99
|
+
"""Function to ensure models are present locally.
|
|
100
|
+
|
|
101
|
+
Convenient function to ensure model downloads before usage.
|
|
102
|
+
Models are re-downloaded if the URL changes (e.g., new Zenodo record).
|
|
103
|
+
|
|
104
|
+
Args:
|
|
105
|
+
default_path (str): Default path for model data
|
|
106
|
+
model_urls (dict): Dictionary containing model names as keys and their corresponding URLs as values
|
|
107
|
+
|
|
108
|
+
Returns:
|
|
109
|
+
dict: A dictionary containing model names as keys and their local paths as values
|
|
110
|
+
"""
|
|
111
|
+
model_paths = {}
|
|
112
|
+
|
|
113
|
+
for model_name, model_url in model_urls.items():
|
|
114
|
+
model_path = os.path.join(default_path, f"{model_name}_model")
|
|
115
|
+
saved_model_file = os.path.join(model_path, "saved_model.pb")
|
|
116
|
+
version_file = os.path.join(model_path, ".model_url")
|
|
117
|
+
|
|
118
|
+
# Check if model needs to be downloaded:
|
|
119
|
+
# 1. saved_model.pb doesn't exist, OR
|
|
120
|
+
# 2. The URL has changed (model was updated)
|
|
121
|
+
needs_download = not os.path.exists(saved_model_file)
|
|
122
|
+
|
|
123
|
+
if not needs_download and os.path.exists(version_file):
|
|
124
|
+
with open(version_file, "r") as f:
|
|
125
|
+
cached_url = f.read().strip()
|
|
126
|
+
if cached_url != model_url:
|
|
127
|
+
needs_download = True
|
|
128
|
+
print(f"Model {model_name} has been updated, re-downloading...")
|
|
129
|
+
|
|
130
|
+
if needs_download:
|
|
131
|
+
# Clean up incomplete/corrupted model directory if it exists
|
|
132
|
+
if os.path.exists(model_path):
|
|
133
|
+
shutil.rmtree(model_path)
|
|
134
|
+
download_trained_weights(model_url, default_path)
|
|
135
|
+
|
|
136
|
+
# Store the URL used for this download
|
|
137
|
+
os.makedirs(model_path, exist_ok=True)
|
|
138
|
+
with open(version_file, "w") as f:
|
|
139
|
+
f.write(model_url)
|
|
140
|
+
|
|
141
|
+
model_paths[model_name] = model_path
|
|
142
|
+
|
|
143
|
+
return model_paths
|