Update README.md

README.md

@@ -1,83 +1,385 @@
+---
+license: apache-2.0
+---
 # **ZymCTRL**
 
-ZymCTRL ([Paper presented @ Machine Learning for Structural Biology workshop](https://www.mlsb.io/papers_2022/ZymCTRL_a_conditional_language_model_for_the_controllable_generation_of_artificial_enzymes.pdf)) is a conditional language model trained on the BRENDA database of enzymes. 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.
+ZymCTRL ([Paper presented @ Machine Learning for Structural Biology workshop](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 BRENDA database of enzymes. 
+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 what 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 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 architecture and contains 36 layers with a model dimensionality of 1280, totalling 738 million parameters. 
+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, totalling 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 can be found here: xx.
+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.
+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.
 
-Because there are stark differences in the number of members among EC classes, 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).
- 
-### **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
+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:
 
-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 '1.1.1.2'. Besides, it can also be fine-tuned on a specific catalytic reaction by providing more sequences for a given EC class, such as sequences obtained with ancestral reconstruction methods.
+![plot](./github1.png)
 
-**Example 1: Generating glucose oxidases (EC 1.1.3.4)**  
  
-In the example below, ZymCTRL generates sequences that catalyze the reaction encoded by the EC number 1.1.3.4. Any other EC class can also be prompted instead. The model will generate the most probable sequences that follow the input.
+## **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 to augment the BRENDA datasets that were using 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 dehalogenases.
+
+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 to generate hundreds or thousands and then only pick the best 5%. 
+With the script below that would mean picking only those that finish in '_0.fasta'**
 
 ```
+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 to update 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 to 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 will need to process your sequences quite a bit. With the scripts below can exactly do that without many
+modifications. The only requisite is to start with an input file 'sequences.fasta' which contain all the sequences in a fasta format. 
+We recommend using at least 200 sequences to obtain best results. But we've seen it working with fewer sequences, so if you don't have
+that many, give it still a go.
+
 
-enzyme_class = 1.1.3.4
-device = torch.device("cuda") # if a GPU is available
-tokenizer = AutoTokenizer.from_pretrained('/path/to/tokenizer')
-model = GPT2LMHeadModel.from_pretrained('/path/to/output').to(device)
-input_ids = tokenizer.encode(enzyme_class,return_tensors='pt').to(device)
-# change max_length or num_return_sequences to your requirements
-output = model.generate(input_ids, top_k=9, repetition_penalty=1.2, max_length=1024,
-                        eos_token_id=1,pad_token_id=0,do_sample=True, num_return_sequences=100)
 ```
+import random
+import transformers
+from transformers import AutoTokenizer
 
-**Example 2: Fine-tuning on a set of user-defined sequences**  
+# 1. Read the source file
+with open('sequences.fasta', 'r') as fn:
+    data = fn.readlines()
+    fn.close()
 
-This alternative option to the zero-shot generation permits further improve the model's confidence for EC number with few members. User-defined training and validation files containing the sequences of interest are provided to the model. After a short update of the model's weights, ZymCTRL will generate sequences that follow the input properties. This might not be necessary in cases where the model has already seen many sequences per EC class.
+# Put sequences into dictionary
+sequences={}
+for line in data:
+    if '>' in line:
+        name = line.strip()
+        sequences[name] = ['2.7.3.12'] # modify with the actual EC class.
+        continue
+    sequences[name].append(line.strip())
 
-To create the validation and training file, it is necessary to 
-(1) remove the FASTA headers for each sequence, 
-(2) prepare the sequences in the format: `EC number<sep><start>S E Q U E N C E<end><|endoftext|>` and (3) split the originating dataset into training and validation files (this is often done with the ratio 90/10, 80/20 or 95/5). Then, to fine-tune the model to the input sequences, we can use the example below. Here we show a learning rate of 1e-06, but ideally, the learning rate should be optimised in separate runs. After training, the fine-tuned model will be stored in the ./output folder. Lastly, ZymCTRL can generate the tailored sequences as shown in Example 1:
+# Process fasta files to be in single string - run this part only if the fastas were formated to 60 characters
+processed_sequences = {}
+for name, sequence in sequences.items():
+    processed_sequences[f"{sequence[0]};{name}"] = ''.join([x for x in sequence[1:]])
 
+# Shuffle sequences
+sequences_list = [(key,value) for key,value in processed_sequences.items()]
+random.shuffle(sequences_list)
+
+# Load tokenizer
+tokenizer = AutoTokenizer.from_pretrained('/path/to/ZymCTRL')
+
+# the objective is to get here strings, that when tokenized, will span a window 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
+    label = i[0].split(';')[0]
+    sequence = i[1].strip()
+    separator = '<sep>'
+    control_code_length = len(tokenizer(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 in BRENDA are over 1024)
+    if len(sequence) > available_space:
+        total_length = control_code_length + len(sequence[:available_space]) + 1
+        seq = f"{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"{label}{separator}<start>{sequence}<end><|endoftext|>"
+        processed_dataset.append((total_length, seq))
+
+# Helper function to 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
+
+# Group sequences
+print("grouping processed dataset")
+grouped_dataset=dict(enumerate(grouper(processed_dataset),1))
+
+# Save the processed file out
+fn = open("./2.7.3.13_processed.txt",'w')
+for key,value in grouped_dataset.items():
+    fn.write(value)
+    fn.write("\n")
+fn.close()    
 ```
-python run_clm.py --model_name_or_path nferruz/ZymCTRL --train_file training.txt --validation_file validation.txt --tokenizer_name nferruz/ZymCTRL
- --do_train --do_eval --output_dir output --learning_rate 1e-06 
+The previous script will prepare a text file with the correct format for tokenization. 
+Now we can use the tokenizer to convert its contents to tokens.
 
 ```
-The HuggingFace script run_clm.py can be found here: https://github.com/huggingface/transformers/blob/master/examples/pytorch/language-modeling/run_clm.py
+from datasets import load_dataset
+import transformers
+from transformers.testing_utils import CaptureLogger
+
+# Load the tokenizer again
+from transformers import AutoTokenizer
+tokenizer = AutoTokenizer.from_pretrained('/agh/projects/noelia/NLP/zymCTRL/dataset_preparation/tokenizer')
+
+
+#Load the data files
+data_files = {}
+dataset_args = {}
+validation_split_percentage = 10 # for a split 90/10
+data_files["train"] = './2.7.3.12_processed.txt'
+extension = "text"
+raw_datasets = load_dataset(extension, data_files=data_files, cache_dir='.', **dataset_args)
+tok_logger = transformers.utils.logging.get_logger("transformers.tokenization_utils_base")
+
+# Load datasets using the HF datasets library:
+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):
+    " This function tokenizes input"
+    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
+
+# tokenize in parallel
+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",
+)
 
-### **How to select the best sequences**
+train_dataset = tokenized_datasets["train"]
+eval_dataset = tokenized_datasets["validation"]
 
-First, we recommend selecting only sequences where the padding token has been emitted. Because the generation occurs with a max_length parameter, the HuggingFace generation function will truncate sequences that surpass that length. Once the sequence has been emitted, select those with at least one <pad> token at the end. Otherwise, you might be seeing truncated sequences by the length limit.
+train_dataset.save_to_disk('./dataset/train')
+eval_dataset.save_to_disk('./dataset/eval')
 
-Besides, we've observed that perplexity values correlate with AlphaFold2's plddt. 
-We recommend computing perplexity for each sequence as follows:
+# This has saved the datasets tokenized. Now we need to group them into the block size of 1024
+from datasets import load_from_disk
 
+train_dataset = load_from_disk('./2.7.3.13/dataset/train')
+eval_dataset = load_from_disk('./2.7.3.13/dataset/eval')
+
+from datasets.dataset_dict import DatasetDict
+tokenized_datasets = DatasetDict()
+
+tokenized_datasets["train"] = train_dataset
+tokenized_datasets["validation"] = eval_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')
 ```
-def calculatePerplexity(sequence, model, tokenizer):
-    with torch.no_grad():
-        outputs = model(sequence, labels=input_ids)
-   loss, logits = outputs[:2]
-    return math.exp(loss)
-    
-# Generate sequences by loading model and tokenizer (previously downloaded)
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-tokenizer = AutoTokenizer.from_pretrained('/path/to/tokenizer') # replace with the actual path
-model = GPT2LMHeadModel.from_pretrained('/path/to/ZymCTRL').to(device) 
-output = model.generate("1.1.1.1", max_length=400, do_sample=True, top_k=8, repetition_penalty=1.2, num_return_sequences=10, eos_token_id=0)
-
-# Take (for example) the first sequence
-sequence = output[0]
-ppl = calculatePerplexity(sequence, model, tokenizer)
+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 /path/to/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_type gpt2 --config_name /path/to/ZymCTRL
+ --gradient_accumulation_steps 4
 
-Where `ppl` is a value with the perplexity for that sequence.
-We do not yet have a threshold as to what perplexity value gives a 'good' or 'bad' sequence, but given the fast inference times, the best is to sample many sequences, order them by perplexity, and select those with the lower values (the lower the better).
+```
+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 8 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.
+The model was trained on 48 NVIDIA A100 GPUs for 8 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 question 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
+