File size: 3,706 Bytes
349aa7a | 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 | #!/bin/bash
# Example training scripts for Protein-LLM project
# =============================================================================
# 1. Contrastive Pre-training (Stage 1)
# Train QFormer projection layer for protein-text alignment
# =============================================================================
echo "Starting contrastive pre-training..."
python train_contrastive.py \
--text_model_name "Qwen/Qwen3-1.7B" \
--protein_model_name "facebook/esm2_t6_8M_UR50D" \
--qformer_model_name "microsoft/BiomedNLP-BiomedBERT-base-uncased-abstract-fulltext" \
--dataset_name "wanglab/protein_descriptions" \
--output_dir "./contrastive_outputs" \
--num_epochs 10 \
--batch_size 32 \
--learning_rate 1e-4 \
--temperature 0.07 \
--freeze_protein_model \
--freeze_text_model \
--max_length_protein 1024 \
--max_length_text 512 \
--eval_dataset \
--use_wandb \
--wandb_project "protein-llm-contrastive" \
--logging_steps 100 \
--eval_steps 500 \
--save_steps 1000
echo "Contrastive pre-training completed!"
# =============================================================================
# 2. Supervised Fine-tuning (Stage 2)
# Fine-tune the entire model on protein function prediction tasks
# =============================================================================
echo "Starting supervised fine-tuning..."
python train_protein_qwen.py \
--model_type "protein-llm" \
--text_model_name "Qwen/Qwen3-1.7B" \
--protein_model_name "facebook/esm2_t6_8M_UR50D" \
--qformer_model_name "microsoft/BiomedNLP-BiomedBERT-base-uncased-abstract-fulltext" \
--dataset_type "protein_function" \
--protein_function_data_dir_huggingface "wanglab/protein_function" \
--text_model_finetune True \
--protein_model_finetune False \
--num_query_tokens 32 \
--seed 23 \
--batch_size 4 \
--max_epochs 5 \
--learning_rate 5e-5 \
--weight_decay 0.01 \
--gradient_accumulation_steps 8 \
--max_length_protein 1024 \
--max_length_text 1024 \
--lora_rank 32 \
--lora_alpha 64 \
--lora_dropout 0.05 \
--num_gpus 1 \
--strategy "ddp" \
--wandb_project "esm2-qwen3-1.7b-finetune" \
--checkpoint_dir "./checkpoints" \
--log_dir "./logs" \
--cache_dir "/model-weights"
echo "Supervised fine-tuning completed!"
# =============================================================================
# 3. GRPO Training (Stage 3)
# Reinforcement learning with Group Relative Policy Optimization
# =============================================================================
echo "Starting GRPO training..."
python protein_reason.py \
--output_dir "./grpo_outputs" \
--model_name_or_path "Qwen/Qwen3-0.6B" \
--protein_model_name_or_path "facebook/esm2_t6_8M_UR50D" \
--qformer_model_name_or_path "microsoft/BiomedNLP-BiomedBERT-base-uncased-abstract-fulltext" \
--dataset_name "wanglab/protein_function" \
--sft_checkpoint "./checkpoints/best_model" \
--per_device_train_batch_size 4 \
--gradient_accumulation_steps 4 \
--num_train_epochs 3 \
--learning_rate 1e-6 \
--beta 0.04 \
--temperature 0.6 \
--top_p 0.95 \
--top_k 20 \
--max_completion_length 800 \
--num_generations 8 \
--reward_funcs "xmlcount" "soft_format" "strict_format" "correctness" \
--lora_r 32 \
--lora_alpha 64 \
--lora_dropout 0.05 \
--freeze_protein_modules \
--logging_steps 2 \
--eval_strategy "steps" \
--eval_steps 100 \
--save_steps 200 \
--report_to "wandb" \
--log_completions
echo "GRPO training completed!"
echo "All training stages completed successfully!" |