Update README.md

README.md

@@ -13,4 +13,220 @@ tags:
 - chemistry
 - biology
 - medical
----
+---
+# DrugGen: Advancing Drug Discovery with Large Language Models and Reinforcement Learning Feedback
+
+DrugGen is a GPT-2 based model specialized for generating drug-like SMILES structures based on protein sequence. The model leverages the characteristics of approved drug targets and has been trained through both supervised fine-tuning and reinforcement learning techniques to enhance its ability to generate chemically valid, safe, and effective structures.
+
+## Model Details
+
+-  Model Name: DrugGen
+-  Training Paradigm: Supervised Fine-Tuning (SFT) + Proximal Policy Optimization (PPO)
+-  Input: Protein Sequence
+-  Output: SMILES Structure
+-  Training Libraries: Hugging Face’s transformers and Transformer Reinforcement Learning (TRL)
+-  Model Sources: liyuesen/druggpt
+
+## How to Get Started with the Model
+```python
+import pandas as pd
+from transformers import AutoTokenizer, GPT2LMHeadModel
+from datasets import load_dataset
+
+class SMILESGeneator:
+    def __init__(self):
+        
+        # Configuration parameters
+        self.config = {
+            "model_name": "alimotahharynia/DrugGen",
+            "dataset_name": "alimotahharynia/approved_drug_target",
+            "dataset_key": "uniprot_sequence",
+            "generation_kwargs": {
+                "do_sample": True,
+                "top_k": 9,
+                "max_length": 1024,
+                "top_p": 0.9,
+                "num_return_sequences": 10
+            },
+            "max_retries": 30  # Max retry limit to avoid infinite loops
+        }
+
+        # Load model and tokenizer
+        self.model_name = self.config["model_name"]
+        self.model, self.tokenizer = self.load_model_and_tokenizer(self.model_name)
+
+        # Load UniProt mapping dataset
+        dataset_name = self.config["dataset_name"]
+        dataset_key = self.config["dataset_key"]
+        self.uniprot_to_sequence = self.load_uniprot_mapping(dataset_name, dataset_key)
+
+        # Adjust generation parameters with token IDs
+        self.generation_kwargs = self.config["generation_kwargs"]
+        self.generation_kwargs["bos_token_id"] = self.tokenizer.bos_token_id
+        self.generation_kwargs["eos_token_id"] = self.tokenizer.eos_token_id
+        self.generation_kwargs["pad_token_id"] = self.tokenizer.pad_token_id
+
+    def load_model_and_tokenizer(self, model_name):
+
+        print(f"Loading model and tokenizer: {model_name}")
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        model = GPT2LMHeadModel.from_pretrained(model_name)
+        return model, tokenizer
+
+    def load_uniprot_mapping(self, dataset_name, dataset_key):
+ 
+        print(f"Loading dataset: {dataset_name}")
+        try:
+            dataset = load_dataset(dataset_name, dataset_key)
+            return {row["UniProt_id"]: row["Sequence"] for row in dataset["uniprot_seq"]}
+        except Exception as e:
+            raise RuntimeError(f"Failed to load dataset {dataset_name}: {e}")
+
+    def generate_smiles(self, sequence, num_generated):
+        """
+        Generate unique SMILES with a retry limit to avoid infinite loops.
+        """
+        generated_smiles_set = set()
+        prompt = f"<|startoftext|><P>{sequence}<L>"
+        encoded_prompt = self.tokenizer(prompt, return_tensors="pt")["input_ids"]
+        retries = 0
+
+        while len(generated_smiles_set) < num_generated:
+            if retries >= self.config["max_retries"]:
+                print("Max retries reached. Returning what has been generated so far.")
+                break
+            
+            sample_outputs = self.model.generate(encoded_prompt, **self.generation_kwargs)
+            for sample_output in sample_outputs:
+                output_decode = self.tokenizer.decode(sample_output, skip_special_tokens=False)
+                try:
+                    generated_smiles = output_decode.split("<L>")[1].split("<|endoftext|>")[0]
+                    if generated_smiles not in generated_smiles_set:
+                        generated_smiles_set.add(generated_smiles)
+                except IndexError:
+                    continue
+            
+            retries += 1
+        
+        return list(generated_smiles_set)
+
+    def generate_smiles_data(self, list_of_sequences=None, list_of_uniprot_ids=None, num_generated=10):
+        """
+        Generate SMILES data for sequences or UniProt IDs.
+        """
+        if not list_of_sequences and not list_of_uniprot_ids:
+            raise ValueError("Either `list_of_sequences` or `list_of_uniprot_ids` must be provided.")
+
+        # Prepare sequences input
+        if list_of_sequences:
+            sequences_input = list_of_sequences
+        else:
+            sequences_input = [
+                self.uniprot_to_sequence[uid]
+                for uid in list_of_uniprot_ids
+                if uid in self.uniprot_to_sequence
+            ]
+
+        data = []
+        for sequence in sequences_input:
+            smiles = self.generate_smiles(sequence, num_generated)
+            uniprot_id = next((uid for uid, seq in self.uniprot_to_sequence.items() if seq == sequence), None)
+            data.append({"UniProt_id": uniprot_id, "sequence": sequence, "smiles": smiles})
+
+        return pd.DataFrame(data)
+```
+Below is an example of how to use DrugGen for generating SMILES. Adjust the `num_generated` parameter to specify the number of unique protein SMILES you wish to generate.
+```python
+if __name__ == "__main__":
+    # Initialize the generator
+    generator = SMILESGeneator()
+
+    # Example input (use either list_of_sequences or list_of_uniprot_ids)
+    list_of_sequences = [
+        "MGAASGRRGPGLLLPLPLLLLLPPQPALALDPGLQPGNFSADEAGAQLFAQSYNSSAEQVLFQSVAASWAHDTNITAENARRQEEAALLSQEFAEAWGQKAKELYEPIWQNFTDPQLRRIIGAVRTLGSANLPLAKRQQYNALLSNMSRIYSTAKVCLPNKTATCWSLDPDLTNILASSRSYAMLLFAWEGWHNAAGIPLKPLYEDFTALSNEAYKQDGFTDTGAYWRSWYNSPTFEDDLEHLYQQLEPLYLNLHAFVRRALHRRYGDRYINLRGPIPAHLLGDMWAQSWENIYDMVVPFPDKPNLDVTSTMLQQGWNATHMFRVAEEFFTSLELSPMPPEFWEGSMLEKPADGREVVCHASAWDFYNRKDFRIKQCTRVTMDQLSTVHHEMGHIQYYLQYKDLPVSLRRGANPGFHEAIGDVLALSVSTPEHLHKIGLLDRVTNDTESDINYLLKMALEKIAFLPFGYLVDQWRWGVFSGRTPPSRYNFDWWYLRTKYQGICPPVTRNETHFDAGAKFHVPNVTPYIRYFVSFVLQFQFHEALCKEAGYEGPLHQCDIYRSTKAGAKLRKVLQAGSSRPWQEVLKDMVGLDALDAQPLLKYFQPVTQWLQEQNQQNGEVLGWPEYQWHPPLPDNYPEGIDLVTDEAEASKFVEEYDRTSQVVWNEYAEANWNYNTNITTETSKILLQKNMQIANHTLKYGTQARKFDVNQLQNTTIKRIIKKVQDLERAALPAQELEEYNKILLDMETTYSVATVCHPNGSCLQLEPDLTNVMATSRKYEDLLWAWEGWRDKAGRAILQFYPKYVELINQAARLNGYVDAGDSWRSMYETPSLEQDLERLFQELQPLYLNLHAYVRRALHRHYGAQHINLEGPIPAHLLGNMWAQTWSNIYDLVVPFPSAPSMDTTEAMLKQGWTPRRMFKEADDFFTSLGLLPVPPEFWNKSMLEKPTDGREVVCHASAWDFYNGKDFRIKQCTTVNLEDLVVAHHEMGHIQYFMQYKDLPVALREGANPGFHEAIGDVLALSVSTPKHLHSLNLLSSEGGSDEHDINFLMKMALDKIAFIPFSYLVDQWRWRVFDGSITKENYNQEWWSLRLKYQGLCPPVPRTQGDFDPGAKFHIPSSVPYIRYFVSFIIQFQFHEALCQAAGHTGPLHKCDIYQSKEAGQRLATAMKLGFSRPWPEAMQLITGQPNMSASAMLSYFKPLLDWLRTENELHGEKLGWPQYNWTPNSARSEGPLPDSGRVSFLGLDLDAQQARVGQWLLLFLGIALLVATLGLSQRLFSIRHRSLHRHSHGPQFGSEVELRHS"
+    ]
+    list_of_uniprot_ids = ["P12821", "P37231"]
+
+    # Generate SMILES data for sequences
+    # df = generator.generate_smiles_data(list_of_sequences=list_of_sequences, num_generated=2)
+
+    # Generate SMILES data for UniProt IDs
+    df = generator.generate_smiles_data(list_of_uniprot_ids=list_of_uniprot_ids, num_generated=2)
+
+    # Save the output
+    output_file = "seq_SMILES.txt"
+    df.to_csv(output_file, sep="\t", index=False)
+    print(f"Generated SMILES saved to {output_file}")
+    print(df)
+```
+
+## Training Details
+### Training Data
+[alimotahharynia/approved_drug_target](https://huggingface.co/datasets/alimotahharynia/approved_drug_target)
+- This dataset contains approved SMILES-protein sequences pairs data. It was used to train the model for generating SMILES strings.
+
+### Training Procedure
+- **Training regime:** fp32
+
+#### Supervised Fine-Tuning
+DrugGen was initially trained using supervised fine-tuning on a curated dataset of approved drug targets.
+- **Training: validation sets** (ratio of 8:2) 
+- **sft_config**
+  - `num_train_epochs= 5`
+  - `per_device_train_batch_size= 8`
+  - `per_device_eval_batch_size= 8`
+  - `evaluation_strategy="steps"`
+  - `save_strategy="epoch"`
+  - `eval_steps=50`
+  - `logging_steps=25`
+  - `logging_strategy="steps"`
+  - `do_eval=True`
+  - `do_train=True`
+  - `learning_rate=5e-4`
+  - `adam_epsilon=1e-08`
+  - `warmup_steps=100`
+  - `eval steps=50`
+  - `dataloader_drop_last=True`
+  - `save_safetensors=False`
+  - `max_seq_length=768`
+
+- **AdamW optimizer**
+  - `lr=5e-4`
+  - `eps=1e-08`
+
+- **scheduler**
+  - get_linear_schedule_with_warmup
+#### Proximal Policy Optimization
+
+- **Rollout:** Generates a response based on an input query. Generation parameters include:
+
+  - `do_sample=True`
+  - `top_k=9`
+  - `max_length=1024`
+  - `top_p=0.9`
+  - `bos_token_id=tokenizer.bos_token_id`
+  - `eos_token_id=tokenizer.eos_token_id`
+  - `pad_token_id=tokenizer.pad_token_id`
+  - `num_return_sequences=10`
+
+In each epoch, generation continued until 30 unique small molecules were generated for each target.
+
+- **Evaluation:** A reward function include:
+
+  - Binding affinity predictor: "Protein-Ligand Binding Affinity Prediction Using Pretrained Transformerswas (PLAPT)"
+  - Customized invalid structure assessor: Based on RDKit library
+  - A multiplicative penalty of "0.7" when a generated SMILES matched a molecule present in the approved SMILES dataset.
+
+- **Optimization:**
+
+- **ppo_config**
+  -  `mini_batch_size=8`
+  -  `batch_size=240`
+  -  `learning_rate=1.41e-5`
+  -  `use_score_scaling=True`
+  -  `use_score_norm=True`
+
+Prompts with a tensor size greater than 768 were omitted, resulting in 2053 sequences (98.09% of the initial dataset).
+
+## Citation
+If you use this model in your research or projects, please cite it as: