File size: 16,508 Bytes
e1a215c
 
4b388b5
 
ee1feb4
cf1b2d8
 
 
 
e1a215c
4c0e7d1
 
4a45550
15610a4
e1a215c
abe5e88
4c0e7d1
abe5e88
4c0e7d1
abe5e88
4c0e7d1
 
e1a215c
abe5e88
4c0e7d1
e1a215c
 
 
4c0e7d1
e1a215c
 
abe5e88
5f81a43
abe5e88
 
e1a215c
 
 
4c0e7d1
e1a215c
4c0e7d1
 
e1a215c
 
 
 
 
 
 
abe5e88
 
 
 
e1a215c
 
 
 
 
030217f
e1a215c
abe5e88
e1a215c
 
 
abe5e88
e1a215c
 
 
 
 
 
030217f
 
4c0e7d1
 
e1a215c
4c0e7d1
e1a215c
 
 
 
 
abe5e88
e1a215c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
abe5e88
e1a215c
 
abe5e88
e1a215c
 
 
 
 
045ef38
 
 
 
 
 
 
 
 
 
 
 
abe5e88
e1a215c
 
4c0e7d1
 
e1a215c
 
4c0e7d1
045ef38
 
 
 
 
 
 
 
 
 
 
 
 
e1a215c
 
4c0e7d1
e1a215c
 
 
 
 
045ef38
e1a215c
 
4c0e7d1
045ef38
 
e1a215c
 
045ef38
 
e1a215c
 
 
 
 
 
045ef38
e1a215c
 
045ef38
e1a215c
 
045ef38
e1a215c
 
 
 
 
 
045ef38
e1a215c
 
045ef38
e1a215c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
045ef38
 
104a6f8
 
 
 
 
 
045ef38
e1a215c
045ef38
 
e1a215c
 
045ef38
 
e1a215c
 
 
045ef38
 
e1a215c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
4c0e7d1
e1a215c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
44a649d
 
045ef38
 
4c0e7d1
e1a215c
 
 
 
 
 
 
 
4c0e7d1
9d4188b
e1a215c
 
 
9d4188b
e1a215c
 
 
4c0e7d1
 
 
abe5e88
 
e1a215c
 
 
 
 
 
abe5e88
4b388b5
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
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
---
license: apache-2.0
pipeline_tag: text-generation
widget:
- text: "1.1.1.21<sep><start>"
inference:
  parameters:
    top_k: 9
    repetition_penalty: 1.2
---
# **ZymCTRL**

ZymCTRL (Enzyme Control) ([Paper presented @ NeurIPS Machine Learning for Structural Biology workshop - December 2022](https://www.mlsb.io/papers_2022/ZymCTRL_a_conditional_language_model_for_the_controllable_generation_of_artificial_enzymes.pdf)) 
is a conditional language model for the generation of artificial functional enzymes. It was trained on the entire BRENDA database of enzymes, comprising over 37 M sequences. 
Given a user-defined Enzymatic Commission (EC) number, the model generates protein sequences that fulfill that catalytic reaction. 
The generated sequences are ordered, globular, and distant to natural ones, while their intended catalytic properties match those defined by users.

If you don't know the EC number of your protein of interest, have a look at the BRENDA webpage: https://www.brenda-enzymes.org/ecexplorer.php?browser=1

See below for information about the model, how to generate sequences, and how to save and rank them by perplexity.

## **Model description**
ZymCTRL is based on the [CTRL Transformer](https://arxiv.org/abs/1909.05858) architecture (which in turn is very similar to ChatGPT) and contains 36 layers
with a model dimensionality of 1280, totaling 738 million parameters. 

ZymCTRL is a decoder-only transformer model pre-trained on the BRENDA database
(version July 2022). The pre-training was done on the raw sequences without FASTA headers,
with the EC classes prepended to each sequence. The databases will be uploaded soon.

ZymCTRL was trained with an autoregressive objective, i.e., the model learns to predict
the next token given a sequence context. Because the first tokens on each sequence encode the EC numbers,
the model learns the dependencies among EC classes and their corresponding sequences and is able to _speak_ the enzyme language.

There are stark differences in the number of members among EC classes, and for this reason, we also tokenized the EC numbers.
In this manner, EC numbers '2.7.1.1' and '2.7.1.2' share the first three tokens (six, including separators), and hence the model can infer that
there are relationships between the two classes.

The figure below summarizes the process of training:

![plot](./github1.png)

 
## **How to use ZymCTRL**
ZymCTRL can be used with the HuggingFace transformer python package.
Detailed installation instructions can be found here: https://huggingface.co/docs/transformers/installation

Since ZymCTRL has been trained on the classical language model objective on enzyme sequences with their EC annotation,
it particularly excels at generating enzyme sequences given a user-defined EC class, such as alcohol dehydrogenases ('1.1.1.2').

The model can generate in two ways: in a zero-shot fashion, i.e., directly generating from the checkpoint weights, or after fine-tuning. 
Fine-tuning allows augmenting the BRENDA datasets that were used during training, for example, 
if you have a curated internal dataset or a set of ancestrally-reconstructed sequences. This is entirely optional. One advantage of 
running the model in zero-shot is that it doesn't require any further training.


### **Example 1: Generating nitrilases (EC 3.5.5.1)**  

The script below will be used for the generation of any BRENDA class in a zero-shot fashion, 
here we showcase the generation of novel nitrilases.

To run this script, you should download ZymCTRL to a local folder in your workstation. 
Then replace the placeholders in the script with your actual folder path.

You can run it directly in the command line (once you have hugging face installed), 
with the following command: `python generate.py`

The script will write each sequence in a fasta file in the folder you specify. In the fasta header,
it will store the sequence's computed perplexity value. Perplexity is a measure of the model's confidence
in that generation, with lower values being better. The sequences are ordered by perplexity before writing them out,
so those that finish in *_0.fasta and *_1.fasta will be the best ones per batch. 

**Given that generation runs so fast, we recommend generating hundreds or thousands and then only picking the best 5% or less. 
With the script below, that would mean picking only those that finish in '_0.fasta'. Good perplexity values for this model so be below 1.75-1.5.** 

```
import torch
from transformers import GPT2LMHeadModel, AutoTokenizer
import os
from tqdm import tqdm
import math

def remove_characters(sequence, char_list):
    "This function removes special tokens used during training."
    columns = sequence.split('<sep>')
    seq = columns[1]
    for char in char_list:
        seq = seq.replace(char, '')
    return seq

def calculatePerplexity(input_ids,model,tokenizer):
    "This function computes perplexities for the generated sequences"
    with torch.no_grad():
        outputs = model(input_ids, labels=input_ids)
    loss, logits = outputs[:2]
    return math.exp(loss)
        
def main(label, model,special_tokens,device,tokenizer):
    # Generating sequences
    input_ids = tokenizer.encode(label,return_tensors='pt').to(device)
    outputs = model.generate(
        input_ids, 
    	top_k=9, #tbd
        repetition_penalty=1.2,
        max_length=1024,
        eos_token_id=1,
        pad_token_id=0,
   	    do_sample=True,
   	    num_return_sequences=20) # Depending non your GPU, you'll be able to generate fewer or more sequences. This runs in an A40.
    
    # Check sequence sanity, ensure sequences are not-truncated.
    # The model will truncate sequences longer than the specified max_length (1024 above). We want to avoid those sequences.
    new_outputs = [ output for output in outputs if output[-1] == 0]
    if not new_outputs:
        print("not enough sequences with short lengths!!")

    # Compute perplexity for every generated sequence in the batch
    ppls = [(tokenizer.decode(output), calculatePerplexity(output, model, tokenizer)) for output in new_outputs ]

    # Sort the batch by perplexity, the lower the better
    ppls.sort(key=lambda i:i[1]) # duplicated sequences?

    # Final dictionary with the results
    sequences={}
    sequences[label] = [(remove_characters(x[0], special_tokens), x[1]) for x in ppls]

    return sequences

if __name__=='__main__':
    device = torch.device("cuda") # Replace with 'cpu' if you don't have a GPU - but it will be slow
    print('Reading pretrained model and tokenizer')
    tokenizer = AutoTokenizer.from_pretrained('/path/to/zymCTRL/') # change to ZymCTRL location
    model = GPT2LMHeadModel.from_pretrained('/path/to/zymCTRL').to(device) # change to ZymCTRL location
    special_tokens = ['<start>', '<end>', '<|endoftext|>','<pad>',' ', '<sep>']

    # change to the appropriate BRENDA EC classes
    labels=['3.5.5.1'] # nitrilases. You can put as many labels as you want.

    for label in tqdm(labels):
        # We'll run 100 batches per label. 20 sequences will be generated per batch.
        for i in range(0,100): 
            sequences = main(label, model, special_tokens, device, tokenizer)
            for key,value in sequences.items():
                for index, val in enumerate(value):
                    # Sequences will be saved with the name of the label followed by the batch index,
                    # and the order of the sequence in that batch.           
                    fn = open(f"/path/to/folder/{label}_{i}_{index}.fasta", "w")
                    fn.write(f'>{label}_{i}_{index}\t{val[1]}\n{val[0]}')
                    fn.close()
```
 

## **Example 2: Fine-tuning on a set of user-defined sequences**  

This alternative to the zero-shot generation allows updating ZymCTRL's weights to new sequences.

This strategy is not strictly necessary, in fact, we have observed good generations even for EC classes where there are 
only 1-2 representatives in Nature. But you might have an internal set of sequences that you'd like to incorporate into the model.
For example, internal datasets after protein engineering efforts, 
ancestrally-reconstructed sets, or after searching against metagenomics databases. In these cases, it is advisable to fine-tune ZymCTRL,
as it will learn new properties from your dataset and potentially improve the generation quality 
(especially for poorly populated EC classes).

To fine-tune ZymCTRL, you can use the script below to process your sequences. The only requisite is to start with an input file,
'sequences.fasta' which contains all the sequences in a fasta format. Please follow the format below. There should not be new lines '\n' or 
any separator between sequences. In the script, change the variable ec_label to the specific BRENDA class you'd like to fine-tune.
The script will produce a file called {ec_label}_processed.txt and a folder with the training and validation datasets (split 10%)
```
>Sequence1
MMMMYMPLKVCD..
>Sequence2
MQWMXMYMPLKVCD..
>Sequence3
MPLKVCWMXMYMPLD..
```
We recommend using at least 200 sequences to obtain the best results. But we've seen it working with fewer sequences, so if you don't have
that many, give it still a go.


```
import random
from transformers import AutoTokenizer

from datasets import load_dataset
import transformers
from transformers.testing_utils import CaptureLogger

## DEFINE THESE VARIABLES
tokenizer = AutoTokenizer.from_pretrained('AI4PD/ZymCTRL')
ec_label = '1.1.1.1' # CHANGE TO YOUR LABEL
validation_split_percentage = 10 # change if you want
sequence_file = 'sequence.fasta'


#Load sequences, Read source file
with open(sequence_file, 'r') as fn: #! CHANGE TO SEQUENCES.FASTA
    data = fn.readlines()
    fn.close()

# Put sequences into dictionary
sequences={}
for line in data:
    if '>' in line:
        name = line.strip()
        sequences[name] = []  #! CHANGE TO corre
        continue
    sequences[name].append(line.strip())

#Pass sequences to list and shuffle their order randomly
sequences_list = [(key,value[0]) for key,value in sequences.items()]
random.shuffle(sequences_list)

#the objective is to get here strings, that when tokenized, would span a length of 1024.
#for each sequence group its length and untokenized string
print("procesing dataset")
processed_dataset = []
for i in sequences_list:
    # length of the control code
    sequence = i[1].strip()
    separator = '<sep>'
    control_code_length = len(tokenizer(ec_label+separator)['input_ids'])
    available_space = 1021 - control_code_length # It is not 1024 because '<|endoftext|>', and start and end

    # Option 1: the sequence is larger than the available space (3-4% of sequences)
    if len(sequence) > available_space:
        total_length = control_code_length + len(sequence[:available_space]) + 1
        seq = f"{ec_label}{separator}{sequence[:available_space]}<|endoftext|>"
        processed_dataset.append((total_length, seq))

    # Option 2 & 3: The sequence fits in the block_size space with or without padding
    else:
        total_length = control_code_length + len(sequence) + 3
        # in this case the sequence does not fit with the start/end tokens
        seq = f"{ec_label}{separator}<start>{sequence}<end><|endoftext|>"
        processed_dataset.append((total_length, seq))

# Group sequences
def grouper(iterable):
    prev = None
    group = ''
    total_sum = 0
    for item in iterable:
        if prev is None or item[0] + total_sum < 1025:
            group += item[1]
            total_sum += item[0]
        else:
            total_sum = item[0]
            yield group
            group = item[1]
        prev = item
    if group:
        total_sum = 0
        yield group

print("grouping processed dataset")
grouped_dataset=dict(enumerate(grouper(processed_dataset),1))

# Write file out for the tokenizer to read
fn = open(f"{ec_label}_processed.txt",'w')
for key,value in grouped_dataset.items():
    padding_len = 1024 - len(tokenizer(value)['input_ids'])
    padding = "<pad>"*padding_len
    fn.write(value+padding)
    fn.write
    fn.write("\n")
fn.close()

##TOKENIZE
# adapted from the trainer file
data_files = {}
dataset_args = {}

data_files["train"] = f"{ec_label}_processed.txt"
extension = "text"
tok_logger = transformers.utils.logging.get_logger("transformers.tokenization_utils_base")

raw_datasets = load_dataset(extension, data_files=data_files, cache_dir='.', **dataset_args)

raw_datasets["train"] = load_dataset(extension,
                data_files=data_files,
                split=f"train[{validation_split_percentage}%:]",
                cache_dir='.',
                **dataset_args,)

raw_datasets["validation"] = load_dataset(extension,
                                          data_files=data_files,
                                          split=f"train[:{validation_split_percentage}%]",
                                          cache_dir='.',
                                          **dataset_args,)

def tokenize_function(examples):
    with CaptureLogger(tok_logger) as cl:
        output = tokenizer(examples["text"])
    # clm input could be much much longer than block_size
    if "Token indices sequence length is longer than the" in cl.out:
        tok_logger.warning(
            "^^^^^^^^^^^^^^^^ Please ignore the warning above - this long input will be chunked into smaller bits before being passed to the model."
        )
    return output

tokenized_datasets = raw_datasets.map(
    tokenize_function,
    batched=True,
    num_proc=32,
    remove_columns=['text'],
    load_from_cache_file = False,
    desc="Running tokenizer on dataset",
)

block_size = 1024
def group_texts(examples):
    # Concatenate all texts.
    concatenated_examples = {k: sum(examples[k], []) for k in examples.keys()}
    total_length = len(concatenated_examples[list(examples.keys())[0]])
    # We drop the small remainder, we could add padding if the model supported it instead of this drop,
    # you can customize this part to your needs.
    if total_length >= block_size:
        total_length = (total_length // block_size) * block_size
    # Split by chunks of max_len.
    result = {
        k: [t[i : i + block_size] for i in range(0, total_length, block_size)]
        for k, t in concatenated_examples.items()
    }
    result["labels"] = result["input_ids"].copy()
    return result

lm_datasets = tokenized_datasets.map(
    group_texts,
    batched=True,
    num_proc=124,
    load_from_cache_file=False,
    desc=f"Grouping texts in chunks of {block_size}",
)

train_dataset = lm_datasets["train"]
eval_dataset = lm_datasets["validation"]

train_dataset.save_to_disk('./dataset/train2')
eval_dataset.save_to_disk('./dataset/eval2')


```
The processed datasets will be inside the folder dataset/, called train2 and eval2.
You could also put the two previous scripts into a single one and run it in one go (that is what we do).

Now you are ready to fine-tune the model. 
To do that, you can take the trainer file that we provide in this repository (5.run_clm-post.py), or use the trainer from Hugging Face. 
The command below shows an example at an specific learning rate, 
but you could try with other hyperparameters to obtain the best training and evaluation losses.

```
python 5.run_clm-post.py --tokenizer_name AI4PD/ZymCTRL
 --do_train --do_eval --output_dir output --evaluation_strategy steps --eval_steps 10
 --logging_steps 5 --save_steps 500 --num_train_epochs 28 --per_device_train_batch_size 1
 --per_device_eval_batch_size 4 --cache_dir '.' --save_total_limit 2 --learning_rate  0.8e-04
 --dataloader_drop_last True --model_name_or_path AI4PD/ZymCTRL
```
In any case, the original HuggingFace script run_clm.py can be found here: 
https://github.com/huggingface/transformers/blob/master/examples/pytorch/language-modeling/run_clm.py


### **Training specs**
The model was trained on 48 NVIDIA A100 GPUs for eight epochs, 
using a block size of 1024 and a total batch size of 768. 
The optimizer used was Adam (beta1 = 0.9, beta2 = 0.999) 
with a learning rate of 0.8e-04.

### **Contact**

We are the AI for Protein Design group at the Institute of Molecular Biology of Barcelona (https://www.aiproteindesign.com/).
For any questions post an issue in this repository so that other people can benefit from the feedback, and I'll get back to you shortly.
We are always open for collaborations, send an email to nfccri [at] ibmb [dot] csic [dot] es