grpo_train.py

print('starting')
import torch
import numpy as np
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, TrainerCallback
#from transformers import T5Tokenizer, T5ForConditionalGeneration
from trl import GRPOConfig, GRPOTrainer
#from trl.core import LengthSampler
from peft import PeftModel, prepare_model_for_kbit_training
from peft import LoraConfig, get_peft_model
import re
import json
import os
import random
import wandb
#from dotenv import load_dotenv
from datasets import Dataset
import time
from rdkit import Chem
from rdkit.Chem import Descriptors
from dockstring import load_target
from tqdm import trange
from rdkit import RDLogger
import multiprocessing
#import torch._dynamo
from pathlib import Path
import deepspeed

import logging
logging.getLogger('transformers').setLevel(logging.ERROR)
#torch._dynamo.config.suppress_errors = True
#torch._dynamo.config.ignore_logger_methods = True

#torch._dynamo.config.ignore_logger_methods.add("info")
#torch._dynamo.config.ignore_logger_methods.add("debug")

import trl.trainer.grpo_trainer as grpo
#accelerator.state.deepspeed_plugin.deepspeed_config = config

os.environ["WANDB_CACHE_DIR"] = "./"

def my_prepare_deepspeed(model, accelerator):
    print(">>> Using custom DeepSpeed config from ds_config.json")
    with open("ds_config.json") as f:
        config = json.load(f)
    config["train_batch_size"] = int(config["train_batch_size"])
    config["train_micro_batch_size_per_gpu"] = int(config["train_micro_batch_size_per_gpu"])
    config["gradient_accumulation_steps"] = int(config["gradient_accumulation_steps"])

    accelerator.state.deepspeed_plugin.deepspeed_config = config
    model_engine, optimizer, _, _ = deepspeed.initialize(model=model, config=config)
#    model_engine, _, _, _ = deepspeed.initialize(model=model, config=config)
    return model_engine

# --- PATCH BEFORE TRAINER IS USED ---
grpo.prepare_deepspeed = my_prepare_deepspeed

NUM_EPOCHS = 100

#model_name = "/rs1/researchers/w/wjpfaend/nisargj/moljet/grpo/models--Qwen--Qwen2-0.5B-Instruct/snapshots/c540970f9e29518b1d8f06ab8b24cba66ad77b6d"
#model_name = "/rs1/researchers/w/wjpfaend/nisargj/moljet/grpo/models--google--txgemma-2b-predict/snapshots/6b0eb81135ad9e73c9c18409c03e315a5a23a18f"
model_name = "/rs1/researchers/w/wjpfaend/nisargj/moljet/grpo/models--google--txgemma-9b-chat/snapshots/bc45665ceb345eea6acfc44d82706e8d8b689209"

os.environ["WANDB_MODE"] = "offline"
#load_dotenv()

def construct_prompt(pro):
         """
         Construct prompt for a single protein target
         """
         prompt = f"""Your task is to design a small, drug-like molecule. The molecule binds to the protein target: {pro}.

         Generate a SMILES (Simplified Molecular Input Line Entry System) for the drug-like molecule.
         The SMILES must be syntactically valid (as accepted by RDKit or OpenBabel).

        Return EXACTLY two items in this order:

        1) The final SMILES string inside <answer>...</answer> tags (and NOTHING else inside those tags).
        2) A concise explanation inside <reasoning>...</reasoning> tags that addresses BOTH points below, specific to the protein target given above:
            1. How the designed drug-like molecule affects (or does not affect) the protein target.
            2. How the structure of the molecule affects binding to the protein target.
         
         IMPORTANT:
         - Output ONLY one valid SMILES string.
         - Put the SMILES inside <answer></answer> tags.
         - Do NOT include anything else inside the tags.
         - Do NOT include multiple SMILES, salts, counter-ions, or explanations.
         - If unsure, return the simplest valid drug-like SMILES you can.

         Examples:
         <answer>C(N(CC)C(=O)OC(C)(C)C)C</answer>
         <answer>C1=CC=CC=C1</answer>
         <reasoning>1) Small polar molecule likely to interact with polar pocket of TARGETPROT. 2) The hydroxyl group can form H-bonds with residue X, improving affinity.</reasoning>
         """
         return prompt

#def save_completions(prompts, completions, **kwargs):
#    with open("completions_checkpoints.json", "w", encoding="utf-8") as f:
#        for prompt, completion in zip(prompts, completions):
#            record = {
#                "prompt": prompt,
#                "completion": completion
#            }
#            f.write(json.dumps(record, ensure_ascii=False) + "\n")

log_file = Path("completions_log.jsonl")

def get_docking_score(ligand, target_name):
    """
    Get docking score for a single ligand-target pair.
    """
    target = load_target(target_name)
    try:
        score, _ = target.dock(ligand, num_cpus=8)
        score = score*(-1)
        if not isinstance(score, (int, float)):
            print(f"Warning: non-numeric score for {ligand}, {target_name}: {score}")
            score = -10.0
#        score, _ = target.dock(ligand, num_cpus=8)
    except Exception as e:
        print(f"Docking error on {ligand}, {target_name}: {e}")
        score = -10.0
    return score

def dockscore_reward_function(prompts, completions, **kwargs):
    rewards = []

    print("length of outputs/prompts:", len(prompts))
#    save_completions(prompts, completions, save_dir='./completions')

    for i in trange(len(prompts), desc="Sequential Docking"):
        prompt = prompts[i]
        completion = completions[i]
        
        ligand_match = re.findall(r"<answer>(.*?)</answer>", completion, re.DOTALL)
        target_match = re.search(r"protein target:\s*([A-Za-z0-9_-]+)", prompt)
#        target_match = re.search(r"protein target:\s*(\S+)", prompt)
        
        score = None
        ligand = None
        
        try:
#      
            if not ligand_match:
                rewards.append(-100.0)
                continue
#            else:
            ligand = ligand_match[-1].strip()
            target_name = target_match.group(1).strip()
            
            mol = Chem.MolFromSmiles(ligand)
            if mol is None:
                rewards.append(-100.0)
                continue

            score = get_docking_score(ligand, target_name)
            if score is not None:
                rewards.append(score)
            else:
                rewards.append(-10.0)

        except ValueError as e:
            print(f"SMILES parsing error: {e}")
            rewards.append(-100.0)

        except Exception as e:
            print(f"Error calculating stability score: {e}")
            rewards.append(-10.0)
       
        log_entry = {
            "prompt": prompt,
            "completion": completion,
        }
        if score is not None: # and score >= 0.0:  # only log valid score
            log_entry["score"] = score

        if ligand:
            log_entry["SMILES"] = ligand

#        if mol_weight:
#            log_entry["MolWt"] = mol_weight

        with log_file.open("a") as f:
            json.dump(log_entry, f)
            f.write("\n") 

    return rewards 

#if __name__ == '__main__':
def main():
    
#    print(torch.cuda.device_count())
#    print(torch.distributed.is_initialized())
#    return

    print("\n Loading and processing dataset...")
    data_load_start = time.time()

    #load the protein targets:
    with open("dockstring_proteins.txt", "r") as t:
        target_list = [l.rstrip('\n') for l in t.readlines()]

    #create prompts for 58 target proteins *100

    all_prompts = []
    for target in target_list:
        all_prompts.append(construct_prompt(target))

    all_prompts = all_prompts*10 # multiplying it to increase the dataset number, change it to 100 after first trial

    data_list = [{"prompt": prompt} for prompt in all_prompts]

    #train_dataset = Dataset.from_list(data_list)
    train_dataset = Dataset.from_list(data_list)
    print(f"Dataset size: {len(train_dataset)}")
    print(f"Data loading and processing completed in {time.time() - data_load_start:.2f} seconds")

    print("\nInitializing wandb...")
    wandb_start = time.time()
    # Initialize wandb
    #wandb.login(key=os.getenv('WANDB_API_KEY'))
    wandb.init(
        project="dockstring_32batch_size_16gen_test",
        name="grpo_4gpus_epoch10_16gen",
        config={
            "model_name": "google/txgemma-9b-chat",
            "num_epochs": NUM_EPOCHS,
#            "batch_size": 2,
#            "learning_rate": 1.41e-5,
#            "num_generations": 4,
        },
#        settings=wandb.Settings(_data_dir="/rs1/researchers/w/wjpfaend/nisargj/moljet/grpo/train/test/lora/gemma_3.1.try/")
    )

    quantization_config = BitsAndBytesConfig(
        load_in_8bit=True,
    )

    print("Loading base model...")
    model_load_start = time.time()
    base_model = AutoModelForCausalLM.from_pretrained(
                model_name,
    #            quantization_config=quantization_config,
#                device_map="auto",
                torch_dtype=torch.float16,
                attn_implementation='eager',
                use_cache=False
    )
    base_model.config.use_cache=False
    print(f"Base model loaded in {time.time() - model_load_start:.2f} seconds")


    print("\nLoading tokenizer...")
    tokenizer_start = time.time()
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    print(f"Tokenizer loaded in {time.time() - tokenizer_start:.2f} seconds")

    #skipping k-bit training and LoRA adapters for now

    base_model = prepare_model_for_kbit_training(base_model)


    lora_config = LoraConfig(
        task_type="CAUSAL_LM",
        r=16,
        lora_alpha=32,
        target_modules=["q_proj","k_proj", "v_proj", "o_proj"],
    )

    model = get_peft_model(base_model, lora_config)
    print(model.print_trainable_parameters())
    #print(f"Loading LoRA adapter from {peft_lora_path}")
    #model = PeftModel.from_pretrained(base_model, peft_lora_path)

    # Enable gradient computation for LoRA parameters
    for name, param in model.named_parameters():
        if "lora" in name.lower():
            param.requires_grad = True
        else:
            param.requires_grad = False

    # Verify some parameters require gradients
    trainable_params = [p for p in model.parameters() if p.requires_grad]
    if not trainable_params:
        raise ValueError("No parameters have requires_grad=True. Training will not work!")

    print(f"Number of trainable parameters: {sum(p.numel() for p in trainable_params)}")

    model.train()  # Ensure model is in training mode
    #base_model.train()
    #Create training arguments
    training_args = GRPOConfig(
        output_dir="./grpo_output",
        run_name="grpo_training_run",  # Add distinct run name
        num_train_epochs=NUM_EPOCHS,
        per_device_train_batch_size=2,
        fp16=True,
        gradient_accumulation_steps=4,
        learning_rate=1e-4, #1.41e-5,
        lr_scheduler_type='cosine',
        warmup_ratio=0.05,
        logging_steps=1,
        num_generations=16,
        max_prompt_length=512,
        max_completion_length=512,
        temperature=0.7,
        beta=0.04,
        log_completions=True,
        remove_unused_columns=False,
        save_strategy='steps',
        save_steps=500,
#        gradient_checkpointing=True,
#        gradient_checkpointing_kwargs={"use_reentrant": False}  # Explicitly set use_reentrant
    )

    class WandBLoggingCallback(TrainerCallback):
        def on_log(self, args, state, control, logs=None, **kwargs):
            if logs:
                # Log all metrics from the trainer
                wandb.log(logs, step=state.global_step)

        def on_evaluate(self, args, state, control, metrics=None, **kwargs):
            if metrics:
                # Log evaluation metrics
                wandb.log({"eval/" + k: v for k, v in metrics.items()}, step=state.global_step)




    # Initialize GRPO trainer
    trainer = GRPOTrainer(
        model=model, #base_model,
        args=training_args,
        train_dataset=train_dataset,
        reward_funcs=dockscore_reward_function,
        processing_class=tokenizer,
#        optimizers=(DummyOptim, DummyScheduler),
        callbacks=[WandBLoggingCallback()],
    )
    #trainer.train(resume_from_checkpoint=True)
    trainer.train()
    # Save the final model
    trainer.save_model("./grpo_output/final_model")
    
    print('saving completions')
    
    wandb.finish()
if __name__ == '__main__':
    main()