Update README.md

README.md

@@ -198,45 +198,156 @@ but this score would enforce a high sequence similarity (thus not *de novo* desi
 We recommend generating many sequences and selecting them by plDDT, as well as other metrics.
 
 ```python
-from datasets import load_from_disk
-from transformers import AutoTokenizer
-from transformers import T5Tokenizer, T5ForConditionalGeneration
-import math
-import torch
-from tqdm import tqdm
-import pickle
-tokenizer_aa = AutoTokenizer.from_pretrained('/path/to//tokenizer_aa')
-tokenizer_smiles = AutoTokenizer.from_pretrained('/path/to//tokenizer_smiles')
-
-model = T5ForConditionalGeneration.from_pretrained("/path/to/REXzyme").cuda()
-print(model.generation_config)
-reactions = ["NC1=NC=NC2=C1N=CN2[C@@H]1O[C@H](COP(=O)([O-])OP(=O)([O-])OP(=O)([O-])[O-])[C@@H](O)[C@H]1O.*N[C@@H](CO)C(*)=O>>NC1=NC=NC2=C1N=CN2[C@@H]1O[C@H](COP(=O)([O-])OP(=O)([O-])[O-])[C@@H](O)[C@H]1O.[H+].*N[C@@H](COP(=O)([O-])[O-])C(*)=O"]
-
-def calculatePerplexity(inputs,model):
-    '''Function to compute perplexity'''
-    a=tokenizer_aa.decode(inputs)
-    b=tokenizer_aa(a, return_tensors="pt").input_ids.to(device='cuda')
-    b = torch.stack([[b[b!=tokenizer_aa.pad_token_id]] for label in b][0])
-    with torch.no_grad():
-        outputs = model(b, labels=b)
-    loss, logits = outputs[:2]
-    return math.exp(loss)
-
-
-for idx,i in tqdm(enumerate(reactions)):
-    input_ids = tokenizer_smiles(f"r2s{i}</s>", return_tensors="pt").input_ids.to(device='cuda')
-    print(f'Generating for {i}')
-    ppls_total = []
-    for _ in range(4):
-        outputs = model.generate(input_ids,
-                top_k=15,
-                top_p = 0.92,
-                repetition_penalty=1.2,
-                max_length=1024,
-                do_sample=True,
-                num_return_sequences=25)
-        ppls = [(tokenizer_aa.decode(output,skip_special_tokens=True), calculatePerplexity(output, model),len(tokenizer_aa.decode(output))) for output in tqdm(outputs)]
-        ppls_total.extend(ppls)
+"""Inference on a SMILES txt. Saved as fastas
+Previously called generate_comparison"""
+
+if __name__ == '__main__':
+    from transformers import AutoTokenizer, AutoModelForSeq2SeqLM,AutoModelForCausalLM  #T5ForConditionalGeneration
+    import argparse
+    import os
+    import torch
+    import json
+
+    parser = argparse.ArgumentParser(description='Mol2Pro inference',
+                                         formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    parser.add_argument('--input_file', default='../inference/random_smiles2.txt', type=str,
+                        help='File with the input molecule SMILES')
+    parser.add_argument('--model_path', default='./output03/checkpoint-60000', type=str, help='Path to model to load')
+    parser.add_argument('--tokenizer_aa',
+                        default='/home/woody/b114cb/b114cb10/mol2pro/1.training-different-sizes/1.all-data-16M-tokenizernuria/tokenizer_aa', type=str,
+                        help='Path to amino acid tokenizer')
+    parser.add_argument('--tokenizer_mol',
+                        default='/home/woody/b114cb/b114cb10/mol2pro/1.training-different-sizes/1.all-data-16M-tokenizernuria/nuria_tokenizer_smiles', type=str,
+                        help='Path to SMILES tokenizer')
+    parser.add_argument('--top_k',
+                        default=15,type=int,
+                        help='K for top-k sampling')
+    parser.add_argument('--output_folder', default='fastas', type=str, help='Folder for saving results')
+    args = parser.parse_args()
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    if 'gatgpt' in args.model_path.lower():
+        GNN = True
+        print('Graph data mode')
+    else:
+        GNN = False
+        print('SMILES/SELFIES data mode')
+
+    # Load protein tokenizer
+    if 'ape' in args.tokenizer_aa:
+        from ape_tokenizer import APETokenizer
+
+        tokenizer_aa = APETokenizer.from_pretrained(args.tokenizer_aa)
+    else:
+        tokenizer_aa = AutoTokenizer.from_pretrained(args.tokenizer_aa)
+
+    # Load molecule tokenizer
+    if GNN:
+        tokenizer_mol = None
+    else:
+        if 'ape' in args.tokenizer_mol:
+            from ape_tokenizer import APETokenizer
+
+            tokenizer_mol = APETokenizer.from_pretrained(args.tokenizer_mol)
+        else:
+            tokenizer_mol = AutoTokenizer.from_pretrained(args.tokenizer_mol)
+
+    # Load model
+    dec_only = False
+    if GNN:
+        from transformers import GPT2Config, Trainer
+        from models import GATGPT2Config, GATGPT2
+        from torch_geometric.data import Batch, Data
+
+        config = GATGPT2Config.from_pretrained(args.model_path)
+
+        # Load model weights
+        model = GATGPT2.from_pretrained(args.model_path, config=config)
+
+        model.eval()
+        model.to("cuda" if torch.cuda.is_available() else "cpu")
+    else:
+        try:
+            print('Attempt Seq2Seq model load... ')
+            model = AutoModelForSeq2SeqLM.from_pretrained(args.model_path).cuda()
+        except:
+            print('Attempt CausalLM model load... ')
+            model = AutoModelForCausalLM.from_pretrained(args.model_path).cuda()
+            model.config.eos_token_id = tokenizer_mol.eos_token_id
+            model.config.pad_token_id = tokenizer_mol.pad_token_id
+            print(
+                f"Set `eos_token_id` to {tokenizer_mol.eos_token_id} and `pad_token_id` to {tokenizer_mol.pad_token_id}.")
+            dec_only = True
+    print('Model Loaded')
+
+
+    smiles_list = []
+    with open(args.input_file, 'r') as input_file:
+        for line in input_file:
+            smiles_list.append(line.strip())
+
+    molecule_json = {}
+    for index,smiles in enumerate(smiles_list):
+        sequences=[]
+        if GNN:
+            from build_tokenized_dataset import convert_smiles_to_graph
+
+            node_feats, edge_index, edge_feats = convert_smiles_to_graph(smiles)
+            node_feats_tensor = torch.tensor(node_feats, dtype=torch.float, device=device)
+            edge_index_tensor = torch.tensor(edge_index, dtype=torch.long, device=device).T.contiguous()
+            edge_feats_tensor = torch.tensor(edge_feats, dtype=torch.float, device=device)
+
+            # Input to decoder is only bos
+            start_token = tokenizer_aa.bos_token_id or tokenizer_aa.convert_tokens_to_ids("▁")  # fallback to the space which is always appended by our tokenizer
+            text_input_ids = torch.tensor([[start_token]], dtype=torch.long, device=device)
+
+            input_ids = {
+                "graph_node_feats": node_feats_tensor,  # shape (N, 3)
+                "graph_edge_index": edge_index_tensor,  # shape (2, E)
+                "graph_edge_feats": edge_feats_tensor,  # shape (E, 2)
+                "batch": torch.full((len(node_feats),), 0, dtype=torch.long, device=device),  # shape (N,)
+                "input_ids": text_input_ids
+            }
+
+        elif 'ape' in args.tokenizer_mol:
+            input_ids = tokenizer_mol(smiles, return_tensors="pt")["input_ids"].to(device='cuda')
+        else:
+            input_ids = tokenizer_mol(smiles, return_tensors="pt").input_ids.to(device='cuda')
+        if not GNN:
+            print(f'Generating for {smiles} (input ids: {input_ids})')
+        else:
+            print(f'Generating for {smiles}')
+
+        # top_k = Choose at random from the first K tokens (weigthed by softmax score)
+        # num_return_sequences = The number of independently computed returned sequences for each element in the batch.
+        if dec_only:
+            attention_mask = torch.ones_like(input_ids).cuda()
+            outputs = model.generate(input_ids, top_k=args.top_k, attention_mask = attention_mask, repetition_penalty=1.2, max_length=1024, do_sample=True, num_return_sequences=25)
+        else:
+            outputs = model.generate(input_ids, top_k=args.top_k, repetition_penalty=1.2, max_length=1024, do_sample=True, num_return_sequences=25)
+
+        sequences = [tokenizer_aa.decode(output, skip_special_tokens=True) for output in outputs]
+
+        if not os.path.exists(args.output_folder):
+            os.makedirs(args.output_folder)
+
+        filename = f'{args.output_folder}/output_topk{args.top_k}_file-{index}.fasta'
+        with open(filename, 'w') as fn:
+            for idx, seq in enumerate(sequences):
+                fn.write(f">{idx}\n{seq}\n")
+
+        # Store molecule name
+        molecule_json[filename] = smiles
+
+    # Save metadata
+    metadata_path = os.path.join(args.output_folder, 'molecule_input_metadata.json')
+    try:
+        with open(metadata_path, 'w') as json_file:
+            json.dump(molecule_json, json_file, indent=4)
+        print(f"Metadata successfully written to {metadata_path}")
+    except Exception as e:
+        print(f"An error occurred while writing to JSON: {e}")
 ```
 
 ## **A word of caution**