Add files using upload-large-folder tool

BIO/sft/qwen-metal_ion_binding-08022141/v0-20250802-215822/checkpoint-250/README.md

@@ -0,0 +1,202 @@
+---
+base_model: /oss/wangyujia/BIO/construction_finetuning/alpaca/v1-20250609-141541/checkpoint-50-merged
+library_name: peft
+---
+
+# Model Card for Model ID
+
+<!-- Provide a quick summary of what the model is/does. -->
+
+
+
+## Model Details
+
+### Model Description
+
+<!-- Provide a longer summary of what this model is. -->
+
+
+
+- **Developed by:** [More Information Needed]
+- **Funded by [optional]:** [More Information Needed]
+- **Shared by [optional]:** [More Information Needed]
+- **Model type:** [More Information Needed]
+- **Language(s) (NLP):** [More Information Needed]
+- **License:** [More Information Needed]
+- **Finetuned from model [optional]:** [More Information Needed]
+
+### Model Sources [optional]
+
+<!-- Provide the basic links for the model. -->
+
+- **Repository:** [More Information Needed]
+- **Paper [optional]:** [More Information Needed]
+- **Demo [optional]:** [More Information Needed]
+
+## Uses
+
+<!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
+
+### Direct Use
+
+<!-- This section is for the model use without fine-tuning or plugging into a larger ecosystem/app. -->
+
+[More Information Needed]
+
+### Downstream Use [optional]
+
+<!-- This section is for the model use when fine-tuned for a task, or when plugged into a larger ecosystem/app -->
+
+[More Information Needed]
+
+### Out-of-Scope Use
+
+<!-- This section addresses misuse, malicious use, and uses that the model will not work well for. -->
+
+[More Information Needed]
+
+## Bias, Risks, and Limitations
+
+<!-- This section is meant to convey both technical and sociotechnical limitations. -->
+
+[More Information Needed]
+
+### Recommendations
+
+<!-- This section is meant to convey recommendations with respect to the bias, risk, and technical limitations. -->
+
+Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. More information needed for further recommendations.
+
+## How to Get Started with the Model
+
+Use the code below to get started with the model.
+
+[More Information Needed]
+
+## Training Details
+
+### Training Data
+
+<!-- This should link to a Dataset Card, perhaps with a short stub of information on what the training data is all about as well as documentation related to data pre-processing or additional filtering. -->
+
+[More Information Needed]
+
+### Training Procedure
+
+<!-- This relates heavily to the Technical Specifications. Content here should link to that section when it is relevant to the training procedure. -->
+
+#### Preprocessing [optional]
+
+[More Information Needed]
+
+
+#### Training Hyperparameters
+
+- **Training regime:** [More Information Needed] <!--fp32, fp16 mixed precision, bf16 mixed precision, bf16 non-mixed precision, fp16 non-mixed precision, fp8 mixed precision -->
+
+#### Speeds, Sizes, Times [optional]
+
+<!-- This section provides information about throughput, start/end time, checkpoint size if relevant, etc. -->
+
+[More Information Needed]
+
+## Evaluation
+
+<!-- This section describes the evaluation protocols and provides the results. -->
+
+### Testing Data, Factors & Metrics
+
+#### Testing Data
+
+<!-- This should link to a Dataset Card if possible. -->
+
+[More Information Needed]
+
+#### Factors
+
+<!-- These are the things the evaluation is disaggregating by, e.g., subpopulations or domains. -->
+
+[More Information Needed]
+
+#### Metrics
+
+<!-- These are the evaluation metrics being used, ideally with a description of why. -->
+
+[More Information Needed]
+
+### Results
+
+[More Information Needed]
+
+#### Summary
+
+
+
+## Model Examination [optional]
+
+<!-- Relevant interpretability work for the model goes here -->
+
+[More Information Needed]
+
+## Environmental Impact
+
+<!-- Total emissions (in grams of CO2eq) and additional considerations, such as electricity usage, go here. Edit the suggested text below accordingly -->
+
+Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700).
+
+- **Hardware Type:** [More Information Needed]
+- **Hours used:** [More Information Needed]
+- **Cloud Provider:** [More Information Needed]
+- **Compute Region:** [More Information Needed]
+- **Carbon Emitted:** [More Information Needed]
+
+## Technical Specifications [optional]
+
+### Model Architecture and Objective
+
+[More Information Needed]
+
+### Compute Infrastructure
+
+[More Information Needed]
+
+#### Hardware
+
+[More Information Needed]
+
+#### Software
+
+[More Information Needed]
+
+## Citation [optional]
+
+<!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->
+
+**BibTeX:**
+
+[More Information Needed]
+
+**APA:**
+
+[More Information Needed]
+
+## Glossary [optional]
+
+<!-- If relevant, include terms and calculations in this section that can help readers understand the model or model card. -->
+
+[More Information Needed]
+
+## More Information [optional]
+
+[More Information Needed]
+
+## Model Card Authors [optional]
+
+[More Information Needed]
+
+## Model Card Contact
+
+[More Information Needed]
+### Framework versions
+
+- PEFT 0.15.2

BIO/sft/qwen-metal_ion_binding-08022141/v0-20250802-215822/checkpoint-250/adapter_config.json

@@ -0,0 +1,39 @@
+{
+  "alpha_pattern": {},
+  "auto_mapping": null,
+  "base_model_name_or_path": "/oss/wangyujia/BIO/construction_finetuning/alpaca/v1-20250609-141541/checkpoint-50-merged",
+  "bias": "none",
+  "corda_config": null,
+  "eva_config": null,
+  "exclude_modules": null,
+  "fan_in_fan_out": false,
+  "inference_mode": true,
+  "init_lora_weights": true,
+  "layer_replication": null,
+  "layers_pattern": null,
+  "layers_to_transform": null,
+  "loftq_config": {},
+  "lora_alpha": 32,
+  "lora_bias": false,
+  "lora_dropout": 0.05,
+  "megatron_config": null,
+  "megatron_core": "megatron.core",
+  "modules_to_save": [],
+  "peft_type": "LORA",
+  "r": 8,
+  "rank_pattern": {},
+  "revision": null,
+  "target_modules": [
+    "up_proj",
+    "o_proj",
+    "gate_proj",
+    "down_proj",
+    "v_proj",
+    "k_proj",
+    "q_proj"
+  ],
+  "task_type": "CAUSAL_LM",
+  "trainable_token_indices": null,
+  "use_dora": false,
+  "use_rslora": false
+}

BIO/sft/qwen-metal_ion_binding-08022141/v0-20250802-215822/checkpoint-250/additional_config.json

@@ -0,0 +1 @@
+{"lora_dtype": null, "lorap_lr_ratio": null, "lorap_emb_lr": 1e-06}

BIO/sft/qwen-metal_ion_binding-08022141/v0-20250802-215822/checkpoint-250/args.json

@@ -0,0 +1,364 @@
+{
+  "model": "/oss/wangyujia/BIO/construction_finetuning/alpaca/v1-20250609-141541/checkpoint-50-merged",
+  "model_type": "qwen2_5",
+  "model_revision": null,
+  "task_type": "causal_lm",
+  "torch_dtype": "bfloat16",
+  "attn_impl": null,
+  "num_labels": null,
+  "problem_type": null,
+  "rope_scaling": null,
+  "device_map": null,
+  "max_memory": {},
+  "local_repo_path": null,
+  "template": "qwen2_5",
+  "system": null,
+  "max_length": 8192,
+  "truncation_strategy": "delete",
+  "max_pixels": null,
+  "tools_prompt": "react_en",
+  "norm_bbox": null,
+  "response_prefix": null,
+  "padding_side": "right",
+  "loss_scale": "default",
+  "sequence_parallel_size": 1,
+  "use_chat_template": true,
+  "template_backend": "swift",
+  "dataset": [
+    "/oss/wangyujia/ProtT3/ProtT3/data/sft/dataset/metal_ion_binding/train.jsonl"
+  ],
+  "val_dataset": [],
+  "split_dataset_ratio": 0.01,
+  "data_seed": 42,
+  "dataset_num_proc": 128,
+  "dataset_shuffle": true,
+  "val_dataset_shuffle": false,
+  "streaming": false,
+  "interleave_prob": null,
+  "stopping_strategy": "first_exhausted",
+  "shuffle_buffer_size": 1000,
+  "enable_cache": false,
+  "download_mode": "reuse_dataset_if_exists",
+  "columns": {},
+  "strict": false,
+  "remove_unused_columns": true,
+  "model_name": [
+    "qwen_bio_sft_deeplocbinary-08022035"
+  ],
+  "model_author": [
+    "swift"
+  ],
+  "custom_dataset_info": [],
+  "quant_method": null,
+  "quant_bits": null,
+  "hqq_axis": null,
+  "bnb_4bit_compute_dtype": "bfloat16",
+  "bnb_4bit_quant_type": "nf4",
+  "bnb_4bit_use_double_quant": true,
+  "bnb_4bit_quant_storage": null,
+  "max_new_tokens": 64,
+  "temperature": 0.0,
+  "top_k": null,
+  "top_p": null,
+  "repetition_penalty": null,
+  "num_beams": 1,
+  "stream": false,
+  "stop_words": [],
+  "logprobs": false,
+  "top_logprobs": null,
+  "ckpt_dir": "/oss/wangyujia/BIO/construction_finetuning/alpaca/v1-20250609-141541/checkpoint-50-merged",
+  "load_dataset_config": null,
+  "lora_modules": [],
+  "tuner_backend": "peft",
+  "train_type": "lora",
+  "adapters": [],
+  "external_plugins": [],
+  "seed": 42,
+  "model_kwargs": {},
+  "load_args": false,
+  "load_data_args": false,
+  "use_hf": false,
+  "hub_token": null,
+  "custom_register_path": [],
+  "ignore_args_error": false,
+  "use_swift_lora": false,
+  "output_dir": "/nas/shared/kilab/wangyujia/BIO/sft/qwen-metal_ion_binding-08022141/v0-20250802-215822",
+  "overwrite_output_dir": false,
+  "do_train": false,
+  "do_eval": false,
+  "do_predict": false,
+  "eval_strategy": "steps",
+  "prediction_loss_only": false,
+  "per_device_train_batch_size": 2,
+  "per_device_eval_batch_size": 2,
+  "per_gpu_train_batch_size": null,
+  "per_gpu_eval_batch_size": null,
+  "gradient_accumulation_steps": 4,
+  "eval_accumulation_steps": null,
+  "eval_delay": 0,
+  "torch_empty_cache_steps": null,
+  "learning_rate": 1e-05,
+  "weight_decay": 0.1,
+  "adam_beta1": 0.9,
+  "adam_beta2": 0.95,
+  "adam_epsilon": 1e-08,
+  "max_grad_norm": 1.0,
+  "num_train_epochs": 3.0,
+  "max_steps": -1,
+  "lr_scheduler_type": "cosine",
+  "lr_scheduler_kwargs": null,
+  "warmup_ratio": 0.05,
+  "warmup_steps": 0,
+  "log_level": "passive",
+  "log_level_replica": "warning",
+  "log_on_each_node": true,
+  "logging_dir": "/nas/shared/kilab/wangyujia/BIO/sft/qwen-metal_ion_binding-08022141/v0-20250802-215822/runs",
+  "logging_strategy": "steps",
+  "logging_first_step": true,
+  "logging_steps": 1,
+  "logging_nan_inf_filter": true,
+  "save_strategy": "steps",
+  "save_steps": 5.0,
+  "save_total_limit": 5,
+  "save_safetensors": true,
+  "save_on_each_node": false,
+  "save_only_model": true,
+  "restore_callback_states_from_checkpoint": false,
+  "no_cuda": false,
+  "use_cpu": false,
+  "use_mps_device": false,
+  "jit_mode_eval": false,
+  "use_ipex": false,
+  "bf16": true,
+  "fp16": false,
+  "fp16_opt_level": "O1",
+  "half_precision_backend": "auto",
+  "bf16_full_eval": false,
+  "fp16_full_eval": false,
+  "tf32": null,
+  "local_rank": 0,
+  "ddp_backend": null,
+  "tpu_num_cores": null,
+  "tpu_metrics_debug": false,
+  "debug": null,
+  "dataloader_drop_last": false,
+  "eval_steps": 5.0,
+  "dataloader_num_workers": 1,
+  "dataloader_prefetch_factor": null,
+  "past_index": -1,
+  "run_name": "construct",
+  "disable_tqdm": null,
+  "label_names": null,
+  "load_best_model_at_end": false,
+  "metric_for_best_model": "loss",
+  "greater_is_better": false,
+  "ignore_data_skip": false,
+  "fsdp": "",
+  "fsdp_min_num_params": 0,
+  "fsdp_config": null,
+  "tp_size": 0,
+  "fsdp_transformer_layer_cls_to_wrap": null,
+  "accelerator_config": {
+    "dispatch_batches": false
+  },
+  "deepspeed": {
+    "fp16": {
+      "enabled": "auto",
+      "loss_scale": 0,
+      "loss_scale_window": 1000,
+      "initial_scale_power": 16,
+      "hysteresis": 2,
+      "min_loss_scale": 1
+    },
+    "bf16": {
+      "enabled": "auto"
+    },
+    "zero_optimization": {
+      "stage": 3,
+      "offload_optimizer": {
+        "device": "none",
+        "pin_memory": true
+      },
+      "offload_param": {
+        "device": "none",
+        "pin_memory": true
+      },
+      "overlap_comm": false,
+      "contiguous_gradients": true,
+      "sub_group_size": 1000000000.0,
+      "reduce_bucket_size": "auto",
+      "zero_quantized_weights": false,
+      "zero_quantized_gradients": false,
+      "stage3_prefetch_bucket_size": "auto",
+      "stage3_param_persistence_threshold": "auto",
+      "stage3_max_live_parameters": 1000000000.0,
+      "stage3_max_reuse_distance": 1000000000.0,
+      "stage3_gather_16bit_weights_on_model_save": true
+    },
+    "gradient_accumulation_steps": "auto",
+    "gradient_clipping": "auto",
+    "steps_per_print": 2000,
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "wall_clock_breakdown": false
+  },
+  "label_smoothing_factor": 0.0,
+  "optim": "adamw_torch",
+  "optim_args": null,
+  "adafactor": false,
+  "group_by_length": false,
+  "length_column_name": "length",
+  "report_to": [
+    "tensorboard"
+  ],
+  "ddp_find_unused_parameters": null,
+  "ddp_bucket_cap_mb": null,
+  "ddp_broadcast_buffers": null,
+  "dataloader_pin_memory": true,
+  "dataloader_persistent_workers": false,
+  "skip_memory_metrics": true,
+  "use_legacy_prediction_loop": false,
+  "push_to_hub": false,
+  "resume_from_checkpoint": null,
+  "hub_model_id": null,
+  "hub_strategy": "every_save",
+  "hub_private_repo": null,
+  "hub_always_push": false,
+  "gradient_checkpointing": true,
+  "gradient_checkpointing_kwargs": null,
+  "include_inputs_for_metrics": false,
+  "include_for_metrics": [],
+  "eval_do_concat_batches": true,
+  "fp16_backend": "auto",
+  "push_to_hub_model_id": null,
+  "push_to_hub_organization": null,
+  "push_to_hub_token": null,
+  "mp_parameters": "",
+  "auto_find_batch_size": false,
+  "full_determinism": false,
+  "torchdynamo": null,
+  "ray_scope": "last",
+  "ddp_timeout": 1800,
+  "torch_compile": false,
+  "torch_compile_backend": null,
+  "torch_compile_mode": null,
+  "include_tokens_per_second": false,
+  "include_num_input_tokens_seen": false,
+  "neftune_noise_alpha": null,
+  "optim_target_modules": null,
+  "batch_eval_metrics": false,
+  "eval_on_start": false,
+  "use_liger_kernel": false,
+  "eval_use_gather_object": false,
+  "average_tokens_across_devices": false,
+  "sortish_sampler": false,
+  "predict_with_generate": false,
+  "generation_max_length": null,
+  "generation_num_beams": null,
+  "generation_config": null,
+  "check_model": true,
+  "acc_strategy": "token",
+  "train_dataloader_shuffle": true,
+  "metric_warmup_step": 0,
+  "fsdp_num": 1,
+  "acc_steps": 1,
+  "eval_use_evalscope": false,
+  "eval_datasets": [],
+  "eval_limit": null,
+  "eval_datasets_args": null,
+  "eval_generation_config": null,
+  "freeze_parameters": [],
+  "freeze_parameters_ratio": 0.0,
+  "trainable_parameters": [],
+  "freeze_llm": false,
+  "freeze_vit": true,
+  "freeze_aligner": true,
+  "target_modules": [
+    "all-linear"
+  ],
+  "target_regex": null,
+  "modules_to_save": [],
+  "lora_rank": 8,
+  "lora_alpha": 32,
+  "lora_dropout": 0.05,
+  "lora_bias": "none",
+  "lora_dtype": null,
+  "lorap_lr_ratio": null,
+  "use_rslora": false,
+  "use_dora": false,
+  "lora_ga_batch_size": 2,
+  "lora_ga_iters": 2,
+  "lora_ga_max_length": 1024,
+  "lora_ga_direction": "ArB2r",
+  "lora_ga_scale": "stable",
+  "lora_ga_stable_gamma": 16,
+  "init_weights": true,
+  "fourier_n_frequency": 2000,
+  "fourier_scaling": 300.0,
+  "boft_block_size": 4,
+  "boft_block_num": 0,
+  "boft_n_butterfly_factor": 1,
+  "boft_dropout": 0.0,
+  "vera_rank": 256,
+  "vera_projection_prng_key": 0,
+  "vera_dropout": 0.0,
+  "vera_d_initial": 0.1,
+  "adapter_act": "gelu",
+  "adapter_length": 128,
+  "use_galore": false,
+  "galore_target_modules": null,
+  "galore_rank": 128,
+  "galore_update_proj_gap": 50,
+  "galore_scale": 1.0,
+  "galore_proj_type": "std",
+  "galore_optim_per_parameter": false,
+  "galore_with_embedding": false,
+  "galore_quantization": false,
+  "galore_proj_quant": false,
+  "galore_proj_bits": 4,
+  "galore_proj_group_size": 256,
+  "galore_cos_threshold": 0.4,
+  "galore_gamma_proj": 2,
+  "galore_queue_size": 5,
+  "adalora_target_r": 8,
+  "adalora_init_r": 12,
+  "adalora_tinit": 0,
+  "adalora_tfinal": 0,
+  "adalora_deltaT": 1,
+  "adalora_beta1": 0.85,
+  "adalora_beta2": 0.85,
+  "adalora_orth_reg_weight": 0.5,
+  "llamapro_num_new_blocks": 4,
+  "llamapro_num_groups": null,
+  "lisa_activated_layers": 0,
+  "lisa_step_interval": 20,
+  "reft_layer_key": null,
+  "reft_layers": null,
+  "reft_rank": 4,
+  "reft_intervention_type": "LoreftIntervention",
+  "reft_args": null,
+  "swanlab_token": null,
+  "swanlab_project": null,
+  "swanlab_workspace": null,
+  "swanlab_exp_name": null,
+  "swanlab_mode": "cloud",
+  "add_version": true,
+  "resume_only_model": false,
+  "create_checkpoint_symlink": false,
+  "packing": false,
+  "lazy_tokenize": false,
+  "loss_type": null,
+  "optimizer": null,
+  "metric": null,
+  "zero_hpz_partition_size": null,
+  "rank": 0,
+  "global_world_size": 8,
+  "local_world_size": 8,
+  "model_suffix": "checkpoint-50-merged",
+  "model_info": "ModelInfo(model_type='qwen2_5', model_dir='/oss/wangyujia/BIO/construction_finetuning/alpaca/v1-20250609-141541/checkpoint-50-merged', torch_dtype=torch.bfloat16, max_model_len=32768, quant_method=None, quant_bits=None, rope_scaling=None, config=None, task_type='causal_lm', num_labels=None)",
+  "model_meta": "ModelMeta(model_type='qwen2_5', model_groups=[ModelGroup(models=[Model(ms_model_id='Qwen/Qwen2.5-0.5B-Instruct', hf_model_id='Qwen/Qwen2.5-0.5B-Instruct', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-1.5B-Instruct', hf_model_id='Qwen/Qwen2.5-1.5B-Instruct', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-3B-Instruct', hf_model_id='Qwen/Qwen2.5-3B-Instruct', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-7B-Instruct', hf_model_id='Qwen/Qwen2.5-7B-Instruct', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-14B-Instruct', hf_model_id='Qwen/Qwen2.5-14B-Instruct', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-32B-Instruct', hf_model_id='Qwen/Qwen2.5-32B-Instruct', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-72B-Instruct', hf_model_id='Qwen/Qwen2.5-72B-Instruct', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-0.5B', hf_model_id='Qwen/Qwen2.5-0.5B', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-1.5B', hf_model_id='Qwen/Qwen2.5-1.5B', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-3B', hf_model_id='Qwen/Qwen2.5-3B', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-7B', hf_model_id='Qwen/Qwen2.5-7B', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-14B', hf_model_id='Qwen/Qwen2.5-14B', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-32B', hf_model_id='Qwen/Qwen2.5-32B', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-72B', hf_model_id='Qwen/Qwen2.5-72B', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-0.5B-Instruct-GPTQ-Int4', hf_model_id='Qwen/Qwen2.5-0.5B-Instruct-GPTQ-Int4', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-1.5B-Instruct-GPTQ-Int4', hf_model_id='Qwen/Qwen2.5-1.5B-Instruct-GPTQ-Int4', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-3B-Instruct-GPTQ-Int4', hf_model_id='Qwen/Qwen2.5-3B-Instruct-GPTQ-Int4', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4', hf_model_id='Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-14B-Instruct-GPTQ-Int4', hf_model_id='Qwen/Qwen2.5-14B-Instruct-GPTQ-Int4', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-32B-Instruct-GPTQ-Int4', hf_model_id='Qwen/Qwen2.5-32B-Instruct-GPTQ-Int4', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-72B-Instruct-GPTQ-Int4', hf_model_id='Qwen/Qwen2.5-72B-Instruct-GPTQ-Int4', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-0.5B-Instruct-GPTQ-Int8', hf_model_id='Qwen/Qwen2.5-0.5B-Instruct-GPTQ-Int8', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-1.5B-Instruct-GPTQ-Int8', hf_model_id='Qwen/Qwen2.5-1.5B-Instruct-GPTQ-Int8', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-3B-Instruct-GPTQ-Int8', hf_model_id='Qwen/Qwen2.5-3B-Instruct-GPTQ-Int8', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-7B-Instruct-GPTQ-Int8', hf_model_id='Qwen/Qwen2.5-7B-Instruct-GPTQ-Int8', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-14B-Instruct-GPTQ-Int8', hf_model_id='Qwen/Qwen2.5-14B-Instruct-GPTQ-Int8', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-32B-Instruct-GPTQ-Int8', hf_model_id='Qwen/Qwen2.5-32B-Instruct-GPTQ-Int8', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-72B-Instruct-GPTQ-Int8', hf_model_id='Qwen/Qwen2.5-72B-Instruct-GPTQ-Int8', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-0.5B-Instruct-AWQ', hf_model_id='Qwen/Qwen2.5-0.5B-Instruct-AWQ', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-1.5B-Instruct-AWQ', hf_model_id='Qwen/Qwen2.5-1.5B-Instruct-AWQ', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-3B-Instruct-AWQ', hf_model_id='Qwen/Qwen2.5-3B-Instruct-AWQ', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-7B-Instruct-AWQ', hf_model_id='Qwen/Qwen2.5-7B-Instruct-AWQ', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-14B-Instruct-AWQ', hf_model_id='Qwen/Qwen2.5-14B-Instruct-AWQ', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-32B-Instruct-AWQ', hf_model_id='Qwen/Qwen2.5-32B-Instruct-AWQ', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-72B-Instruct-AWQ', hf_model_id='Qwen/Qwen2.5-72B-Instruct-AWQ', model_path=None, ms_revision=None, hf_revision=None)], ignore_patterns=None, requires=None, tags=[]), ModelGroup(models=[Model(ms_model_id='Qwen/Qwen2.5-Coder-0.5B-Instruct', hf_model_id='Qwen/Qwen2.5-Coder-0.5B-Instruct', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-Coder-1.5B-Instruct', hf_model_id='Qwen/Qwen2.5-Coder-1.5B-Instruct', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-Coder-3B-Instruct', hf_model_id='Qwen/Qwen2.5-Coder-3B-Instruct', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-Coder-7B-Instruct', hf_model_id='Qwen/Qwen2.5-Coder-7B-Instruct', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-Coder-14B-Instruct', hf_model_id='Qwen/Qwen2.5-Coder-14B-Instruct', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-Coder-32B-Instruct', hf_model_id='Qwen/Qwen2.5-Coder-32B-Instruct', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-Coder-0.5B', hf_model_id='Qwen/Qwen2.5-Coder-0.5B', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-Coder-1.5B', hf_model_id='Qwen/Qwen2.5-Coder-1.5B', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-Coder-3B', hf_model_id='Qwen/Qwen2.5-Coder-3B', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-Coder-7B', hf_model_id='Qwen/Qwen2.5-Coder-7B', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-Coder-14B', hf_model_id='Qwen/Qwen2.5-Coder-14B', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-Coder-32B', hf_model_id='Qwen/Qwen2.5-Coder-32B', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-Coder-0.5B-Instruct-AWQ', hf_model_id='Qwen/Qwen2.5-Coder-0.5B-Instruct-AWQ', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-Coder-1.5B-Instruct-AWQ', hf_model_id='Qwen/Qwen2.5-Coder-1.5B-Instruct-AWQ', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-Coder-3B-Instruct-AWQ', hf_model_id='Qwen/Qwen2.5-Coder-3B-Instruct-AWQ', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-Coder-7B-Instruct-AWQ', hf_model_id='Qwen/Qwen2.5-Coder-7B-Instruct-AWQ', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-Coder-14B-Instruct-AWQ', hf_model_id='Qwen/Qwen2.5-Coder-14B-Instruct-AWQ', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-Coder-32B-Instruct-AWQ', hf_model_id='Qwen/Qwen2.5-Coder-32B-Instruct-AWQ', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-Coder-0.5B-Instruct-GPTQ-Int4', hf_model_id='Qwen/Qwen2.5-Coder-0.5B-Instruct-GPTQ-Int4', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-Coder-0.5B-Instruct-GPTQ-Int8', hf_model_id='Qwen/Qwen2.5-Coder-0.5B-Instruct-GPTQ-Int8', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-Coder-1.5B-Instruct-GPTQ-Int4', hf_model_id='Qwen/Qwen2.5-Coder-1.5B-Instruct-GPTQ-Int4', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-Coder-1.5B-Instruct-GPTQ-Int8', hf_model_id='Qwen/Qwen2.5-Coder-1.5B-Instruct-GPTQ-Int8', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-Coder-3B-Instruct-GPTQ-Int4', hf_model_id='Qwen/Qwen2.5-Coder-3B-Instruct-GPTQ-Int4', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-Coder-3B-Instruct-GPTQ-Int8', hf_model_id='Qwen/Qwen2.5-Coder-3B-Instruct-GPTQ-Int8', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-Coder-7B-Instruct-GPTQ-Int4', hf_model_id='Qwen/Qwen2.5-Coder-7B-Instruct-GPTQ-Int4', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-Coder-7B-Instruct-GPTQ-Int8', hf_model_id='Qwen/Qwen2.5-Coder-7B-Instruct-GPTQ-Int8', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-Coder-14B-Instruct-GPTQ-Int4', hf_model_id='Qwen/Qwen2.5-Coder-14B-Instruct-GPTQ-Int4', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-Coder-14B-Instruct-GPTQ-Int8', hf_model_id='Qwen/Qwen2.5-Coder-14B-Instruct-GPTQ-Int8', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-Coder-32B-Instruct-GPTQ-Int4', hf_model_id='Qwen/Qwen2.5-Coder-32B-Instruct-GPTQ-Int4', model_path=None, ms_revision=None, hf_revision=None), Model(ms_model_id='Qwen/Qwen2.5-Coder-32B-Instruct-GPTQ-Int8', hf_model_id='Qwen/Qwen2.5-Coder-32B-Instruct-GPTQ-Int8', model_path=None, ms_revision=None, hf_revision=None)], ignore_patterns=None, requires=None, tags=['coding'])], template='qwen2_5', get_function=<function get_model_tokenizer_with_flash_attn at 0x7fbf7497de10>, model_arch='llama', architectures=['Qwen2ForCausalLM'], additional_saved_files=[], torch_dtype=None, is_multimodal=False, is_reward=False, task_type=None, ignore_patterns=[], requires=['transformers>=4.37'], tags=[])",
+  "model_dir": "/oss/wangyujia/BIO/construction_finetuning/alpaca/v1-20250609-141541/checkpoint-50-merged",
+  "hub": "<class 'swift.hub.hub.MSHub'>",
+  "evaluation_strategy": "steps",
+  "training_args": "Seq2SeqTrainingArguments(output_dir='/nas/shared/kilab/wangyujia/BIO/sft/qwen-metal_ion_binding-08022141/v0-20250802-215822', overwrite_output_dir=False, do_train=False, do_eval=True, do_predict=False, eval_strategy=<IntervalStrategy.STEPS: 'steps'>, prediction_loss_only=False, per_device_train_batch_size=2, per_device_eval_batch_size=2, per_gpu_train_batch_size=None, per_gpu_eval_batch_size=None, gradient_accumulation_steps=4, eval_accumulation_steps=None, eval_delay=0, torch_empty_cache_steps=None, learning_rate=1e-05, weight_decay=0.1, adam_beta1=0.9, adam_beta2=0.95, adam_epsilon=1e-08, max_grad_norm=1.0, num_train_epochs=3.0, max_steps=-1, lr_scheduler_type=<SchedulerType.COSINE: 'cosine'>, lr_scheduler_kwargs=None, warmup_ratio=0.05, warmup_steps=0, log_level='passive', log_level_replica='warning', log_on_each_node=True, logging_dir='/nas/shared/kilab/wangyujia/BIO/sft/qwen-metal_ion_binding-08022141/v0-20250802-215822/runs', logging_strategy=<IntervalStrategy.STEPS: 'steps'>, logging_first_step=True, logging_steps=1, logging_nan_inf_filter=True, save_strategy=<SaveStrategy.STEPS: 'steps'>, save_steps=5, save_total_limit=5, save_safetensors=True, save_on_each_node=False, save_only_model=True, restore_callback_states_from_checkpoint=False, no_cuda=False, use_cpu=False, use_mps_device=False, seed=42, data_seed=42, jit_mode_eval=False, use_ipex=False, bf16=True, fp16=False, fp16_opt_level='O1', half_precision_backend='auto', bf16_full_eval=False, fp16_full_eval=False, tf32=None, local_rank=0, ddp_backend=None, tpu_num_cores=None, tpu_metrics_debug=False, debug=[], dataloader_drop_last=False, eval_steps=5, dataloader_num_workers=1, dataloader_prefetch_factor=10, past_index=-1, run_name='/nas/shared/kilab/wangyujia/BIO/sft/qwen-metal_ion_binding-08022141/v0-20250802-215822', disable_tqdm=False, remove_unused_columns=False, label_names=None, load_best_model_at_end=False, metric_for_best_model='loss', greater_is_better=False, ignore_data_skip=False, fsdp=[], fsdp_min_num_params=0, fsdp_config={'min_num_params': 0, 'xla': False, 'xla_fsdp_v2': False, 'xla_fsdp_grad_ckpt': False}, tp_size=0, fsdp_transformer_layer_cls_to_wrap=None, accelerator_config=AcceleratorConfig(split_batches=False, dispatch_batches=False, even_batches=True, use_seedable_sampler=True, non_blocking=False, gradient_accumulation_kwargs=None, use_configured_state=False), deepspeed={'fp16': {'enabled': 'auto', 'loss_scale': 0, 'loss_scale_window': 1000, 'initial_scale_power': 16, 'hysteresis': 2, 'min_loss_scale': 1}, 'bf16': {'enabled': 'auto'}, 'zero_optimization': {'stage': 3, 'offload_optimizer': {'device': 'none', 'pin_memory': True}, 'offload_param': {'device': 'none', 'pin_memory': True}, 'overlap_comm': False, 'contiguous_gradients': True, 'sub_group_size': 1000000000.0, 'reduce_bucket_size': 'auto', 'zero_quantized_weights': False, 'zero_quantized_gradients': False, 'stage3_prefetch_bucket_size': 'auto', 'stage3_param_persistence_threshold': 'auto', 'stage3_max_live_parameters': 1000000000.0, 'stage3_max_reuse_distance': 1000000000.0, 'stage3_gather_16bit_weights_on_model_save': True}, 'gradient_accumulation_steps': 'auto', 'gradient_clipping': 'auto', 'steps_per_print': 2000, 'train_batch_size': 'auto', 'train_micro_batch_size_per_gpu': 'auto', 'wall_clock_breakdown': False}, label_smoothing_factor=0.0, optim=<OptimizerNames.ADAMW_TORCH: 'adamw_torch'>, optim_args=None, adafactor=False, group_by_length=False, length_column_name='length', report_to=['tensorboard'], ddp_find_unused_parameters=None, ddp_bucket_cap_mb=None, ddp_broadcast_buffers=None, dataloader_pin_memory=True, dataloader_persistent_workers=False, skip_memory_metrics=True, use_legacy_prediction_loop=False, push_to_hub=False, resume_from_checkpoint=None, hub_model_id=None, hub_strategy=<HubStrategy.EVERY_SAVE: 'every_save'>, hub_token=None, hub_private_repo=None, hub_always_push=False, gradient_checkpointing=True, gradient_checkpointing_kwargs=None, include_inputs_for_metrics=False, include_for_metrics=[], eval_do_concat_batches=True, fp16_backend='auto', push_to_hub_model_id=None, push_to_hub_organization=None, push_to_hub_token=None, mp_parameters='', auto_find_batch_size=False, full_determinism=False, torchdynamo=None, ray_scope='last', ddp_timeout=1800, torch_compile=False, torch_compile_backend=None, torch_compile_mode=None, include_tokens_per_second=None, include_num_input_tokens_seen=None, neftune_noise_alpha=None, optim_target_modules=None, batch_eval_metrics=False, eval_on_start=False, use_liger_kernel=False, eval_use_gather_object=False, average_tokens_across_devices=None, sortish_sampler=False, predict_with_generate=False, generation_max_length=None, generation_num_beams=None, generation_config=None, check_model=True, acc_strategy='token', train_dataloader_shuffle=True, metric_warmup_step=0, fsdp_num=1, acc_steps=1, eval_use_evalscope=False, eval_datasets=[], eval_limit=None, eval_datasets_args=None, eval_generation_config=None, train_type='lora', optimizer=None, local_repo_path=None, galore_config=None)"
+}

BIO/sft/qwen-metal_ion_binding-08022141/v0-20250802-215822/val_dataset.jsonl

@@ -0,0 +1,57 @@
+{"messages": [{"role": "user", "content": "Based on the following protein's amino acid sequence, is the protein located on the membrane?.\nProtein amino acid sequence: TVPRFGVQTDQEEQLAKELEDTNKWGLDVFKVAELSGNRPLTAIIFSIFQERDLLKTFQIPADTLATYLLMLEGHYHANVAYHNSLHAADVAQSTHVLLATPALEAVFTDLEILAALFASAIHDVDHPGVSNQFLINTNSELALMYNDASVLENHHLAVGFKLLQAENCDIFQNLSAKQRLSLRRMVIDMVLATDMSKHMNLLADLKTMVETKKVTSLGVLLLDNYSDRIQVLQNLVHCADLSNPTKPLPLYRQWTDRIMAEFFQQGDRERESGLDISPMCDKHTASVEKSQVGFIDYIAHPLWETWADLVHPDAQDLLDTLEDNREWYQSKIPRSPSDLTNPERDGPDRFQFELTLE\nOptions:\n0.\"Yes\"\n1.\"No\"\n"}, {"role": "assistant", "content": "1"}]}
+{"messages": [{"role": "user", "content": "Based on the following protein's amino acid sequence, is the protein located on the membrane?.\nProtein amino acid sequence: MLTAEDKKLIQQAWEKAASHQEEFGAEALTRMFTTYPQTKTYFPHFDLSPGSDQVRGHGKKVLGALGNAVKNVDNLSQAMAELSNLHAYNLRVDPVNFKLLSQCIQVVLAVHMGKDYTPEVHAAFDKFLSAVSAVLAEKYR\nOptions:\n0.\"Yes\"\n1.\"No\"\n"}, {"role": "assistant", "content": "1"}]}
+{"messages": [{"role": "user", "content": "Based on the following protein's amino acid sequence, is the protein located on the membrane?.\nProtein amino acid sequence: PQVTLWQRPLVTIKIGGQLKEALLDTGADDTVLEEMSLPGRWKPKMIGGIGGFIKVRQYDQILIEICGHKAIGTVLVGPTPVNIIGRDLLTQIGMTLNFGGSSGPQVTLWQRPLVTIKIGGQLKEALLDTGADDTVLEEMSLPGRWKPKMIGGIGGFIKVRQYDQILIEICGHKAIGTVLVGPTPVNIIGRDLLTQIGMTLNF\nOptions:\n0.\"Yes\"\n1.\"No\"\n"}, {"role": "assistant", "content": "1"}]}
+{"messages": [{"role": "user", "content": "Based on the following protein's amino acid sequence, is the protein located on the membrane?.\nProtein amino acid sequence: ADKELKFLVVDDESTMRRIVRNLLKELGFNNVEEAEDGVDALNKLQAGGYGFVISDWMMPNMDGLELLKTIRADGAMSALPVLMVTALAKKENIIAAAQAGASGYVVKPFTAATLEEKLNKIFEKLGM\nOptions:\n0.\"Yes\"\n1.\"No\"\n"}, {"role": "assistant", "content": "0"}]}
+{"messages": [{"role": "user", "content": "Based on the following protein's amino acid sequence, is the protein located on the membrane?.\nProtein amino acid sequence: MRTRSTISTPNGITWYYEQEGTGPDIVLVPDGLGECQMFDSSVSQIAAQGFRVTTFDMPGMSRSAKAPAETYTEVTAQKLASYVISILDALDIKHATVWGCSSGASTVVALLLGYPDRIRNAMCHELPTKLLDHLSNTAVLEDEEISNILANVMLNDVSGGSEAWQALGVEVHARLHKNYPVWARGYPRTIPPSAPVQDVEALRGKPLDWTVGAATPTESFFDNIVTATKAGVNIGLLPGMHFPYVSHPDVFAKYVVETTQKHLWNSSSVDKLAAALEHHHHHH\nOptions:\n0.\"Yes\"\n1.\"No\"\n"}, {"role": "assistant", "content": "0"}]}
+{"messages": [{"role": "user", "content": "Based on the following protein's amino acid sequence, is the protein located on the membrane?.\nProtein amino acid sequence: MGSSHHHHHHSSGLVPRGSHMTEIATTSGARSVGLLSVGAYRPERVVTNDEICQHIDSSDEWIYTRTGIKTRRFAADDESAASMATEACRRALSNAGLSAADIDGVIVTTNTHFLQTPPAAPMVAASLGAKGILGFDLSAGCAGFGYALGAAADMIRGGGAATMLVVGTEKLSPTIDMYDRGNCFIFADGAAAVVVGETPFQGIGPTVAGSDGEQADAIRQDIDWITFAQNPSGPRPFVRLEGPAVFRWAAFKMGDVGRRAMDAAGVRPDQIDVFVPHQANSRINELLVKNLQLRPDAVVANDIEHTGNTSAASIPLAMAELLTTGAAKPGDLALLIGYGAGLSYAAQVVRMPKG\nOptions:\n0.\"Yes\"\n1.\"No\"\n"}, {"role": "assistant", "content": "0"}]}
+{"messages": [{"role": "user", "content": "Based on the following protein's amino acid sequence, is the protein located on the membrane?.\nProtein amino acid sequence: GLKLDLTWFDKSTEDFKGEEYSKDFGDDGSVMESLGVPFKDNVNNGCFDVIAEWVPLLQPYFNHQIDISDNEYFVSFDYRDGDW\nOptions:\n0.\"Yes\"\n1.\"No\"\n"}, {"role": "assistant", "content": "0"}]}
+{"messages": [{"role": "user", "content": "Based on the following protein's amino acid sequence, is the protein located on the membrane?.\nProtein amino acid sequence: MGSSHHHHHHSSGLVPRGSHMPAKKPYNKIVSHLLVAEPEKIYAMPDPTVPDSDIKALTTLCDLADRELVVIIGWAKHIPGFSTLSLADQMSLLQSAWMEILILGVVYRSLSFEDELVYADDYIMDEDQSKLAGLLDLNNAILQLVKKYKSMKLEKEEFVTLKAIALANSDSMHIEDVEAVQKLQDVLHEALQDYEAGQHMEDPRRAGKMLMTLPLLRQTSTKAVQHFYNIKLEGKVPMHKLFLEMLEAKV\nOptions:\n0.\"Yes\"\n1.\"No\"\n"}, {"role": "assistant", "content": "1"}]}
+{"messages": [{"role": "user", "content": "Based on the following protein's amino acid sequence, is the protein located on the membrane?.\nProtein amino acid sequence: VLSGTDKTNVKSIFSKIGGQADDYGAEALERMFVTYPQTKTYFPHFDVSPGSAQVKAHGKKVAGGLSEAANHIDDIATSLSKLSDLHAQKLRVDPVNFKLLGQCFLVVVAIHNPSALTPEAHASLDKFLCAVGLVLTAKYR\nOptions:\n0.\"Yes\"\n1.\"No\"\n"}, {"role": "assistant", "content": "1"}]}
+{"messages": [{"role": "user", "content": "Based on the following protein's amino acid sequence, is the protein located on the membrane?.\nProtein amino acid sequence: TTPLVHVASVEKGRSYEDFQKVYNAIALKLREDDEYDNYIGYGPVLVRLAWHTSGTWDKHDNTGGSYGGTYRFKKEFNDPSNAGLQNGFKFLEPIHKEFPWISSGDLFSLGGVTAVQEMQGPKIPWRCGRVDTPEDTTPDNGRLPDADKDADYVRTFFQRLNMNDREVVALMGAHTLGKTHLKNSGYEGPWTANNNVFDNSFYLNLLNEDWKLEKNDANNEQWDSKSGYLQLPTDYSLIQDPKYLSIVKEYANDQDKFFKDFSKAFEKLLENGITFPKDAPSPFIFKTLEEQGL\nOptions:\n0.\"Yes\"\n1.\"No\"\n"}, {"role": "assistant", "content": "1"}]}
+{"messages": [{"role": "user", "content": "Based on the following protein's amino acid sequence, is the protein located on the membrane?.\nProtein amino acid sequence: DFDCIPGWSAYDRYCYQAFSKPKNWEDAESFCEEGVKTSHLVSIESSGEGDFVAQLVAEKIKTSFQYVWIGLRIQNKEQQCRSEWSDASSVNYENLVKQFSKKCYALKKGTELRTWFNVYCGTENPEVCKYTPEC\nOptions:\n0.\"Yes\"\n1.\"No\"\n"}, {"role": "assistant", "content": "0"}]}
+{"messages": [{"role": "user", "content": "Based on the following protein's amino acid sequence, is the protein located on the membrane?.\nProtein amino acid sequence: QNNVPNTFTDPDSGITFNTWGLDEDSPQTQGGFTFGVALPSDALTTDASEFIGYLKCARNDESGWCGISLGGPMTNSLLITAWPHEDTVYTSLRFATGYAMPDVYEGDAEITQVSSSVNSTHFSLIFRCKNCLQWSHGGSSGGASTSGGVLVLGWVQAFDDPGNPTCPEQITLQQHDNGMGIWGAQLNTDAASPSYTDWAAQATKTVT\nOptions:\n0.\"Yes\"\n1.\"No\"\n"}, {"role": "assistant", "content": "1"}]}
+{"messages": [{"role": "user", "content": "Based on the following protein's amino acid sequence, is the protein located on the membrane?.\nProtein amino acid sequence: MFTGSIVAIVTPMDEKGNVCRASLKKLIDYHVASGTSAIVSVGTTGESATLNHDEHADVVMMTLDLADGRIPVIAGTGANATAEAISLTQRFNDSGIVGCLTVTPYWNRPSQEGLYQHFKAIAEHTDLPQILYNVPSRTGCDLLPETVGRLAKVKNIIGIKEATGNLTRVNQIKELVSDDFVLLSGDDASALDFMQLGGHGVISVTANVAARDMAQMCKLAAEGHFAEARVINQRLMPLHNKLFVEPNPIPVKWACKELGLVATDTLRLPMTPITDSGRETVRAALKHAGLL\nOptions:\n0.\"Yes\"\n1.\"No\"\n"}, {"role": "assistant", "content": "0"}]}
+{"messages": [{"role": "user", "content": "Based on the following protein's amino acid sequence, is the protein located on the membrane?.\nProtein amino acid sequence: MQNDAGEFVDLYVPRKCSASNRIIGAKDHASIQMNVAEVDKVTGRFNGQFKTYAICGAIRRMGESDDSILRLAKADGIVSKNF\nOptions:\n0.\"Yes\"\n1.\"No\"\n"}, {"role": "assistant", "content": "0"}]}
+{"messages": [{"role": "user", "content": "Based on the following protein's amino acid sequence, is the protein located on the membrane?.\nProtein amino acid sequence: GMRIPSAIQLHKASKTLTLRYGEDSYDLPAEFLRVHSPSAEVQGHGNPVLQYGKLNVGLVGVEPAGQYALKLSFDDGHDSGLFTWDYLYELATRKDQLWADYLAELASAGKSRDPDESVVKLML\nOptions:\n0.\"Yes\"\n1.\"No\"\n"}, {"role": "assistant", "content": "1"}]}
+{"messages": [{"role": "user", "content": "Based on the following protein's amino acid sequence, is the protein located on the membrane?.\nProtein amino acid sequence: ATSTKKLHKEPATLIKAIDGDTVKLMYKGQPMTFRLLLVDTPEFNEKYGPEASAFTKKMEENAKKIEVEFDKGQRTDKYGRGLAYIYADGKMVNEALVRQGLAKVAYVYKGNNTHEQLLRKAEAQAKKEKLNIWSEDNADSGQ\nOptions:\n0.\"Yes\"\n1.\"No\"\n"}, {"role": "assistant", "content": "1"}]}
+{"messages": [{"role": "user", "content": "Based on the following protein's amino acid sequence, is the protein located on the membrane?.\nProtein amino acid sequence: LTESSTSKFVKINEKGFSDFNIHYNEAGNGETVIMLHGGGPGAGGWSNYYRNVGPFVDAGYRVILKDSPGFNKSDAVVMDEQRGLVNARAVKGLMDALDIDRAHLVGNAMGGATALNFALEYPDRIGKLILMGPGGLGPSMFAPMPMEGIKLLFKLYAEPSYETLKQMLQVFLYDQSLITEELLQGRWEAIQRQPEHLKNFLISAQKAPLSTWDVTARLGEIKAKTFITWGRDDRFVPLDHGLKLLWNIDDARLHVFSKCGHWAQWEHADEFNRLVIDFLRHA\nOptions:\n0.\"Yes\"\n1.\"No\"\n"}, {"role": "assistant", "content": "0"}]}
+{"messages": [{"role": "user", "content": "Based on the following protein's amino acid sequence, is the protein located on the membrane?.\nProtein amino acid sequence: CGTSNGGQNTSPSTSSSSAKGEASALPKGQTITVWSWQTGPELQDVKQIAAQWAKAHGDKVIVVDQSSNPKGFQFYATAARTGKGPDVVFGMPHDNNGVFAEEGLMAPVPSGVLNTGLYAPNTIDAIKVNGTMYSVPVSVQVAAIYYNKKLVPQPPQTWAEFVKDANAHGFMYDQANLYFDYAIIGGYGGYVFKDNNGTLDPNNIGLDTPGAVQAYTLMRDMVSKYHWMTPSTNGSIAKAEFLAGKIGMYVSGPWDTADIEKAKIDFGVTPWPTLPNGKHATPFLGVITAFVNKESKTQAADWSLVQALTSAQAQQMYFRDSQQIPALLSVQRSSAVQSSPTFKAFVEQLRYAVPMPNIPQMQAVWQAMSILQNIIAGKVSPEQGAKDFVQNIQKGIMAQGS\nOptions:\n0.\"Yes\"\n1.\"No\"\n"}, {"role": "assistant", "content": "0"}]}
+{"messages": [{"role": "user", "content": "Based on the following protein's amino acid sequence, is the protein located on the membrane?.\nProtein amino acid sequence: MAVKVEYDLKRLRNIGIAAHIDAGKTTTTERILYYTGRIHKIGEVHEGAATMDFMEQERERGITITAAVTTCFWKDHRINIIDTPGHVDFTIEVERSMRVLDGAIVVFDSSQGVEPQSETVWRQAEKYKVPRIAFANKMDKTGADLWLVIRTMQERLGARPVVMQLPIGREDTFSGIIDVLRMKAYTYGNDLGTDIREIPIPEEYLDNAREYHEKLVEVAADFDENIMLKYLEGEEPTEEELVAAIRKGTIDLKITPVFLGSALKNKGVQLLLDAVVDYLPSPLDIPPIKGTTPEGEVVEIHPDPNGPLAALAFKIMADPYVGRLTFIRVYSGTLTSGSYVYNTTKGRKERVARLLRMHANHREEVEELKAGDLGAVVGLKETITGDTLVGEDAPRVILESIEVPEPVIDVAIEPKTKADQEKLSQALARLAEEDPTFRVSTHPETGQTIISGMGELHLEIIVDRLKREFKVDANVGKPQVAYRETITKPVDVEGKFIRQTGGRGQYGHVKIKVEPLPRGSGFEFVNAIVGGVIPKEYIPAVQKGIEEAMQSGPLIGFPVVDIKVTLYDGSYHEVDSSEMAFKIAGSMAIKEAVQKGDPVILEPIMRVEVTTPEEYMGDVIGDLNARRGQILGMEPRGNAQVIRAFVPLAEMFGYATDLRSKTQGRGSFVMFFDHYQEVPKQVQEKLIKGQ\nOptions:\n0.\"Yes\"\n1.\"No\"\n"}, {"role": "assistant", "content": "0"}]}
+{"messages": [{"role": "user", "content": "Based on the following protein's amino acid sequence, is the protein located on the membrane?.\nProtein amino acid sequence: MSLSIPTIKQVRAFVLRGGGADYHDQGDGHWIDDHISTPMGKYPEYRQSRRSFGINVLGTLVVEIEASDGNVGFAVTTGGEPAAYIVEKHLARFLEGARVTDIERIWDQMYNSTLYYGRKGLVINTISGVDLALWDLLGKVRREPVHQLLGGAVRDELQFYATGARPDLAQKMGFIGGKMPLHHGPSEGEEGLKKNLEELATMRERVGPDFWLMFDCWMSLDLNYATRLARGAREYGLKWIEEALPPDDYWGYAELRRNAPTGMMVTTGEHEATRWGFRMLLEMGCCDIIQPDVGWCGGVTELLKISALADAHNALVVPHGSSVYSYHFVATRQNSPFAEFLMMAPKADQVVPMFHPQLLGEPVPENGRMRLSRLDQPGFGVTLNPECQLHRPYTHEGHHHHHH\nOptions:\n0.\"Yes\"\n1.\"No\"\n"}, {"role": "assistant", "content": "0"}]}
+{"messages": [{"role": "user", "content": "Based on the following protein's amino acid sequence, is the protein located on the membrane?.\nProtein amino acid sequence: GSAKGNFCPLCDKCYDDDDYESKMMQCGKCDRWVHSKCENLSDEMYEILSNLPESVAYTCVNCTERH\nOptions:\n0.\"Yes\"\n1.\"No\"\n"}, {"role": "assistant", "content": "1"}]}
+{"messages": [{"role": "user", "content": "Based on the following protein's amino acid sequence, is the protein located on the membrane?.\nProtein amino acid sequence: NEAISVEKSIQEQKLNGYGVGSLIKFPVSSTAPTLDAKSFYKYFQLRDTLDDRLTAVTATEVSLEGTTLDPTDYKVDTKGQTVTVTFTAEGLKRIKAAPGKKVSAVFQGKVTEARNGAITNRAQVISDTVYAEQPPTPEEPPANPENPPTSNEVTSRWGDLLIKKVDNHQQGQDKAGLQGAQFQLYKAKNAYAGTCTKDKEGDPIAINGETTLTTDAQGAINVKGLFISDSIDGANRDNQKDATARCYVLVETKAPAGYVLPAGDGAVTPVKIEVGAVTTDNVTIENTKQ\nOptions:\n0.\"Yes\"\n1.\"No\"\n"}, {"role": "assistant", "content": "1"}]}
+{"messages": [{"role": "user", "content": "Based on the following protein's amino acid sequence, is the protein located on the membrane?.\nProtein amino acid sequence: MVTSPPTSSPSQRSYSPQDWLRGYQSQPQEWDYWVEDVEGSIPPDLQGTLYRNGPGLLEIGDRPLKHPFDGDGMVTAFKFPGDGRVHFQSKFVRTQGYVEEQKAGKMIYRGVFGSQPAGGWLKTIFDLRLKNIANTNITYWGDRLLALWQGGQPHRLEPSNLATIGLDDLGGILAEGQPLSAHPRIDPASTFDGGQPCYVTFSIKSSLSSTLTLLELDPQGKLLRQKTETFPGFAFIHDFAITPHYAIFLQNNVTLNGLPYLFGLRGAGECVQFHPDKPAQIILVPRDGGEIKRIPVQAGFVFHHANAFEENGKIILDSICYNSLPQVDTDGDFRSTNFDNLDPGQLWRFTIDPAAATVEKQLMVSRCCEFPVVHPQQVGRPYRYVYMGAAHHSTGNAPLQAILKVDLESGTETLRSFAPHGFAGEPIFVPRPGGVAEDDGWLLCLIYKADLHRSELVILDAQDITAPAIATLKLKHHIPYPLHGSWAQT\nOptions:\n0.\"Yes\"\n1.\"No\"\n"}, {"role": "assistant", "content": "1"}]}
+{"messages": [{"role": "user", "content": "Based on the following protein's amino acid sequence, is the protein located on the membrane?.\nProtein amino acid sequence: MAAGSITTLPALPEDGGSGAFPPGHFKDPKRLYCKNGGFFLRIHPDGRVDGVREKSDPHIKLQLQAEERGVVSIKGVSANRYLAMKEDGRLLASKSVTDECFFFERLESNNYNTYRSRKYTSWYVALKRTGQYKLGSKTGPGQKAILFLPMSAKS\nOptions:\n0.\"Yes\"\n1.\"No\"\n"}, {"role": "assistant", "content": "0"}]}
+{"messages": [{"role": "user", "content": "Based on the following protein's amino acid sequence, is the protein located on the membrane?.\nProtein amino acid sequence: GPMGGFQRGKYGTMAEGRSEDNLSATPPALRIILVGKTGCGKSATGNSILGQPVFESKLRAQSVTRTCQVKTGTWNGRKVLVVDTPSIFESQADTQELYKNIGDCYLLSAPGPHVLLLVIQLGRFTAQDTVAIRKVKEVFGTGAMRHVVILFTHKEDLGGQALDDYVANTDNCSLKDLVRECERRYCAFNNWGSVEEQRQQQAELLAVIERLGREREGSFHSNDLFLDAQLLQRTGAGACQEDYRQYQAKVEWQVEKHKQELRENESNWAYKALLRVK\nOptions:\n0.\"Yes\"\n1.\"No\"\n"}, {"role": "assistant", "content": "0"}]}
+{"messages": [{"role": "user", "content": "Based on the following protein's amino acid sequence, is the protein located on the membrane?.\nProtein amino acid sequence: METVEMVAIATVAGLLSLATVTGNILLMLSIKVNRQLQTVNNYFAFSLACADLIIGAFSMNLYTVYIIMGHWALGALACDLALALDYVASNAAVMNLLLISFDRYFSVTRPLSYRAKRTPRRALLMIGLAWLVSFVLWAPAILFWQYLVGERTVLAGQCYIQFLSQPIITFGTAMATFYLPVTVMCTLYWRIYRETENRANIFEMLRIDEGLRLKIYKDTEGYYTIGIGHLLTKSPSLNAAKSELDKAIGRNTNGVITKDEAEKLFNQDVDAAVRGILRNAKLKPVYDSLDAVRRAALINMVFQMGETGVAGFTNSLRMLQQKRWDEAAVNLAKSRWYNQTPNRAKRVITTFRTGTWDAYTFSLVKEKAALRTLSAILLAFILTWTPYNIMVLVSTFCKDCVPETLWELGYWLCYVNATINPMCYALCNKAFRDTFRLLLLARWDHHHHHHHHHH\nOptions:\n0.\"Yes\"\n1.\"No\"\n"}, {"role": "assistant", "content": "0"}]}
+{"messages": [{"role": "user", "content": "Based on the following protein's amino acid sequence, is the protein located on the membrane?.\nProtein amino acid sequence: GPGSEFMADDDVLFEDVYELCEVIGKGPFSVVRRCINRETGQQFAVKIVDVAKFTSSPGLSTEDLKREASICHMLKHPHIVELLETYSSDGMLYMVFEFMDGADLCFEIVKRADAGFVYSEAVASHYMRQILEALRYCHDNNIIHRDVKPHCVLLASKENSAPVKLGGFGVAIQLGESGLVAGGRVGTPHFMAPEVVKREPYGKPVDVWGCGVILFILLSGCLPFYGTKERLFEGIIKGKYKMNPRQWSHISESAKDLVRRMLMLDPAERITVYEALNHPWLKERDRYAYKIHLPETVEQLRKFNARRKLKGAVLAAVSSHKFNS\nOptions:\n0.\"Yes\"\n1.\"No\"\n"}, {"role": "assistant", "content": "0"}]}
+{"messages": [{"role": "user", "content": "Based on the following protein's amino acid sequence, is the protein located on the membrane?.\nProtein amino acid sequence: TPEAPYASLTEIEHLVQSVSKSYRETCQLRLEDLLRQRSNIFSREEVTGYQRKSMWEMWERCAHHLTEAIQYVVEFAKRLSGFMELSQNDQIVLLKAGAMEVVLVRMCRAYNADNRTVFFEGKYGGMELFRALGCSELISSIFDFSHSLSALHFSEDEIALYTALVLINAHRPGLQEKRKVEQLQYNLELAFHHHLCKTHRQSILAKLPPKGKLRSLCSQHVERLQIFQHLHPIVVQAAFPPLYKELFSTETESPVGLS\nOptions:\n0.\"Yes\"\n1.\"No\"\n"}, {"role": "assistant", "content": "1"}]}
+{"messages": [{"role": "user", "content": "Based on the following protein's amino acid sequence, is the protein located on the membrane?.\nProtein amino acid sequence: QVQLVQSGAEVKKPGASVKVSCKASGYTFTGYYMHWVRQAPGQGLEWMGWINPNSGGTNYAQKFQGRVTMTRDTSISTAYMELSRLRSDDTAVYYCARDGSGDDTSWHLHPWGQGTLVTVSSASTKGPSVFPLAPSSKSTSGGTAALGCLVKDYFPEPVTVSWNSGALTSGVHTFPAVLQSSGLYSLSSVVTVPSSSLGTQTYICNVNHKPSNTKVDKRVEPKSCDKTHHHHHH\nOptions:\n0.\"Yes\"\n1.\"No\"\n"}, {"role": "assistant", "content": "0"}]}
+{"messages": [{"role": "user", "content": "Based on the following protein's amino acid sequence, is the protein located on the membrane?.\nProtein amino acid sequence: QDGEALFKSKPCAACHSIDAKMVGPALKEVAAKYAGQEGAADLLAGHIKNGTQGNWGPIPMPPNPVTEEEAKTLAEWVLSLK\nOptions:\n0.\"Yes\"\n1.\"No\"\n"}, {"role": "assistant", "content": "1"}]}
+{"messages": [{"role": "user", "content": "Based on the following protein's amino acid sequence, is the protein located on the membrane?.\nProtein amino acid sequence: LTNSSLMPTLNPMIQQLALAIAASWQSLPLKPYQLPEDLGYVEGRLEGEKLVIENRCYQTPQFRKMQLELAKVGKGLDILHCVMFPEPLYGLPLFGCDIVAGPGGVSAAIADLSPTQSDRQLPAAYQKSLAELGQPEFEQQRELPPWGEIFSEYCLFIRPSNVTEEERFVQRVVDFLQIHCHQSIVAEPLSEAQTLEHRQGQIHYCQQQQKNDKTRRVLEKAFGEAWAERYMSQVLFDV\nOptions:\n0.\"Yes\"\n1.\"No\"\n"}, {"role": "assistant", "content": "1"}]}
+{"messages": [{"role": "user", "content": "Based on the following protein's amino acid sequence, is the protein located on the membrane?.\nProtein amino acid sequence: GIQGLAKLIADVAPSAIRENDIKSYFGRKVAIDASMSIYQFLIAVRQGGDVLQNEEGETTSHLMGMFYRTIRMMENGIKPVYVFDGKPPQLKSGELAKRSERRAEAEKQLQQAQAAGAEQEVEKFTKRLVKVTKQHNDECKHLLSLMGIPYLDAPSEAEASCAALVKAGKVYAAATEDMDCLTFGSPVLMRHLTASEAKKLPIQEFHLSRILQELGLNQEQFVDLCILLGSNYCESIRGIGPKRAVDLIQKHKSIEEIVRRLDPNKYPVPENWLHKEAHQLFLEPEVLDPESVELKWSEPNEEELIKFMCGEKQFSEERIRSGVKRLSKSRQGSTLEVLFQ\nOptions:\n0.\"Yes\"\n1.\"No\"\n"}, {"role": "assistant", "content": "1"}]}
+{"messages": [{"role": "user", "content": "Based on the following protein's amino acid sequence, is the protein located on the membrane?.\nProtein amino acid sequence: MHHHHHHSSGVDLGTENLYFQSMVDATRVPMDERFRTLKKKLEEGMVFTEYEQIPKKKANGIFSTAALPENAERSRIREVVPYEENRVELIPTKENNTGYINASHIKVVVGGAEWHYIATQGPLPHTCHDFWQMVWEQGVNVIAMVTAEEEGGRTKSHRYWPKLGSKHSSATYGKFKVTTKFRTDSVCYATTGLKVKHLLSGQERTVWHLQYTDWPDHGCPEDVQGFLSYLEEIQSVRRHTNSMLEGTKNRHPPIVVHCSAGVGRTGVLILSELMIYCLEHNEKVEVPMMLRLLREQRMFMIQTIAQYKFVYQVLIQFLQNSRLI\nOptions:\n0.\"Yes\"\n1.\"No\"\n"}, {"role": "assistant", "content": "0"}]}
+{"messages": [{"role": "user", "content": "Based on the following protein's amino acid sequence, is the protein located on the membrane?.\nProtein amino acid sequence: QAQITGRPEWIWLALGTALMGLGTLYFLVKGMGVSDPDAKKFYAITTLVAAIAFTMYLSMLLGYGLTMVPFGGEQNPIYWARYADWLFTTPLLLLDLALLVDADQGTILALVGADGIMIGTGLVGALTKVYSYRFVWWAISTAAMLYILYVLFFGFTSKAESMRPEVASTFKVLRNVTVVLWSAYPVVWLIGSEGAGIVPLNIETLLFMVLDVSAKVGFGLILLRSRAIFGEAEAPEPSAGDGAAATSD\nOptions:\n0.\"Yes\"\n1.\"No\"\n"}, {"role": "assistant", "content": "0"}]}
+{"messages": [{"role": "user", "content": "Based on the following protein's amino acid sequence, is the protein located on the membrane?.\nProtein amino acid sequence: MEKDGLCRADQQYECVAEIGEGAYGKVFKARDLKNGGRFVALKRVRVQTGEEGMPLSTIREVAVLRHLETFEHPNVVRLFDVCTVSRTDRETKLTLVFEHVDQDLTTYLDKVPEPGVPTETIKDMMFQLLRGLDFLHSHRVVHRDLKPQNILVTSSGQIKLADFGLARIYSFQMALTSVVVTLWYRAPEVLLQSSYATPVDLWSVGCIFAEMFRRKPLFRGSSDVDQLGKILDVIGLPGEEDWPRDVALPRQAFHSKSAQPIEKFVTDIDELGKDLLLKCLTFNPAKRISAYSALSHPYFQDLERCKENLDSHLPPSQNTSELNTA\nOptions:\n0.\"Yes\"\n1.\"No\"\n"}, {"role": "assistant", "content": "0"}]}
+{"messages": [{"role": "user", "content": "Based on the following protein's amino acid sequence, is the protein located on the membrane?.\nProtein amino acid sequence: AFELPPLPYAHDALQPHISKETLEFHHDKHHNTYVVNLNNLVPGTEFEGKTLEEIVKTSSGGIFNNAAQVWNHTFYWNCLSPNAGGQPTGALADAINAAFGSFDKFKEEFTKTSVGTFGSGWGWLVKKADGSLALASTIGAGCPLTIGDTPLLTCDVWEHAYYIDYRNLRPKYVEAFWNLVNWAFVAEQFEGKTYKV\nOptions:\n0.\"Yes\"\n1.\"No\"\n"}, {"role": "assistant", "content": "1"}]}
+{"messages": [{"role": "user", "content": "Based on the following protein's amino acid sequence, is the protein located on the membrane?.\nProtein amino acid sequence: AVGIGAVFLGFLGAAGSTMGAASMTLTVQARNLLSGIVQQQSNLLRAIEAQQHLLKLTVWGIKQLQARVLAVERYLRDQQLLGIWGCSGKLICCTNVPWNSSWSNRNLSEIWDNMTWLQWDKEISNYTQIIYGLLEESQNQQEKNEQDLLALD\nOptions:\n0.\"Yes\"\n1.\"No\"\n"}, {"role": "assistant", "content": "0"}]}
+{"messages": [{"role": "user", "content": "Based on the following protein's amino acid sequence, is the protein located on the membrane?.\nProtein amino acid sequence: SNADLLVPKLTASVTDGAVGVTVDAPVSVTAADGVLAAVTMVNDNGRPVAGRLSPDGLRWSTTEQLGYNRRYTLNATALGLGGAATRQLTFQTSSPAHLTMPYVMPGDGEVVGVGEPVAIRFDENIADRGAAEKAIKITTNPPVEGAFYWLNNREVRWRPEHFWKPGTAVDVAVNTYGVDLGEGMFGEDNVQTHFTIGDEVIATADDNTKILTVRVNGEVVKSMPTSMGKDSTPTANGIYIVGSRYKHIIMDSSTYGVPVNSPNGYRTDVDWATQISYSGVFVHSAPWSVGAQGHTNTSHGCLNVSPSNAQWFYDHVKRGDIVEVVNTVGGTLPGIDGLGDWNIPWDQWRAGNAKA\nOptions:\n0.\"Yes\"\n1.\"No\"\n"}, {"role": "assistant", "content": "1"}]}
+{"messages": [{"role": "user", "content": "Based on the following protein's amino acid sequence, is the protein located on the membrane?.\nProtein amino acid sequence: QAPTAVLNGNEVISGVLEGKVDTFKGIPFADPPLNDLRFKHPQPFTGSYQGLKANDFSPACMQLDPGNSLTLLDKALGLAKVIPEEFRGPLYDMAKGTVSMNEDCLYLNVFRPAGTKPDAKLPVMVWIYGGAFVYGSSAAYPGNSYVKESINMGQPVVFVSINYRTGPFGFLGGDAITAEGNTNAGLHDQRKGLEWVSDNIANFGGDPDKVMIFGESAGAMSVAHQLIAYGGDNTYNGKKLFHSAILQSGGPLPYHDSSSVGPDISYNRFAQYAGCDTSASANDTLECLRSKSSSVLHDAQNSYDLKDLFGLLPQFLGFGPRPDGNIIPDAAYELFRSGRYAKVPYISGNQEDEGTAFAPVALNATTTPHVKKWLQYIFYDASEASIDRVLSLYPQTLSVGSPFRTGILNALTPQFKRVAAILSDMLFQSPRRVMLSATKDVNRWTYLSTHLHNLVPFLGTFHGNELIFQFNVNIGPANSYLRYFISFANHHDPNVGTNLLQWDQYTDEGKEMLEIHMTDNVMRTDDYRIEGISNFETDVNLYG\nOptions:\n0.\"Yes\"\n1.\"No\"\n"}, {"role": "assistant", "content": "0"}]}
+{"messages": [{"role": "user", "content": "Based on the following protein's amino acid sequence, is the protein located on the membrane?.\nProtein amino acid sequence: MTGSLNRHSLLNGVKKMRIILCDTNEVVTNLWQESIPHAYIQNDKYLCIHHGHLQSLMDSMRKGDAIHHGHSYAIVSPGNSYGYLGGGFDKALYNYFGGKPFETWFRNQLGGRYHTVGSATVVDLQRCLEEKTIECRDGIRYIIHVPTVVAPSAPIFNPQNPLKTGFEPVFNAMWNALMHSPKDIDGLIIPGLCTGYAGVPPIISCKSMAFALRLYMAGDHISKELKNVLIMYYLQYPFEPFFPESCKIECQKLGIDIEMLKSFNVEKDAIELLIPRRILTLDL\nOptions:\n0.\"Yes\"\n1.\"No\"\n"}, {"role": "assistant", "content": "0"}]}
+{"messages": [{"role": "user", "content": "Based on the following protein's amino acid sequence, is the protein located on the membrane?.\nProtein amino acid sequence: GSSGSSGQPPLKNLLSLLKAYYALNAQPSAEELSKIADSVNLPLDVVKKWFEKMQAGQISVQSS\nOptions:\n0.\"Yes\"\n1.\"No\"\n"}, {"role": "assistant", "content": "1"}]}
+{"messages": [{"role": "user", "content": "Based on the following protein's amino acid sequence, is the protein located on the membrane?.\nProtein amino acid sequence: YKLICYYTSWSQYREGDGSCFPDAIDPFLCTHVIYSFANISNNEIDTWEWNDVTLYDTLNTLKNRNPNLKTLLSVGGWNFGSERFSKIASKTQSRRTFIKSVPPFLRTHGFDGLDLAWLYPGWRDKRHLTTLVKEMKAEFVREAQAGTEQLLLSAAVPAGKIAIDRGYDIAQISRHLDFISLLTYDFHGGWRGTVGHHSPLFRGNSDGSSRFSNADYAVSYMLRLGAPANKLVMGIPTFGRSYTLASSKTDVGAPISGPGIPGQFTKEKGTLAYYEICDFLHGATTHRFRDQQVPYATKGNQWVAYDDQESVKNKARYLKNRQLAGAMVWALDLDDFRGTFCGQNLTFPLTSAIKDVLARV\nOptions:\n0.\"Yes\"\n1.\"No\"\n"}, {"role": "assistant", "content": "0"}]}
+{"messages": [{"role": "user", "content": "Based on the following protein's amino acid sequence, is the protein located on the membrane?.\nProtein amino acid sequence: MSIEPVFILVRPQMGENIGAAARAMLNFGLGRLRIVDPRDGWPNPKAVAMASGAGRLLDHAGLFPTVAEAIRDCDYVFATTARGRELTKPVMTPERAMAHGRALTGEGRRVGILFGPERTGLENEDVALANAIVTVPVNPEFFSLNLAQCVLLLAYEWRRQHDETPPEVIDMARVDFASGLEVEKLGDHFEEKLEAAGFFFPPEKAPGMKLNLRNMWARLPLTRADVQTLHGMLRQIAWKLKQENLYFQ\nOptions:\n0.\"Yes\"\n1.\"No\"\n"}, {"role": "assistant", "content": "0"}]}
+{"messages": [{"role": "user", "content": "Based on the following protein's amino acid sequence, is the protein located on the membrane?.\nProtein amino acid sequence: GIKQYSQEELKEMALVEIAHELFEEHKKPVPFQELLNEIASLLGVKKEELGDRIAQFYTDLNIDGRFLALSDQTWGLRSWYPYDQLDEETQPTVKAKKKKAKKAVEEDLDLDEFEEIDEDDLDLDEVEEELDLEADDFDEEDLDEDDDDLEIEEDIIDEDDEDYDDEEEEIK\nOptions:\n0.\"Yes\"\n1.\"No\"\n"}, {"role": "assistant", "content": "0"}]}
+{"messages": [{"role": "user", "content": "Based on the following protein's amino acid sequence, is the protein located on the membrane?.\nProtein amino acid sequence: MGSSHHHHHHSSGLVPRGSHMENLYFQSSPNYDKWEMERTDITMKHKLGGGQYGEVYEGVWKKYSLTVAVKTLKEDTMEVEEFLKEAAVMKEIKHPNLVQLLGVCTREPPFYIITEFMTYGNLLDYLRECNRQEVNAVVLLYMATQISSAMEYLEKKNFIHRDLAARNCLVGENHLVKVADFGLSRLMTGDTYTAHAGAKFPIKWTAPESLAYNKFSIKSDVWAFGVLLWEIATYGMSPYPGIDLSQVYELLEKDYRMERPEGCPEKVYELMRACWQWNPSDRPSFAEIHQAFETMFQESSISDEVEKEL\nOptions:\n0.\"Yes\"\n1.\"No\"\n"}, {"role": "assistant", "content": "1"}]}
+{"messages": [{"role": "user", "content": "Based on the following protein's amino acid sequence, is the protein located on the membrane?.\nProtein amino acid sequence: MKIEEGKLVIWINGDKGYNGLAEVGKKFEKDTGIKVTVEHPDKLEEKFPQVAATGDGPDIIFWAHDRFGGYAQSGLLAEITPDKAFQDKLYPFTWDAVRYNGKLIAYPIAVEALSLIYNKDLLPNPPKTWEEIPALDKELKAKGKSALMFNLQEPYFTWPLIAADGGYAFKYENGKYDIKDVGVDNAGAKAGLTFLVDLIKNKHMNADTDYSIAEAAFNKGETAMTINGPWAWSNIDTSKVNYGVTVLPTFKGQPSKPFVGVLSAGINAASPNKELAKEFLENYLLTDEGLEAVNKDKPLGAVALKSYEEELAKDPRIAATMENAQKGEIMPNIPQMSAFWYAVRTAVINAASGRQTVDEALKDAQTNAAAAAGCKAA\nOptions:\n0.\"Yes\"\n1.\"No\"\n"}, {"role": "assistant", "content": "0"}]}
+{"messages": [{"role": "user", "content": "Based on the following protein's amino acid sequence, is the protein located on the membrane?.\nProtein amino acid sequence: MRVIFSEDHKLRNAKTELYGGELVPPFEAPFRAEWILAAVKEAGFDDVVAPARHGLETVLKVHDAGYLNFLETAWDRWKAAGYKGEAIATSFPVRRTSPRIPTDIEGQIGYYCNAAETAISPGTWEAALSSMASAIDGADLIAAGHKAAFSLCRPPGHHAGIDMFGGYCFINNAAVAAQRLLDKGAKKIAILDVDFHHGNGTQDIFYERGDVFFASLHGDPAEAFPHFLGYAEETGKGAGAGTTANYPMGRGTPYSVWGEALTDSLKRIAAFGAEAIVVSLGVDTFEQDPISFFKLTSPDYITMGRTIAASGVPLLVVMEGGYGVPEIGLNVANVLKGVAG\nOptions:\n0.\"Yes\"\n1.\"No\"\n"}, {"role": "assistant", "content": "1"}]}
+{"messages": [{"role": "user", "content": "Based on the following protein's amino acid sequence, is the protein located on the membrane?.\nProtein amino acid sequence: MAAECATQRAPGSVVELLGKSYPQDDHSNLTRKVLTRVGRNLHNQQHHPLWLIKERVKEHFYKQYVGRFGTPLFSVYDNLSPVVTTWQNFDSLLIPADHPSRKKGDNYYLNRTHMLRAHTSAHQWDLLHAGLDAFLVVGDVYRRDQIDSQHYPIFHQLEAVRLFSKHELFAGIKDGESLQLFEQSSRSAHKQETHTMEAVKLVEFDLKQTLTRLMAHLFGDELEIRWVDCYFPFTHPSFEMEINFHGEWLEVLGCGVMEQQLVNSAGAQDRIGWAFGLGLERLAMILYDIPDIRLFWCEDERFLKQFCVSNINQKVKFQPLSKYPAVINDISFWLPSENYAENDFYDLVRTIGGDLVEKVDLIDKFVHPKTHKTSHCYRITYRHMERTLSQREVRHIHQALQEAAVQLLGVEGRF\nOptions:\n0.\"Yes\"\n1.\"No\"\n"}, {"role": "assistant", "content": "0"}]}
+{"messages": [{"role": "user", "content": "Based on the following protein's amino acid sequence, is the protein located on the membrane?.\nProtein amino acid sequence: ATINNAGFESGFSNWNETDPAAISSDAYSGSKSLKIQGSPARVYQVVDIQPNTEYTLSAYVLGKGQIGVNDLNGLFKNQTFNVSSWTKVTKTFTSANTNSLQVFAKHYNNTSDVRFDNFSLIEGSGSNDGGSDGGSDNSNGSTIPSSITSGSIFDLEGDNPNPLVDDSTLVFVPLEAQHITPNGNGWRHEYKVKESLRVAMTQTYEVFEATVKVEMSDGGKTIISQHHASDTGTISKVYVSDTDESGFNDSVANNGIFDVYVRLRNTSGNEEKFALGTMTSGETFNLRVVNNYGDVEVTAFGNSFGIPVEDDSQSYFKFGNYLQSQDPYTLDKCGEAGNSNSFKNCFEDLGITESKVTMTNVSYTRETNHHHHHH\nOptions:\n0.\"Yes\"\n1.\"No\"\n"}, {"role": "assistant", "content": "1"}]}
+{"messages": [{"role": "user", "content": "Based on the following protein's amino acid sequence, is the protein located on the membrane?.\nProtein amino acid sequence: KIQVKCSAPNSVTITNASGGLYLVEYPEGYVAYSKATEVTGKLVHANFGTKKDFEDLDYAVNGSIVIVRAGKITIAEKVANAQSFNAIGVLIYKDRTKYPISRADEPLQGHSGLPSIPVQTISREAAEKLFQNMERDCPRSWNTDSSCKLELLQNRNVKLTVNNCLKEGTSGHHHHHH\nOptions:\n0.\"Yes\"\n1.\"No\"\n"}, {"role": "assistant", "content": "1"}]}
+{"messages": [{"role": "user", "content": "Based on the following protein's amino acid sequence, is the protein located on the membrane?.\nProtein amino acid sequence: MDIKINDITLGNNSPFVLFGGINVLESLDSTLQTCAHYVEVTRKLGIPYIFKASFDKANRSSIHSYRGVGLEEGLKIFEKVKAEFGIPVITDVHEPHQCQPVAEVCDVIQLPAFLARQTDLVVAMAKTGNVVNIKKPQFLSPSQMKNIVEKFHEAGNGKLILCERGSSFGYDNLVVDMLGFGVMKQTCGNLPVIFDVTHSLATRDAGSAASGGRRAQALDLALAGMATRLAGLFLESHPDPKLAKCDGPSALPLHLLEDFLIRIKALDDLIKSQPILTIE\nOptions:\n0.\"Yes\"\n1.\"No\"\n"}, {"role": "assistant", "content": "0"}]}
+{"messages": [{"role": "user", "content": "Based on the following protein's amino acid sequence, is the protein located on the membrane?.\nProtein amino acid sequence: MPPGKCLFSGVFCNMAEKQTAKRNRREEILQSLALMLESSDGSQRITTAKLAASVGVSEAALYRHFPSKTRMFDSLIEFIEDSLITRINLILKDEKDTTARLRLIVLLLLGFGERNPGLTRILTGHALMFEQDRLQGRINQLFERIEAQLRQVLREKRMREGEGYATDETLLASQILAFCEGMLSRFVRSEFKYRPTDDFDARWPLIAAQLQ\nOptions:\n0.\"Yes\"\n1.\"No\"\n"}, {"role": "assistant", "content": "0"}]}
+{"messages": [{"role": "user", "content": "Based on the following protein's amino acid sequence, is the protein located on the membrane?.\nProtein amino acid sequence: GQRQLLLASENPQQFMDYFSEEFRNDFLELLRRRFGTKRVHNNIVYNEYISHREHIHMNATQWETLTDFTKWLGREGLCKVDETPKGWYIQYIDRDPETIRRQLELEKKKK\nOptions:\n0.\"Yes\"\n1.\"No\"\n"}, {"role": "assistant", "content": "1"}]}
+{"messages": [{"role": "user", "content": "Based on the following protein's amino acid sequence, is the protein located on the membrane?.\nProtein amino acid sequence: MSYYHHHHHHLESTSLYKKAGFMSLQRIVRVSLEHPTSAVCVAGVETLVDIYGSVPEGTEMFEVYGTPGVDIYISPNMERGRERADTRRWRFDATLEIIVVMNSPSNDLNDSHVQISYHSSHEPLPLAYAVLYLTCVDISLDCDLNCEGRQDRNFVDKRQWVWGPSGYGGILLVNCDRDDPSCDVQDNCDQHVHCLQDLEDMSVMVLRTQGPAALFDDHKLVLHTSSYDAKRAQVFHICGPEDVCEAYRHVLGQDKVSYEVPRLHGDEERFFVEGLSFPDAGFTGLISFHVTLLDDSNEDFSASPIFTDTVVFRVAPWIMTPSTLPPLEVYVCRVRNNTCFVDAVAELARKAGCKLTICPQAENRNDRWIQDEMELGYVQAPHKTLPVVFDSPRNGELQDFPYKRILGPDFGYVTREPRDRSVSGLDSFGNLEVSPPVVANGKEYPLGRILIGGNLPGSSGRRVTQVVRDFLHAQKVQPPVELFVDWLAVGHVDEFLSFVPAPDGKGFRMLLASPGACFKLFQEKQKCGHGRALLFQGVVDDEQVKTISINQVLSNKDLINYNKFVQSCIDWNREVLKRELGLAECDIIDIPQLFKTERKKATAFFPDLVNMLVLGKHLGIPKPFGPIINGCCCLEEKVRSLLEPLGLHCTFIDDFTPYHMLHGEVHCGTNVCRKPFSFKWWNMVP\nOptions:\n0.\"Yes\"\n1.\"No\"\n"}, {"role": "assistant", "content": "1"}]}
+{"messages": [{"role": "user", "content": "Based on the following protein's amino acid sequence, is the protein located on the membrane?.\nProtein amino acid sequence: MANGANRVDLDGKPIQPLTICMIGAGGFIGSHLCEKLLTETPHKVLALDVYNDKIKHLLEPDTVEWSGRIQFHRINIKHDSRLEGLVKMADLIINLAAICTPADYNTRPLDTIYSNFIDALPVVKYCSENNKRLIHFSTCEVYGKTIGSFLPKDHPLRDDPAFYVLKEDISPCIFGSIEKQRWSYACAKQLIERLVYAEGAENGLEFTIVRPFNWIGPRMDFIPGIDGPSEGVPRVLACFSNNLLRREPLKLVDGGESQRTFVYINDAIEAVLLMIENPERANGHIFNVGNPNNEVTVRQLAEMMTEVYAKVSGEGAIESPTVDVSSKEFYGEGYDDSDKRIPDMTIINRQLGWNPKTSLWDLLESTLTYQHRTYAEAVKKATSKPVAS\nOptions:\n0.\"Yes\"\n1.\"No\"\n"}, {"role": "assistant", "content": "0"}]}
+{"messages": [{"role": "user", "content": "Based on the following protein's amino acid sequence, is the protein located on the membrane?.\nProtein amino acid sequence: PISPIETVPVKLKPGMDGPKVKQWPLTEEKIKALVEICTELEKEGKISKIGPENPYNTPVFAIKKKDSTKWRKLVDFRELNKRTQDFWEVQLGIPHPAGLKKKKSVTVLDVGDAYFSVPLDEDFRKYTAFTIPSINNETPGIRYQYNVLPQGWKGSPAIFQSSMTKILEPFRKQNPDIVIYQYMDDLYVGSDLEIGQHRTKIEELRQHLLRWGLYTPDKKHQKEPPFLWMGYELHPDKWTVQPIVLPEKDSWTVNDIQKLVGKLNWASQIYPGIKVRQLCKLLRGTKALTEVIPLTEEAELELAENREILKEPVHGVYYDPSKDLIAEIQKQGQGQWTYQIYQEPFKNLKTGKYARMRGAHTNDVKQLTEAVQKITTESIVIWGKTPKFKLPIQKETWETWWTEYWQATWIPEWEFVNTPPLVKLWYQLEKEPIVGAETFYVDGAANRETKLGKAGYVTNRGRQKVVTLTDTTNQKTELQAIYLALQDSGLEVNIVTDSQYALGIIQAQPDQSESELVNQIIEQLIKKEKVYLAWVPAHKGIGGNEQVDKLVSAGIRKVL\nOptions:\n0.\"Yes\"\n1.\"No\"\n"}, {"role": "assistant", "content": "1"}]}
+{"messages": [{"role": "user", "content": "Based on the following protein's amino acid sequence, is the protein located on the membrane?.\nProtein amino acid sequence: HHHHHHFNLPPGNYKKGGTVDGTRDRSDTHIQFQISPEGNGEVLLKSTETGQYLRINPDGTVDGTRDRSDTHIQFQISPEGNGEVLLKSTETGQYLRINPDGTVDGTRDRSDTHIQFQISPEGNGEVLLKSTETGQYLRINP\nOptions:\n0.\"Yes\"\n1.\"No\"\n"}, {"role": "assistant", "content": "0"}]}

BioReason/.gitignore

@@ -0,0 +1,180 @@
+# Byte-compiled / optimized / DLL files
+.idea/
+__pycache__/
+*.py[cod]
+*$py.class
+wandb/
+.DS_Store
+.vscode/
+.venv/
+.env
+.pytest_cache/
+
+# C extensions
+*.so
+
+outputs/
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# UV
+#   Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#uv.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/latest/usage/project/#working-with-version-control
+.pdm.toml
+.pdm-python
+.pdm-build/
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+# PyPI configuration file
+.pypirc

BioReason/LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

BioReason/README.md

@@ -0,0 +1,148 @@
+<h1 align="center">
+🧬 BioReason<br>Incentivizing Multimodal Biological Reasoning<br>within a DNA-LLM Model
+</h1>
+
+<p align="center">
+  <a href="https://www.arxiv.org/abs/2505.23579" target="_blank"><img src="https://img.shields.io/badge/arXiv-2505.23579-FF6B6B?style=for-the-badge&logo=arxiv&logoColor=white" alt="arXiv"></a>
+  <a href="https://github.com/bowang-lab/BioReason"><img src="https://img.shields.io/badge/GitHub-Code-4A90E2?style=for-the-badge&logo=github&logoColor=white" alt="GitHub"></a>
+  <a href="https://bowang-lab.github.io/BioReason/"><img src="https://img.shields.io/badge/Website-Online-00B89E?style=for-the-badge&logo=internet-explorer&logoColor=white" alt="Website"></a>
+  <a href="https://huggingface.co/collections/wanglab/bioreason-683cd17172a037a31d208f70"><img src="https://img.shields.io/badge/HuggingFace-Dataset-FFBF00?style=for-the-badge&logo=huggingface&logoColor=white" alt="HuggingFace Dataset"></a>
+</p>
+
+<br>
+
+## Updates [Jun 10, 2025]
+- We are integrating vLLM to improve the speed and efficiency of the GRPO pipeline. We expect this to be pushed by end of week.
+- Checkpoints along with the custom DNA-LLM model class will be released on HuggingFace by end of week.
+- More training results with GRPO will be shared soon.
+
+<br>
+
+## Abstract
+
+Unlocking deep, interpretable biological reasoning from complex genomic data is a major AI challenge hindering scientific discovery. Current DNA foundation models, despite strong sequence representation, struggle with multi-step reasoning and lack inherent transparent, biologically intuitive explanations. We introduce BioReason, a pioneering architecture that, for the first time, deeply integrates a DNA foundation model with a large language model (LLM). This novel connection enables the LLM to directly process and reason with genomic information as a fundamental input, fostering a new form of multimodal biological understanding. BioReason's sophisticated multi-step reasoning is developed through supervised fine-tuning and targeted reinforcement learning, guiding the system to generate logical, biologically coherent deductions. On biological reasoning benchmarks including KEGG-based disease pathway prediction—where accuracy improves from 88% to 97%—and variant effect prediction, BioReason demonstrates an average 15% performance gain over strong single-modality baselines.
+
+<br>
+
+## Key Contributions
+
+• **Novel multimodal architecture**: The first successful integration of a DNA foundation model with an LLM, establishing a new methodology for AI-driven biological studies.
+
+• **Advanced reasoning methodology**: A systematic training approach combining supervised fine-tuning and reinforcement learning that incentivizes multi-step biological reasoning.
+
+• **New biological reasoning benchmarks**: Development and curation of novel benchmarks for evaluating biological reasoning capabilities, including an annotated reasoning dataset for gene pathway and disease prediction from KEGG.
+
+• **Empirical performance improvements**: Demonstration that BioReason outperforms both DNA foundation models and LLMs used independently or in simple combination, with average performance gains of 15%+ over baseline.
+
+• **Interpretable reasoning traces**: A mechanism for generating step-by-step biological reasoning traces that provide interpretable predictions, enhancing scientific insight and hypothesis generation.
+
+<br>
+
+## Datasets
+
+The datasets used to train and evaluate BioReason can be found on our [HuggingFace collection](https://huggingface.co/collections/wanglab/bioreason-683cd17172a037a31d208f70) with detailed download and usage instructions.
+
+<br>
+
+## Checkpoints
+
+We will release the checkpoints soon!
+
+<br>
+
+## Installation
+
+### Prerequisites
+- Python 3.11+
+- CUDA/GPU for best performance
+
+### Installation Steps
+```bash
+# Clone the repository
+git clone https://github.com/bowang-lab/BioReason.git
+cd BioReason
+
+# Install package
+pip install -e .
+```
+
+<br>
+
+## Results
+
+### KEGG-Derived Biological Reasoning Task
+Performance comparison on 290 test datapoints for multi-step mechanistic reasoning:
+
+| Model | Accuracy | F1-Score | Precision | Recall |
+|-------|----------|----------|-----------|---------|
+| [DNA] NT - 500M | 86.55 | 69.76 | 73.23 | 66.61 |
+| [DNA] Evo2 - 1B | 88.28 | 72.43 | 75.23 | 69.83 |
+| [LLM] Qwen3 - 1B | 85.17 | 65.71 | 71.39 | 64.19 |
+| [LLM] Qwen3 - 4B | 93.48 | 85.44 | 88.31 | 86.72 |
+| [DNA-LLM] NT + Qwen3 - 1B | 88.42 | 72.13 | 75.42 | 71.91 |
+| [DNA-LLM] NT + Qwen3 - 1B (+RL) | 89.66 | 74.11 | 78.82 | 72.96 |
+| [DNA-LLM] NT + Qwen3 - 4B | 96.90 | **89.03** | **90.99** | **89.38** |
+| [DNA-LLM] Evo2 + Qwen3 - 1B | 90.42 | 75.62 | 77.42 | 73.91 |
+| [DNA-LLM] Evo2 + Qwen3 - 4B | **97.24** | 86.30 | 86.75 | 87.25 |
+
+### Variant Effect Prediction Benchmarks
+Performance on pathogenic/benign classification:
+
+| Model | Variant Effect - Coding | | Variant Effect - Non-SNV | |
+|-------|------------|----------|------------|----------|
+| | Accuracy | F1-Score | Accuracy | F1-Score |
+| [DNA] NT - 500M | 60.91 | 45.20 | 67.93 | 65.97 |
+| [DNA] Evo2 - 1B | 70.07 | 49.19 | 76.17 | 66.51 |
+| [LLM] Qwen3 - 1B | 46.55 | 34.82 | 70.67 | 76.21 |
+| [LLM] Qwen3 - 4B | 48.99 | 39.58 | 61.86 | 67.60 |
+| [DNA-LLM] NT + Qwen3 - 1B | 55.58 | 54.50 | 72.82 | 76.93 |
+| [DNA-LLM] NT + Qwen3 - 4B | 60.94 | 55.66 | 65.59 | 73.00 |
+| [DNA-LLM] Evo2 + Qwen3 - 1B | 72.83 | 68.90 | **88.20** | **89.91** |
+| [DNA-LLM] Evo2 + Qwen3 - 4B | **80.21** | **80.00** | 83.85 | 85.02 |
+
+<br>
+
+## Citation
+
+If you find this work useful, please cite our paper:
+
+```bibtex
+@misc{fallahpour2025bioreasonincentivizingmultimodalbiological,
+      title={BioReason: Incentivizing Multimodal Biological Reasoning within a DNA-LLM Model}, 
+      author={Adibvafa Fallahpour and Andrew Magnuson and Purav Gupta and Shihao Ma and Jack Naimer and Arnav Shah and Haonan Duan and Omar Ibrahim and Hani Goodarzi and Chris J. Maddison and Bo Wang},
+      year={2025},
+      eprint={2505.23579},
+      archivePrefix={arXiv},
+      primaryClass={cs.LG},
+      url={https://arxiv.org/abs/2505.23579}, 
+}
+```
+
+<br>
+
+## Authors
+
+- **Adibvafa Fallahpour**¹²³⁵ * (adibvafa.fallahpour@mail.utoronto.ca)
+- **Andrew Magnuson**¹² *
+- **Purav Gupta**¹² *
+- **Shihao Ma**¹²³
+- **Jack Naimer**¹²³
+- **Arnav Shah**¹²³
+- **Haonan Duan**¹²
+- **Omar Ibrahim**³
+- **Hani Goodarzi**†⁴⁶
+- **Chris J. Maddison**†¹²⁷
+- **Bo Wang**†¹²³
+
+¹ University of Toronto ² Vector Institute ³ University Health Network (UHN) <br>
+⁴ Arc Institute ⁵ Cohere ⁶ University of California, San Francisco ⁷ Google DeepMind
+
+<br>
+* Equal contribution <br>
+† Equal advising
+
+---
+
+<p align="center">
+Made with ❤️ at University of Toronto, Vector Institute, and University Health Network
+</p>

BioReason/bioreason/dataset/__init__.py

@@ -0,0 +1,11 @@
+from .kegg import KEGGDataset, split_kegg_dataset
+from .utils import torch_to_hf_dataset, truncate_dna
+from .variant_effect import get_format_variant_effect_function
+
+__all__ = [
+    "KEGGDataset",
+    "split_kegg_dataset",
+    "torch_to_hf_dataset",
+    "truncate_dna",
+    "get_format_variant_effect_function",
+]

BioReason/bioreason/dataset/kegg.py

@@ -0,0 +1,382 @@
+import json
+import os
+import random
+import sys
+import torch
+from torch.utils.data import Dataset, DataLoader
+from typing import Any, Dict, List, Tuple
+
+from bioreason.dataset.utils import torch_to_hf_dataset
+from bioreason.models.dl.processing_dl import DLProcessor
+from trl.data_utils import maybe_apply_chat_template
+
+
+class KEGGDataset(Dataset):
+    """Dataset for KEGG data."""
+
+    def __init__(self, data_dir: str):
+        """
+        Initialize the dataset by loading all JSON files from the given directory.
+
+        Args:
+            data_dir: Path to the directory containing JSON files
+        """
+        self.data_dir = data_dir
+        self.data = []
+
+        # Load all JSON files
+        json_files = sorted([f for f in os.listdir(data_dir) if f.endswith(".json")])
+
+        # Process each file
+        for filename in json_files:
+            file_path = os.path.join(data_dir, filename)
+            kegg_id = filename.split("_")[1]
+
+            with open(file_path, "r", encoding="utf-8") as f:
+                item = json.load(f)
+                item["kegg_id"] = kegg_id
+                processed_item = self._process_item(item)
+                self.data.append(processed_item)
+
+    def _process_item(self, item: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Process a single data item to format fields as required.
+
+        Args:
+            item: Original data item from JSON
+
+        Returns:
+            Processed data item
+        """
+        # Extract question as is
+        question = item.get("question", "")
+
+        # Convert answer to lowercase and strip whitespace
+        answer = item.get("answer", "").lower().strip()
+
+        # Combine reasoning steps into a single paragraph with newlines
+        reasoning_steps = item.get("reasoning", {}).get("reasoning_steps", [])
+        reasoning = "\n".join(reasoning_steps)
+
+        # Convert sequences to uppercase and strip whitespace
+        reference_sequence = item.get("reference_sequence", "").upper().strip()
+        variant_sequence = item.get("variant_sequence", "").upper().strip()
+
+        return {
+            "question": question,
+            "answer": answer,
+            "reasoning": reasoning,
+            "reference_sequence": reference_sequence,
+            "variant_sequence": variant_sequence,
+        }
+
+    def __len__(self) -> int:
+        """Return the number of items in the dataset."""
+        return len(self.data)
+
+    def __getitem__(self, idx: int) -> Dict[str, Any]:
+        """Return a specific item from the dataset."""
+        return self.data[idx]
+
+
+def split_kegg_dataset(
+    dataset: KEGGDataset,
+    train_ratio: float = 0.8,
+    val_ratio: float = 0.1,
+    test_ratio: float = 0.1,
+    seed: int = 42,
+) -> Tuple[KEGGDataset, KEGGDataset, KEGGDataset]:
+    """
+    Split a KEGG dataset into train, validation, and test sets.
+
+    Args:
+        dataset: The dataset to split
+        train_ratio: Proportion of data for training
+        val_ratio: Proportion of data for validation
+        test_ratio: Proportion of data for testing
+        batch_size: Batch size for the dataloaders
+        seed: Random seed for reproducibility
+
+    Returns:
+        Tuple of (train_dataset, val_dataset, test_dataset)
+    """
+    # Calculate the size of each split
+    dataset_size = len(dataset)
+    train_size = int(train_ratio * dataset_size)
+    val_size = int(val_ratio * dataset_size)
+    test_size = dataset_size - train_size - val_size
+    assert train_ratio + val_ratio + test_ratio == 1.0, "Ratios must sum to 1"
+
+    # Set the random seed
+    torch.manual_seed(seed)
+    random.seed(seed)
+
+    # Split the dataset
+    train_dataset, val_dataset, test_dataset = torch.utils.data.random_split(
+        dataset, [train_size, val_size, test_size]
+    )
+
+    return train_dataset, val_dataset, test_dataset
+
+
+def create_kegg_dataloader(
+    data_dir: str,
+    batch_size: int = 2,
+    shuffle: bool = True,
+    num_workers: int = 2,
+    pin_memory: bool = True,
+) -> DataLoader:
+    """
+    Create a DataLoader for the KEGG dataset.
+
+    Args:
+        data_dir: Path to the directory containing JSON files
+        batch_size: Batch size for the dataloader
+        shuffle: Whether to shuffle the data
+        num_workers: Number of worker processes for loading data
+        pin_memory: Whether to pin memory for faster data transfer
+
+    Returns:
+        DataLoader for the KEGG dataset
+    """
+    dataset = KEGGDataset(data_dir)
+    return DataLoader(
+        dataset,
+        batch_size=batch_size,
+        shuffle=shuffle,
+        num_workers=num_workers,
+        pin_memory=pin_memory,
+    )
+
+
+def get_format_kegg_function(model_name: str) -> Any:
+    """
+    Get the appropriate format function for a given model name.
+    """
+    if model_name.lower() == "llm":
+        return format_kegg_for_llm
+    elif model_name.lower() == "dna-llm":
+        return format_kegg_for_dna_llm
+    else:
+        raise ValueError(f"Unsupported model name: {model_name}")
+
+
+def format_kegg_for_dna_llm(example: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Format a KEGG example into the required chat format for DNA-LLM.
+    """
+    return {
+        "prompt": [
+            {
+                "role": "user",
+                "content": [
+                    *({"type": "dna", "text": None} for _ in range(2)),
+                    {"type": "text", "text": example["question"].strip()},
+                ],
+            },
+            {
+                "role": "assistant",
+                "reasoning_content": example["reasoning"].strip(),
+                "content": [
+                    {"type": "text", "text": f"Answer: {example['answer'].strip()}"},
+                ],
+            },
+        ],
+        "dna_sequences": [
+            example["reference_sequence"],
+            example["variant_sequence"],
+        ],
+        "answer": example["answer"],
+    }
+
+
+def format_kegg_for_llm(example: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Format a KEGG example into the required chat format for LLM.
+    """
+    question = f"Reference sequence: {example['reference_sequence']}\nVariant sequence: {example['variant_sequence']}\nQuestion: {example['question']}"
+    return {
+        "prompt": [
+            {
+                "role": "user",
+                "content": [
+                    *({"type": "dna", "text": None} for _ in range(2)),
+                    {"type": "text", "text": question.strip()},
+                ],
+            },
+            {
+                "role": "assistant",
+                "reasoning_content": example["reasoning"].strip(),
+                "content": [
+                    {"type": "text", "text": f"Answer: {example['answer'].strip()}"},
+                ],
+            },
+        ],
+        "dna_sequences": [
+            "",
+            "",
+        ],
+        "answer": example["answer"],
+    }
+
+
+def qwen_dna_collate_fn(
+    examples: List[Dict],
+    processor: DLProcessor,
+    max_length_text: int,
+    max_length_dna: int,
+    return_answer_in_batch: bool = False,
+) -> Dict:
+    """
+    Custom collate function for Qwen DNA models.
+
+    Creates a batch with proper labels for supervised fine-tuning where only
+    the assistant responses contribute to the loss calculation.
+    """
+    prompts_text = [
+        maybe_apply_chat_template(example, processor)["prompt"] for example in examples
+    ]
+    batch_dna_sequences = [example["dna_sequences"] for example in examples]
+
+    batch = processor(
+        text=prompts_text,
+        batch_dna_sequences=batch_dna_sequences,
+        return_tensors="pt",
+        padding=True,
+        padding_side="left",
+        add_special_tokens=False,
+        max_length_text=max_length_text,
+        max_length_dna=max_length_dna,
+    )
+
+    # Create labels tensor filled with -100 (ignored in loss calculation)
+    labels = torch.full_like(batch["input_ids"], -100)
+
+    # Get token IDs for special markers
+    assistant_start_marker = "<|im_start|>assistant\n"
+    im_end_marker = "<|im_end|>"
+
+    assistant_start_token_ids = processor.tokenizer.encode(
+        assistant_start_marker, add_special_tokens=False
+    )
+    im_end_token_ids = processor.tokenizer.encode(
+        im_end_marker, add_special_tokens=False
+    )
+
+    # Convert token arrays to tensors for faster comparison
+    assistant_marker_tensor = torch.tensor(
+        assistant_start_token_ids, device=batch["input_ids"].device
+    )
+    im_end_marker_tensor = torch.tensor(
+        im_end_token_ids, device=batch["input_ids"].device
+    )
+
+    # Get dimensions for easier reference
+    assistant_marker_len = len(assistant_start_token_ids)
+    im_end_marker_len = len(im_end_token_ids)
+
+    # For each sequence in the batch
+    for i in range(batch["input_ids"].shape[0]):
+        input_ids = batch["input_ids"][i]
+        seq_len = input_ids.size(0)
+
+        # Track assistant sections
+        assistant_sections = []
+
+        # Find all assistant start markers
+        start_positions = []
+        for pos in range(seq_len - assistant_marker_len + 1):
+            if torch.all(
+                input_ids[pos : pos + assistant_marker_len] == assistant_marker_tensor
+            ):
+                start_positions.append(
+                    pos + assistant_marker_len
+                )  # Store position after marker
+
+        # Find all end markers
+        end_positions = []
+        for pos in range(seq_len - im_end_marker_len + 1):
+            if torch.all(
+                input_ids[pos : pos + im_end_marker_len] == im_end_marker_tensor
+            ):
+                end_positions.append(pos)  # Store position at start of end marker
+
+        # Match start and end markers to create sections
+        for start_pos in start_positions:
+            # Find the next end marker after this start position
+            valid_ends = [pos for pos in end_positions if pos > start_pos]
+            if valid_ends:
+                end_pos = min(valid_ends)  # Take the first end marker after start
+                # Only include content between markers (not the markers themselves)
+                if start_pos < end_pos:
+                    assistant_sections.append((start_pos, end_pos))
+            else:
+                # If no end marker, assume the section runs to the end of the sequence
+                assistant_sections.append((start_pos, seq_len))
+
+        # Set labels for all identified assistant sections
+        for start_pos, end_pos in assistant_sections:
+            if start_pos < end_pos and start_pos < seq_len:
+                end_pos = min(end_pos, seq_len)  # Safety check
+                labels[i, start_pos:end_pos] = input_ids[start_pos:end_pos]
+
+    # Also mask padding tokens
+    labels[batch["input_ids"] == processor.tokenizer.pad_token_id] = -100
+
+    # Add labels to batch
+    batch["labels"] = labels
+
+    # Add answer to batch
+    if return_answer_in_batch:
+        batch["answer"] = [example["answer"].strip() for example in examples]
+
+    return batch
+
+
+def dna_collate_fn(
+    batch: List[Dict[str, Any]],
+    dna_tokenizer: Any,
+    label2id: Dict[str, int],
+    max_length: int = 2048,
+) -> Dict[str, Any]:
+    """
+    Custom collate function for DNA models.
+    """
+    ref_sequences = [item["reference_sequence"] for item in batch]
+    alt_sequences = [item["variant_sequence"] for item in batch]
+
+    # Tokenize DNA sequences separately
+    tokenized_ref = dna_tokenizer(
+        ref_sequences,
+        padding=True,
+        truncation=True,
+        max_length=max_length,
+        return_tensors="pt",
+    )
+
+    tokenized_alt = dna_tokenizer(
+        alt_sequences,
+        padding=True,
+        truncation=True,
+        max_length=max_length,
+        return_tensors="pt",
+    )
+
+    # Get labels
+    labels = []
+    for item in batch:
+        label = label2id[item["answer"]]
+        labels.append(label)
+
+    # Create labels tensor
+    labels_tensor = torch.tensor(labels, dtype=torch.long)
+
+    tokenized_batch = {
+        "ref_ids": tokenized_ref.input_ids,
+        "ref_attention_mask": tokenized_ref.attention_mask,
+        "alt_ids": tokenized_alt.input_ids,
+        "alt_attention_mask": tokenized_alt.attention_mask,
+        "labels": labels_tensor,
+    }
+
+    return tokenized_batch

BioReason/bioreason/dataset/utils.py

@@ -0,0 +1,59 @@
+from datasets import Dataset as HFDataset
+from torch.utils.data import Dataset as TorchDataset
+from typing import Dict, Any, Union, List
+
+
+def truncate_dna(
+    example: Dict[str, Any], truncate_dna_per_side: int = 1024
+) -> Dict[str, Any]:
+    """
+    Truncate DNA sequences by removing a specified number of base pairs from both ends.
+    If the sequence is too short, it will return the middle portion.
+    """
+    for key in ["reference_sequence", "variant_sequence"]:
+        sequence = example[key]
+        seq_len = len(sequence)
+
+        if seq_len > 2 * truncate_dna_per_side + 8:
+            example[key] = sequence[truncate_dna_per_side:-truncate_dna_per_side]
+
+    return example
+
+
+def torch_to_hf_dataset(torch_dataset: TorchDataset) -> HFDataset:
+    """
+    Convert a PyTorch Dataset to a Hugging Face Dataset.
+
+    This function takes a PyTorch Dataset and converts it to a Hugging Face Dataset
+    by extracting all items and organizing them into a dictionary structure that
+    can be used to create a Hugging Face Dataset.
+
+    Args:
+        torch_dataset: A PyTorch Dataset object to be converted
+
+    Returns:
+        A Hugging Face Dataset containing the same data as the input PyTorch Dataset
+    """
+    # Get first item to determine structure
+    if len(torch_dataset) == 0:
+        return HFDataset.from_dict({})
+
+    first_item = torch_dataset[0]
+
+    # Initialize dictionary based on first item's keys
+    data_dict = (
+        {k: [] for k in first_item.keys()}
+        if isinstance(first_item, dict)
+        else {"data": []}
+    )
+
+    # Populate dictionary
+    for i in range(len(torch_dataset)):
+        item = torch_dataset[i]
+        if isinstance(item, dict):
+            for k in data_dict:
+                data_dict[k].append(item[k])
+        else:
+            data_dict["data"].append(item)
+
+    return HFDataset.from_dict(data_dict)

BioReason/bioreason/dataset/variant_effect.py

@@ -0,0 +1,98 @@
+import json
+import os
+import random
+import sys
+import torch
+from torch.utils.data import Dataset, DataLoader
+from typing import Any, Dict, List, Tuple
+
+from bioreason.dataset.utils import torch_to_hf_dataset
+from bioreason.models.dl.processing_dl import DLProcessor
+from trl.data_utils import maybe_apply_chat_template
+
+
+def get_format_variant_effect_function(model_name: str) -> Any:
+    """
+    Get the appropriate format function for a given model name.
+    """
+    if model_name.lower() == "llm":
+        return format_variant_effect_for_llm
+    elif model_name.lower() == "dna-llm":
+        return format_variant_effect_for_dna_llm
+    else:
+        raise ValueError(f"Unsupported model name: {model_name}")
+    
+
+def clean_variant_effect_example(example: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Clean a variant effect example.
+    """
+    example['answer'] = example['answer'].split(";")[0].strip().lower()
+    return example
+
+
+def clean_variant_effect_non_snv_example(example: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Clean a variant effect non-SNV example.
+    """
+    example['answer'] = example['answer'].replace("[", "").replace("]", "").replace("'", "").replace("_", " ").strip()
+    return example
+
+
+def format_variant_effect_for_dna_llm(example: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Format a VEP example into the required chat format for DNA-LLM.
+    """
+    return {
+        "prompt": [
+            {
+                "role": "user",
+                "content": [
+                    *({"type": "dna", "text": None} for _ in range(2)),
+                    {"type": "text", "text": example["question"].strip()},
+                ],
+            },
+            {
+                "role": "assistant",
+                "reasoning_content": f"Answer: {example['answer'].strip()}",
+                "content": [
+                    {"type": "text", "text": f"Answer: {example['answer'].strip()}"},
+                ],
+            },
+        ],
+        "dna_sequences": [
+            example["reference_sequence"],
+            example["variant_sequence"],
+        ],
+        "answer": example["answer"].strip(),
+    }
+
+
+def format_variant_effect_for_llm(example: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Format a VEP example into the required chat format for LLM.
+    """
+    question = f"Reference sequence: {example['reference_sequence']}\nVariant sequence: {example['variant_sequence']}\nQuestion: {example['question']}"
+    return {
+        "prompt": [
+            {
+                "role": "user",
+                "content": [
+                    *({"type": "dna", "text": None} for _ in range(2)),
+                    {"type": "text", "text": question.strip()},
+                ],
+            },
+            {
+                "role": "assistant",
+                "reasoning_content": f"Answer: {example['answer'].strip()}",
+                "content": [
+                    {"type": "text", "text": f"Answer: {example['answer'].strip()}"},
+                ],
+            },
+        ],
+        "dna_sequences": [
+            "",
+            "",
+        ],
+        "answer": example["answer"].strip(),
+    }

BioReason/bioreason/dna_modules/__init__.py

@@ -0,0 +1,4 @@
+from .dna_module import DNABaseModule
+from .nucleotide_module import NucleotideDNAModule
+
+__all__ = ["DNABaseModule", "NucleotideDNAModule"]

BioReason/bioreason/dna_modules/dna_module.py

@@ -0,0 +1,49 @@
+from abc import ABC, abstractmethod
+from typing import Dict, Any, Union
+import torch
+
+class DNABaseModule(ABC):
+    def __init__(self):
+        super().__init__()
+    
+    @abstractmethod
+    def get_dnallm_key(self):
+        pass
+
+    @abstractmethod
+    def get_model_class(self, model_id: str, model_init_kwargs: dict):
+        pass
+
+    def post_model_init(self, model, processing_class):
+        pass
+
+    def is_embeds_input(self):
+        return False
+    
+    @abstractmethod
+    def get_processing_class(self):
+        pass
+
+    @abstractmethod
+    def get_dnallm_modules_keywords(self):
+        pass
+
+    @abstractmethod
+    def get_custom_multimodal_keywords(self):
+        pass
+
+    @abstractmethod
+    def get_non_generate_params(self):
+        pass
+
+    @abstractmethod
+    def get_custom_processing_keywords(self):
+        pass
+
+    @abstractmethod
+    def prepare_prompt(self, processing_class, inputs: dict[str, Union[torch.Tensor, Any]]):
+        pass
+    
+    @abstractmethod
+    def prepare_model_inputs(self, processing_class, prompts_text, images, return_tensors, padding, padding_side, add_special_tokens):
+        pass

BioReason/bioreason/dna_modules/esm_protein_module.py

@@ -0,0 +1,649 @@
+from typing import Dict, Any, Union, List, Optional, Callable, Type
+from trl.data_utils import maybe_apply_chat_template
+import torch
+import re
+import os
+from datetime import datetime
+
+from bioreason.protein_modules.protein_base_module import ProteinBaseModule
+from bioreason.models.protein_llm.protein_llm import ProteinLLMModel
+from bioreason.models.protein_llm.processing_protein import ProteinLLMProcessor
+
+
+class ESMQFormerModule(ProteinBaseModule):
+    """
+    Protein module implementation for ESM2 + Q-Former based models.
+
+    This module provides the interface between Protein-LLM models and the training
+    infrastructure, handling model loading, processing setup, and reward functions
+    specifically for the ESM2 + Q-Former + LLM architecture.
+    """
+
+    def __init__(self):
+        """Initialize the ESMQFormerModule."""
+        super().__init__()
+
+    def get_protein_llm_key(self) -> str:
+        """
+        Get the key identifier for this protein-LLM implementation.
+
+        Returns:
+            String identifier for this module type
+        """
+        return "esm_qformer"
+
+    def get_model_class(self, model_id: str, model_init_kwargs: Dict[str, Any]) -> Type:
+        """
+        Return the appropriate model class based on model ID.
+
+        Args:
+            model_id: Identifier for the model
+            model_init_kwargs: Initialization arguments for the model
+
+        Returns:
+            The model class to instantiate
+
+        Raises:
+            ValueError: If the model is not supported
+        """
+        if "ProteinLLM" in model_id or "ESM" in model_id:
+            return ProteinLLMModel
+        else:
+            raise ValueError(f"Unsupported model for ESM-QFormer module: {model_id}")
+
+    def post_model_init(self, model: Any, processing_class: Any) -> None:
+        """
+        Perform any post-initialization setup on the model.
+
+        Args:
+            model: The initialized model
+            processing_class: The processor for the model
+        """
+        # Ensure special tokens are properly set
+        if hasattr(model, 'text_tokenizer') and hasattr(model.text_tokenizer, 'pad_token_id'):
+            if model.text_tokenizer.pad_token_id is None:
+                model.text_tokenizer.pad_token_id = model.text_tokenizer.eos_token_id
+        
+        # Set up Q-Former query token count for batch processing
+        if hasattr(processing_class, 'set_query_tokens'):
+            processing_class.set_query_tokens(model.qformer_num_query_tokens)
+
+    def get_processing_class(self) -> Type:
+        """
+        Get the processing class to use with this protein-LLM model.
+
+        Returns:
+            The processing class
+        """
+        return ProteinLLMProcessor
+
+    def get_protein_llm_modules_keywords(self) -> List[str]:
+        """
+        Get keywords to identify protein-specific modules in the model.
+
+        Used to exclude protein modules from LoRA adaptation during training.
+
+        Returns:
+            List of keywords that identify protein modules
+        """
+        return [
+            "protein",
+            "esm",
+            "qformer",
+            "protein_projection",
+            "query_tokens",
+            "protein_model",
+        ]
+
+    def get_custom_multimodal_keywords(self) -> List[str]:
+        """
+        Get keywords for multimodal inputs that should be passed to the model.
+
+        Returns:
+            List of input keywords for multimodal processing
+        """
+        return ["protein_tokenized", "batch_idx_map"]
+
+    def get_non_generate_params(self) -> List[str]:
+        """
+        Get parameter names that should be excluded from generation.
+
+        Returns:
+            List of parameter names to exclude from generation calls
+        """
+        return ["labels"]
+
+    def get_custom_processing_keywords(self) -> List[tuple]:
+        """
+        Get custom processing keywords for the processor.
+
+        Returns:
+            List of (component, parameter) tuples for custom processing
+        """
+        return [
+            ("protein_tokenizer", "max_length_protein"),
+            ("text_tokenizer", "max_length_text"),
+            ("qformer", "num_query_tokens"),
+        ]
+
+    def prepare_prompt(
+        self, 
+        processing_class: Any, 
+        inputs: List[Dict[str, Union[torch.Tensor, Any]]]
+    ) -> List[str]:
+        """
+        Prepare prompts from input examples.
+
+        Args:
+            processing_class: The processor to use
+            inputs: List of input examples
+
+        Returns:
+            List of prepared prompts
+        """
+        prompts_text = []
+        for example in inputs:
+            # Apply chat template if available
+            if hasattr(processing_class, 'tokenizer') and hasattr(processing_class.tokenizer, 'chat_template'):
+                formatted_prompt = maybe_apply_chat_template(example, processing_class)["prompt"]
+            else:
+                # Fallback to simple formatting
+                formatted_prompt = example.get("prompt", example.get("text", str(example)))
+            prompts_text.append(formatted_prompt)
+        
+        return prompts_text
+
+    def prepare_model_inputs(
+        self,
+        processing_class: Any,
+        model: Any,
+        prompts_text: List[str],
+        batch_protein_sequences: List[List[str]],
+        return_tensors: str = "pt",
+        padding: bool = True,
+        padding_side: str = "left",
+        add_special_tokens: bool = False,
+        **kwargs
+    ) -> Dict[str, Any]:
+        """
+        Prepare inputs for the model.
+
+        Args:
+            processing_class: The processor to use
+            model: The model to prepare inputs for
+            prompts_text: List of text prompts
+            batch_protein_sequences: List of lists of protein sequences
+            return_tensors: Return format for tensors
+            padding: Whether to pad inputs
+            padding_side: Side to pad on
+            add_special_tokens: Whether to add special tokens
+            **kwargs: Additional arguments
+
+        Returns:
+            Processed inputs for the model
+        """
+        # Handle DataParallel wrapped models
+        actual_model = model.module if hasattr(model, 'module') else model
+        
+        # Get model parameters
+        max_length_text = getattr(actual_model, 'max_length_text', 512)
+        max_length_protein = getattr(actual_model, 'max_length_protein', 1024)
+        
+        # Process inputs using the processor
+        prompt_inputs = processing_class(
+            text=prompts_text,
+            proteins=batch_protein_sequences,
+            return_tensors=return_tensors,
+            padding=padding,
+            padding_side=padding_side,
+            add_special_tokens=add_special_tokens,
+            max_length_text=max_length_text,
+            max_length_protein=max_length_protein,
+            **kwargs
+        )
+
+        return prompt_inputs
+
+    @staticmethod
+    def get_question_template() -> str:
+        """
+        Get the template for formatting questions.
+
+        Returns:
+            String template for questions
+        """
+        return "Question: {Question}\nAnswer:"
+
+    @staticmethod
+    def accuracy_reward(
+        completions: List[Dict[str, Any]], 
+        ground_truth: List[str], 
+        **kwargs
+    ) -> List[float]:
+        """
+        Accuracy reward function for protein tasks.
+        
+        Compares model outputs with ground truth using exact match and semantic similarity.
+
+        Args:
+            completions: List of model completions
+            ground_truth: List of ground truth answers
+            **kwargs: Additional arguments including:
+                - match_type: 'exact', 'partial', 'semantic' (default: 'partial')
+                - case_sensitive: bool (default: False)
+
+        Returns:
+            List of reward scores (0.0 to 1.0)
+        """
+        completion_contents = [completion[0]["content"] for completion in completions]
+        match_type = kwargs.get('match_type', 'partial')
+        case_sensitive = kwargs.get('case_sensitive', False)
+        rewards = []
+        
+        for content, truth in zip(completion_contents, ground_truth):
+            if not case_sensitive:
+                content = content.lower().strip()
+                truth = truth.lower().strip()
+            else:
+                content = content.strip()
+                truth = truth.strip()
+            
+            if match_type == 'exact':
+                # Exact string match
+                reward = 1.0 if content == truth else 0.0
+            elif match_type == 'partial':
+                # Partial matching using token overlap
+                content_tokens = set(content.split())
+                truth_tokens = set(truth.split())
+                if len(truth_tokens) == 0:
+                    reward = 0.0
+                else:
+                    overlap = len(content_tokens & truth_tokens)
+                    reward = overlap / len(truth_tokens)
+            elif match_type == 'semantic':
+                # Simple semantic matching using key biological terms
+                bio_terms_content = set(re.findall(r'\b(?:protein|enzyme|binding|domain|function|structure|fold|helix|sheet|active|site|residue|amino|acid)\b', content.lower()))
+                bio_terms_truth = set(re.findall(r'\b(?:protein|enzyme|binding|domain|function|structure|fold|helix|sheet|active|site|residue|amino|acid)\b', truth.lower()))
+                
+                if len(bio_terms_truth) == 0:
+                    # Fallback to partial matching
+                    content_tokens = set(content.split())
+                    truth_tokens = set(truth.split())
+                    overlap = len(content_tokens & truth_tokens)
+                    reward = overlap / len(truth_tokens) if len(truth_tokens) > 0 else 0.0
+                else:
+                    overlap = len(bio_terms_content & bio_terms_truth)
+                    reward = overlap / len(bio_terms_truth)
+            else:
+                reward = 0.0
+            
+            rewards.append(min(1.0, max(0.0, reward)))  # Clamp between 0 and 1
+        
+        return rewards
+
+    @staticmethod
+    def repetition_reward(
+        completions: List[Dict[str, Any]], 
+        **kwargs
+    ) -> List[float]:
+        """
+        Repetition penalty reward function.
+        
+        Penalizes outputs with excessive repetition of words, phrases, or patterns.
+
+        Args:
+            completions: List of model completions
+            **kwargs: Additional arguments including:
+                - window_size: Size of n-gram window to check (default: 3)
+                - repetition_threshold: Maximum allowed repetition ratio (default: 0.3)
+                - penalty_weight: Weight of penalty applied (default: 1.0)
+
+        Returns:
+            List of reward scores (0.0 to 1.0, lower for more repetitive text)
+        """
+        completion_contents = [completion[0]["content"] for completion in completions]
+        window_size = kwargs.get('window_size', 3)
+        repetition_threshold = kwargs.get('repetition_threshold', 0.3)
+        penalty_weight = kwargs.get('penalty_weight', 1.0)
+        rewards = []
+        
+        for content in completion_contents:
+            if not content.strip():
+                rewards.append(0.0)
+                continue
+            
+            # Tokenize content
+            tokens = content.lower().split()
+            if len(tokens) < window_size:
+                rewards.append(1.0)  # Too short to have meaningful repetition
+                continue
+            
+            # Calculate n-gram repetition
+            ngrams = []
+            for i in range(len(tokens) - window_size + 1):
+                ngram = ' '.join(tokens[i:i + window_size])
+                ngrams.append(ngram)
+            
+            if len(ngrams) == 0:
+                rewards.append(1.0)
+                continue
+            
+            # Count unique vs total n-grams
+            unique_ngrams = len(set(ngrams))
+            total_ngrams = len(ngrams)
+            repetition_ratio = 1.0 - (unique_ngrams / total_ngrams)
+            
+            # Calculate word-level repetition
+            unique_words = len(set(tokens))
+            total_words = len(tokens)
+            word_repetition_ratio = 1.0 - (unique_words / total_words)
+            
+            # Combine repetition metrics
+            combined_repetition = (repetition_ratio + word_repetition_ratio) / 2
+            
+            # Apply penalty
+            if combined_repetition > repetition_threshold:
+                penalty = (combined_repetition - repetition_threshold) * penalty_weight
+                reward = max(0.0, 1.0 - penalty)
+            else:
+                reward = 1.0
+            
+            rewards.append(reward)
+        
+        # Log repetition analysis if in debug mode
+        if os.getenv("DEBUG_MODE") == "true":
+            current_time = datetime.now().strftime("%d-%H-%M-%S-%f")
+            log_path = os.getenv("LOG_PATH", "debug.log")
+            with open(
+                log_path.replace(".txt", "_repetition.txt"), "a", encoding="utf-8"
+            ) as f:
+                f.write(f"------------- {current_time} Repetition Reward -------------\n")
+                for content, reward in zip(completion_contents, rewards):
+                    f.write(f"Content: {content[:200]}...\n")
+                    f.write(f"Repetition reward: {reward:.3f}\n")
+        
+        return rewards
+
+    @staticmethod
+    def format_accuracy_reward(
+        completions: List[Dict[str, Any]], 
+        **kwargs
+    ) -> List[float]:
+        """
+        Format accuracy reward function for protein analysis outputs.
+        
+        Checks if the model output follows the expected format and structure.
+
+        Args:
+            completions: List of model completions
+            **kwargs: Additional arguments including:
+                - required_format: Format pattern to match (default: structured analysis)
+                - strict_mode: Whether to require exact format compliance (default: False)
+                - format_elements: List of required elements to check for
+
+        Returns:
+            List of reward scores (0.0 to 1.0 based on format compliance)
+        """
+        completion_contents = [completion[0]["content"] for completion in completions]
+        required_format = kwargs.get('required_format', 'structured_analysis')
+        strict_mode = kwargs.get('strict_mode', False)
+        format_elements = kwargs.get('format_elements', ['analysis', 'answer'])
+        rewards = []
+        
+        for content in completion_contents:
+            if not content.strip():
+                rewards.append(0.0)
+                continue
+            
+            score = 0.0
+            total_checks = 0
+            
+            if required_format == 'structured_analysis':
+                # Check for structured analysis format
+                patterns = {
+                    'analysis_section': r'<analysis>.*?</analysis>',
+                    'answer_section': r'<answer>.*?</answer>',
+                    'proper_nesting': r'<analysis>.*?</analysis>\s*<answer>.*?</answer>'
+                }
+                
+                for pattern_name, pattern in patterns.items():
+                    total_checks += 1
+                    if re.search(pattern, content, re.DOTALL | re.IGNORECASE):
+                        score += 1
+                        
+            elif required_format == 'qa_format':
+                # Check for Q&A format
+                patterns = {
+                    'question_indicator': r'(?:question|q):\s*',
+                    'answer_indicator': r'(?:answer|a):\s*',
+                    'proper_structure': r'(?:question|q):\s*.*?(?:answer|a):\s*'
+                }
+                
+                for pattern_name, pattern in patterns.items():
+                    total_checks += 1
+                    if re.search(pattern, content, re.DOTALL | re.IGNORECASE):
+                        score += 1
+                        
+            elif required_format == 'json_format':
+                # Check for JSON-like format
+                try:
+                    import json
+                    # Try to find JSON-like structures
+                    json_pattern = r'\{[^{}]*\}'
+                    json_matches = re.findall(json_pattern, content)
+                    total_checks = 1
+                    if json_matches:
+                        score = 1
+                except:
+                    total_checks = 1
+                    score = 0
+                    
+            elif required_format == 'custom_elements':
+                # Check for custom format elements
+                total_checks = len(format_elements)
+                for element in format_elements:
+                    if element.lower() in content.lower():
+                        score += 1
+            
+            # Additional general format checks
+            general_checks = {
+                'proper_capitalization': bool(re.search(r'[A-Z]', content)),
+                'proper_punctuation': bool(re.search(r'[.!?]', content)),
+                'reasonable_length': len(content.split()) >= 5,
+                'no_excessive_whitespace': not bool(re.search(r'\s{5,}', content))
+            }
+            
+            for check_name, check_result in general_checks.items():
+                total_checks += 1
+                if check_result:
+                    score += 1
+            
+            # Calculate final reward
+            if total_checks > 0:
+                reward = score / total_checks
+                if strict_mode and reward < 1.0:
+                    reward = 0.0  # All or nothing in strict mode
+            else:
+                reward = 0.0
+            
+            rewards.append(reward)
+        
+        # Log format analysis if in debug mode
+        if os.getenv("DEBUG_MODE") == "true":
+            current_time = datetime.now().strftime("%d-%H-%M-%S-%f")
+            log_path = os.getenv("LOG_PATH", "debug.log")
+            with open(
+                log_path.replace(".txt", "_format_accuracy.txt"), "a", encoding="utf-8"
+            ) as f:
+                f.write(f"------------- {current_time} Format Accuracy Reward -------------\n")
+                for content, reward in zip(completion_contents, rewards):
+                    f.write(f"Content: {content[:200]}...\n")
+                    f.write(f"Format accuracy: {reward:.3f}\n")
+                    f.write(f"Required format: {required_format}\n")
+        
+        return rewards
+
+    def get_reward_functions(self) -> Dict[str, callable]:
+        """
+        Get available reward functions for protein tasks.
+        
+        Returns:
+            Dictionary mapping function names to callables
+        """
+        return {
+            "accuracy": self.accuracy_reward,
+            "repetition": self.repetition_reward,
+            "format_accuracy": self.format_accuracy_reward,
+        }
+
+    @staticmethod
+    def select_reward_func(func: str, task_type: str = None, **kwargs) -> Callable:
+        """
+        Select the appropriate reward function based on function name.
+
+        Args:
+            func: The type of reward function ('accuracy', 'repetition', 'format_accuracy')
+            task_type: The type of task (optional, for compatibility)
+            **kwargs: Additional arguments to pass to the reward function
+
+        Returns:
+            The reward function to use
+
+        Raises:
+            ValueError: If the function is not supported
+        """
+        module = ESMQFormerModule()
+        reward_funcs = module.get_reward_functions()
+        
+        if func in reward_funcs:
+            reward_func = reward_funcs[func]
+            
+            # Create a wrapper that includes the kwargs
+            def wrapped_reward_func(completions, ground_truth=None, **additional_kwargs):
+                # Merge kwargs with additional_kwargs
+                merged_kwargs = {**kwargs, **additional_kwargs}
+                
+                # Handle functions that don't need ground_truth
+                if func == 'repetition':
+                    return reward_func(completions, **merged_kwargs)
+                elif func == 'format_accuracy':
+                    return reward_func(completions, **merged_kwargs)
+                else:  # accuracy
+                    if ground_truth is None:
+                        raise ValueError(f"Ground truth required for {func} reward function")
+                    return reward_func(completions, ground_truth, **merged_kwargs)
+            
+            return wrapped_reward_func
+        else:
+            raise ValueError(f"Unsupported reward function: {func}. Available functions: {list(reward_funcs.keys())}")
+
+    @staticmethod
+    def combine_rewards(
+        completions: List[Dict[str, Any]], 
+        ground_truth: List[str] = None,
+        reward_weights: Dict[str, float] = None,
+        **kwargs
+    ) -> List[float]:
+        """
+        Combine multiple reward functions with specified weights.
+        
+        Args:
+            completions: List of model completions
+            ground_truth: List of ground truth answers (optional)
+            reward_weights: Dictionary of reward function names and their weights
+                          Default: {'accuracy': 0.5, 'repetition': 0.3, 'format_accuracy': 0.2}
+            **kwargs: Additional arguments for individual reward functions
+            
+        Returns:
+            List of combined reward scores
+        """
+        if reward_weights is None:
+            reward_weights = {
+                'accuracy': 0.5,
+                'repetition': 0.3,
+                'format_accuracy': 0.2
+            }
+        
+        # Normalize weights
+        total_weight = sum(reward_weights.values())
+        if total_weight > 0:
+            reward_weights = {k: v / total_weight for k, v in reward_weights.items()}
+        
+        combined_rewards = [0.0] * len(completions)
+        module = ESMQFormerModule()
+        
+        for func_name, weight in reward_weights.items():
+            if weight <= 0:
+                continue
+                
+            try:
+                if func_name == 'accuracy' and ground_truth is not None:
+                    rewards = module.accuracy_reward(completions, ground_truth, **kwargs)
+                elif func_name == 'repetition':
+                    rewards = module.repetition_reward(completions, **kwargs)
+                elif func_name == 'format_accuracy':
+                    rewards = module.format_accuracy_reward(completions, **kwargs)
+                else:
+                    continue  # Skip unknown or incompatible functions
+                
+                # Add weighted rewards
+                for i, reward in enumerate(rewards):
+                    combined_rewards[i] += weight * reward
+                    
+            except Exception as e:
+                print(f"Warning: Error calculating {func_name} reward: {e}")
+                continue
+        
+        return combined_rewards
+
+    def validate_model_config(self, config: Dict[str, Any]) -> bool:
+        """
+        Validate model configuration parameters for ESM+Q-Former architecture.
+        
+        Args:
+            config: Configuration dictionary
+            
+        Returns:
+            True if valid, False otherwise
+        """
+        required_keys = [
+            "text_model_name",
+            "protein_model_name",
+        ]
+        
+        for key in required_keys:
+            if key not in config:
+                print(f"Missing required config key: {key}")
+                return False
+        
+        # Validate Q-Former configuration
+        if "qformer_num_query_tokens" in config:
+            if not isinstance(config["qformer_num_query_tokens"], int) or config["qformer_num_query_tokens"] <= 0:
+                print("qformer_num_query_tokens must be a positive integer")
+                return False
+        
+        # Validate max lengths
+        for length_key in ["max_length_protein", "max_length_text"]:
+            if length_key in config:
+                if not isinstance(config[length_key], int) or config[length_key] <= 0:
+                    print(f"{length_key} must be a positive integer")
+                    return False
+        
+        return True
+
+    def get_default_generation_config(self) -> Dict[str, Any]:
+        """
+        Get default generation configuration for protein-LLM models.
+        
+        Returns:
+            Dictionary of default generation parameters
+        """
+        return {
+            "max_new_tokens": 512,
+            "temperature": 0.7,
+            "do_sample": True,
+            "top_p": 0.9,
+            "repetition_penalty": 1.1,
+            "pad_token_id": None,  # Will be set by processor
+            "eos_token_id": None,  # Will be set by processor
+        }

BioReason/bioreason/dna_modules/nucleotide_module.py

@@ -0,0 +1,263 @@
+from transformers import (
+    Qwen2_5_VLForConditionalGeneration,
+    Qwen2VLForConditionalGeneration,
+    AutoProcessor,
+)
+from typing import Dict, Any, Union, List, Optional, Callable, Type
+from trl.data_utils import maybe_apply_chat_template
+from trl import SFTTrainer
+import torch
+
+from bioreason.dna_modules.dna_module import DNABaseModule
+from bioreason.models.dna_llm import DNALLMModel
+from bioreason.models.dl.processing_dl import DLProcessor
+
+
+class NucleotideDNAModule(DNABaseModule):
+    """
+    DNA module implementation for NucleotideTransformer-based models.
+
+    This module provides the interface between DNA-LLM models and the training
+    infrastructure, handling model loading, processing setup, and reward functions.
+    """
+
+    def __init__(self):
+        """Initialize the NucleotideDNAModule."""
+        super().__init__()
+
+    def get_dnallm_key(self) -> str:
+        """
+        Get the key identifier for this DNA-LLM implementation.
+
+        Returns:
+            String identifier for this module type
+        """
+        return "qwen"
+
+    def get_model_class(self, model_id: str, model_init_kwargs: Dict[str, Any]) -> Type:
+        """
+        Return the appropriate model class based on model ID.
+
+        Args:
+            model_id: Identifier for the model
+            model_init_kwargs: Initialization arguments for the model
+
+        Returns:
+            The model class to instantiate
+
+        Raises:
+            ValueError: If the model is not supported
+        """
+        if "DNALLM" in model_id:
+            model_cls = DNALLMModel
+        else:
+            raise ValueError(f"Unsupported model: {model_id}")
+        return model_cls
+
+    def post_model_init(self, model: Any, processing_class: Any) -> None:
+        """
+        Perform any post-initialization setup on the model.
+
+        Args:
+            model: The initialized model
+            processing_class: The processor for the model
+        """
+        # No post-init needed for this implementation
+        pass
+
+    def get_processing_class(self) -> Type:
+        """
+        Get the processing class to use with this DNA-LLM model.
+
+        Returns:
+            The processing class
+        """
+        return DLProcessor
+
+    def get_dnallm_modules_keywords(self) -> List[str]:
+        """
+        Get keywords to identify DNA-specific modules in the model.
+
+        Used to exclude DNA modules from LoRA adaptation during training.
+
+        Returns:
+            List of keywords that identify DNA modules
+        """
+        return ["dna"]
+
+    def get_custom_multimodal_keywords(self) -> List[str]:
+        """
+        Get keywords for multimodal inputs that should be passed to the model.
+
+        Returns:
+            List of input keywords for multimodal processing
+        """
+        return ["dna_tokenized", "batch_idx_map"]
+
+    def get_non_generate_params(self) -> List[str]:
+        """
+        Get parameter names that should be excluded from generation.
+
+        Returns:
+            List of parameter names to exclude from generation calls
+        """
+        return []
+
+    def get_custom_processing_keywords(self) -> List[tuple]:
+        """
+        Get custom processing keywords for the processor.
+
+        Returns:
+            List of (component, parameter) tuples for custom processing
+        """
+        return [("dna_tokenizer", "max_length")]
+
+    def prepare_prompt(
+        self, processing_class: Any, inputs: List[Dict[str, Union[torch.Tensor, Any]]]
+    ) -> List[str]:
+        """
+        Prepare prompts from input examples.
+
+        Args:
+            processing_class: The processor to use
+            inputs: List of input examples
+
+        Returns:
+            List of prepared prompts
+        """
+        prompts_text = [
+            maybe_apply_chat_template(example, processing_class)["prompt"]
+            for example in inputs
+        ]
+        return prompts_text
+
+    def prepare_model_inputs(
+        self,
+        processing_class: Any,
+        model: Any,
+        prompts_text: List[str],
+        batch_dna_sequences: List[List[str]],
+        return_tensors: str = "pt",
+        padding: bool = True,
+        padding_side: str = "left",
+        add_special_tokens: bool = False,
+    ) -> Dict[str, Any]:
+        """
+        Prepare inputs for the model.
+
+        Args:
+            processing_class: The processor to use
+            model: The model to prepare inputs for
+            prompts_text: List of text prompts
+            batch_dna_sequences: List of lists of DNA sequences
+            return_tensors: Return format for tensors
+            padding: Whether to pad inputs
+            padding_side: Side to pad on
+            add_special_tokens: Whether to add special tokens
+
+        Returns:
+            Processed inputs for the model
+        """
+        # Handle DataParallel wrapped models by accessing the module attribute if needed
+        max_length_text = model.max_length_text if not hasattr(model, 'module') else model.module.max_length_text
+        max_length_dna = model.max_length_dna if not hasattr(model, 'module') else model.module.max_length_dna
+        
+        prompt_inputs = processing_class(
+            text=prompts_text,
+            batch_dna_sequences=batch_dna_sequences,
+            return_tensors=return_tensors,
+            padding=padding,
+            padding_side=padding_side,
+            add_special_tokens=add_special_tokens,
+            max_length_text=max_length_text,
+            max_length_dna=max_length_dna,
+        )
+
+        return prompt_inputs
+
+    def is_embeds_input(self) -> bool:
+        """
+        Whether the model uses embeddings as input (instead of token IDs).
+
+        Returns:
+            Boolean indicating if the model takes embedding inputs
+        """
+        return True
+
+    @staticmethod
+    def get_question_template() -> str:
+        """
+        Get the template for formatting questions.
+
+        Returns:
+            String template for questions
+        """
+        return "{Question}"
+
+    @staticmethod
+    def format_reward_rec(completions: List[Dict[str, Any]], **kwargs) -> List[float]:
+        """
+        Check if the Qwen model output matches a specific format.
+
+        Args:
+            completions: List of model completions
+            **kwargs: Additional arguments
+
+        Returns:
+            List of reward scores (1.0 for match, 0.0 for no match)
+        """
+        import re
+        import os
+        from datetime import datetime
+
+        # Pattern to match the expected output format
+        pattern = r"<think>.*?</think>\s*<answer>.*?\{.*\[\d+,\s*\d+,\s*\d+,\s*\d+\].*\}.*?</answer>"
+        completion_contents = [completion[0]["content"] for completion in completions]
+        matches = [
+            re.search(pattern, content, re.DOTALL) is not None
+            for content in completion_contents
+        ]
+
+        # Log format results if in debug mode
+        current_time = datetime.now().strftime("%d-%H-%M-%S-%f")
+        if os.getenv("DEBUG_MODE") == "true":
+            log_path = os.getenv("LOG_PATH")
+            with open(
+                log_path.replace(".txt", "_format.txt"), "a", encoding="utf-8"
+            ) as f:
+                f.write(f"------------- {current_time} Format reward -------------\n")
+                for content, match in zip(completion_contents, matches):
+                    f.write(f"Content: {content}\n")
+                    f.write(f"Has format: {bool(match)}\n")
+
+        return [1.0 if match else 0.0 for match in matches]
+
+    @staticmethod
+    def select_reward_func(func: str, task_type: str) -> Callable:
+        """
+        Select the appropriate reward function based on function name and task type.
+
+        Args:
+            func: The type of reward function ('accuracy', 'format', etc.)
+            task_type: The type of task ('rec', etc.)
+
+        Returns:
+            The reward function to use
+
+        Raises:
+            ValueError: If the function or task type is not supported
+        """
+        if func == "accuracy":
+            match task_type:
+                case "rec":
+                    return NucleotideDNAModule.iou_reward
+                case _:
+                    raise ValueError(f"Unsupported reward function: {func}")
+        elif func == "format":
+            match task_type:
+                case "rec":
+                    return NucleotideDNAModule.format_reward_rec
+                case _:
+                    raise ValueError(f"Unsupported reward function: {func}")
+        else:
+            raise ValueError(f"Unsupported reward function: {func}")

BioReason/bioreason/dna_modules/protein_module.py

@@ -0,0 +1,200 @@
+from abc import ABC, abstractmethod
+from typing import Dict, Any, Union, List, Optional, Type
+import torch
+
+class ProteinBaseModule(ABC):
+    """
+    Abstract base class for protein-language model modules.
+    
+    This class defines the interface that all protein-LLM implementations
+    must follow, providing standardized methods for model loading,
+    processing, and training integration.
+    """
+    
+    def __init__(self):
+        """Initialize the protein module."""
+        super().__init__()
+    
+    @abstractmethod
+    def get_protein_llm_key(self) -> str:
+        """
+        Get the unique identifier for this protein-LLM implementation.
+        
+        Returns:
+            String identifier for this module type
+        """
+        pass
+
+    @abstractmethod
+    def get_model_class(self, model_id: str, model_init_kwargs: Dict[str, Any]) -> Type:
+        """
+        Return the appropriate model class based on model ID.
+        
+        Args:
+            model_id: Identifier for the model
+            model_init_kwargs: Initialization arguments for the model
+            
+        Returns:
+            The model class to instantiate
+        """
+        pass
+
+    def post_model_init(self, model: Any, processing_class: Any) -> None:
+        """
+        Perform any post-initialization setup on the model.
+        
+        Args:
+            model: The initialized model
+            processing_class: The processor for the model
+        """
+        # Default implementation does nothing
+        pass
+
+    def is_embeds_input(self) -> bool:
+        """
+        Whether the model uses embeddings as input (instead of token IDs).
+        
+        Returns:
+            Boolean indicating if the model takes embedding inputs
+        """
+        # Default for protein-LLM models is True due to Q-Former integration
+        return True
+    
+    @abstractmethod
+    def get_processing_class(self) -> Type:
+        """
+        Get the processing class to use with this protein-LLM model.
+        
+        Returns:
+            The processing class
+        """
+        pass
+
+    @abstractmethod
+    def get_protein_llm_modules_keywords(self) -> List[str]:
+        """
+        Get keywords to identify protein-specific modules in the model.
+        
+        Used to exclude protein modules from LoRA adaptation during training
+        or to identify components for specific training strategies.
+        
+        Returns:
+            List of keywords that identify protein modules
+        """
+        pass
+
+    @abstractmethod
+    def get_custom_multimodal_keywords(self) -> List[str]:
+        """
+        Get keywords for multimodal inputs that should be passed to the model.
+        
+        Returns:
+            List of input keywords for multimodal processing
+        """
+        pass
+
+    @abstractmethod
+    def get_non_generate_params(self) -> List[str]:
+        """
+        Get parameter names that should be excluded from generation.
+        
+        Returns:
+            List of parameter names to exclude from generation calls
+        """
+        pass
+
+    @abstractmethod
+    def get_custom_processing_keywords(self) -> List[tuple]:
+        """
+        Get custom processing keywords for the processor.
+        
+        Returns:
+            List of (component, parameter) tuples for custom processing
+        """
+        pass
+
+    @abstractmethod
+    def prepare_prompt(
+        self, 
+        processing_class: Any, 
+        inputs: List[Dict[str, Union[torch.Tensor, Any]]]
+    ) -> List[str]:
+        """
+        Prepare prompts from input examples.
+        
+        Args:
+            processing_class: The processor to use
+            inputs: List of input examples
+            
+        Returns:
+            List of prepared prompts
+        """
+        pass
+    
+    @abstractmethod
+    def prepare_model_inputs(
+        self,
+        processing_class: Any,
+        model: Any,
+        prompts_text: List[str],
+        batch_protein_sequences: List[List[str]],
+        return_tensors: str = "pt",
+        padding: bool = True,
+        padding_side: str = "left",
+        add_special_tokens: bool = False,
+        **kwargs
+    ) -> Dict[str, Any]:
+        """
+        Prepare inputs for the model.
+        
+        Args:
+            processing_class: The processor to use
+            model: The model to prepare inputs for
+            prompts_text: List of text prompts
+            batch_protein_sequences: List of lists of protein sequences
+            return_tensors: Return format for tensors
+            padding: Whether to pad inputs
+            padding_side: Side to pad on
+            add_special_tokens: Whether to add special tokens
+            **kwargs: Additional arguments
+            
+        Returns:
+            Processed inputs for the model
+        """
+        pass
+
+    def get_reward_functions(self) -> Dict[str, callable]:
+        """
+        Get available reward functions for this module.
+        
+        Returns:
+            Dictionary mapping function names to callables
+        """
+        return {}
+
+    def validate_model_config(self, config: Dict[str, Any]) -> bool:
+        """
+        Validate model configuration parameters.
+        
+        Args:
+            config: Configuration dictionary
+            
+        Returns:
+            True if valid, False otherwise
+        """
+        return True
+
+    def get_default_generation_config(self) -> Dict[str, Any]:
+        """
+        Get default generation configuration for this model type.
+        
+        Returns:
+            Dictionary of default generation parameters
+        """
+        return {
+            "max_new_tokens": 512,
+            "temperature": 0.7,
+            "do_sample": True,
+            "top_p": 0.9,
+            "pad_token_id": None,  # Will be set by processor
+        }

BioReason/bioreason/models/__init__.py

@@ -0,0 +1,9 @@
+from .dna_only import DNAClassifierModel
+from .dna_llm import DNALLMModel
+from .evo2_tokenizer import Evo2Tokenizer
+
+__all__ = [
+    "DNAClassifierModel",
+    "DNALLMModel",
+    "Evo2Tokenizer",
+]

BioReason/bioreason/models/dl/__init__.py

@@ -0,0 +1 @@
+

BioReason/bioreason/models/dl/chat_template_dl.py

@@ -0,0 +1 @@
+CHAT_TEMPLATE = "{%- set dna_count = namespace(value=0) %}{%- if tools %}\n    {{- '<|im_start|>system\\n' }}\n    {%- if messages[0].role == 'system' %}\n        {{- messages[0].content + '\\n\\n' }}\n    {%- endif %}\n    {{- \"# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n    {%- for tool in tools %}\n        {{- \"\\n\" }}\n        {{- tool | tojson }}\n    {%- endfor %}\n    {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}\n{%- else %}\n    {%- if messages[0].role == 'system' %}\n        {{- '<|im_start|>system\\n' + messages[0].content + '<|im_end|>\\n' }}\n    {%- endif %}\n{%- endif %}\n{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}\n{%- for message in messages[::-1] %}\n    {%- set index = (messages|length - 1) - loop.index0 %}\n    {%- if ns.multi_step_tool and message.role == \"user\" and not(message.content is string and message.content.startswith('<tool_response>') and message.content.endswith('</tool_response>')) %}\n        {%- set ns.multi_step_tool = false %}\n        {%- set ns.last_query_index = index %}\n    {%- endif %}\n{%- endfor %}\n{%- for message in messages %}\n    {%- if (message.role == \"user\") or (message.role == \"system\" and not loop.first) %}\n        {{- '<|im_start|>' + message.role + '\\n' }} {%- if message.content is string %}{{- message.content + '<|im_end|>' + '\\n' }}{%- else %}{%- for content in message.content %}{%- if content.type == 'dna' or 'dna' in content %}{%- set dna_count.value = dna_count.value + 1 %}{%- if add_dna_id %}DNA Sequence {{- dna_count.value }}: {%- endif %}<|dna_start|><|dna_pad|><|dna_end|>{%- elif 'text' in content %}{{- content.text }}{%- endif %}{%- endfor %}{{- '<|im_end|>' + '\\n' }}{%- endif %}{%- elif message.role == \"assistant\" %}\n        {%- set content = message.content[0].text %}\n        {%- set reasoning_content = '' %}\n        {%- if message.reasoning_content is defined and message.reasoning_content is not none %}\n            {%- set reasoning_content = message.reasoning_content %}\n        {%- else %}\n            {%- if '</think>' in message.content %}\n                {%- set content = message.content[0].text.split('</think>')[-1].lstrip('\\n') %}\n                {%- set reasoning_content = message.content[0].text.split('</think>')[0].rstrip('\\n').split('<think>')[-1].lstrip('\\n') %}\n            {%- endif %}\n        {%- endif %}\n        {%- if loop.index0 > ns.last_query_index %}\n            {%- if loop.last or (not loop.last and reasoning_content) %}\n                {{- '<|im_start|>' + message.role + '\\n<think>\\n' + reasoning_content.strip('\\n') + '\\n</think>\\n\\n' + content.lstrip('\\n') }}\n            {%- else %}\n                {{- '<|im_start|>' + message.role + '\\n' + content }}\n            {%- endif %}\n        {%- else %}\n            {{- '<|im_start|>' + message.role + '\\n' + content }}\n        {%- endif %}\n        {%- if message.tool_calls %}\n            {%- for tool_call in message.tool_calls %}\n                {%- if (loop.first and content) or (not loop.first) %}\n                    {{- '\\n' }}\n                {%- endif %}\n                {%- if tool_call.function %}\n                    {%- set tool_call = tool_call.function %}\n                {%- endif %}\n                {{- '<tool_call>\\n{\"name\": \"' }}\n                {{- tool_call.name }}\n                {{- '\", \"arguments\": ' }}\n                {%- if tool_call.arguments is string %}\n                    {{- tool_call.arguments }}\n                {%- else %}\n                    {{- tool_call.arguments | tojson }}\n                {%- endif %}\n                {{- '}\\n</tool_call>' }}\n            {%- endfor %}\n        {%- endif %}\n        {{- '<|im_end|>\\n' }}\n    {%- elif message.role == \"tool\" %}\n        {%- if loop.first or (messages[loop.index0 - 1].role != \"tool\") %}\n            {{- '<|im_start|>user' }}\n        {%- endif %}\n        {{- '\\n<tool_response>\\n' }}\n        {{- message.content }}\n        {{- '\\n</tool_response>' }}\n        {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n            {{- '<|im_end|>\\n' }}\n        {%- endif %}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {{- '<|im_start|>assistant\\n' }}\n    {%- if enable_thinking is defined and enable_thinking is false %}\n        {{- '<think>\\n\\n</think>\\n\\n' }}\n    {%- endif %}\n{%- endif %}"

BioReason/bioreason/models/dl/configuration_dl.py

@@ -0,0 +1,232 @@
+from transformers import PretrainedConfig
+
+class DLDNAConfig(PretrainedConfig):
+    model_type = "dl"
+    base_config_key = "dna_config"
+
+    def __init__(
+        self,
+        depth=32,
+        hidden_size=3584,
+        hidden_act="silu",
+        intermediate_size=3420,
+        num_heads=16,
+        in_channels=3,
+        patch_size=14,
+        spatial_merge_size=2,
+        temporal_patch_size=2,
+        tokens_per_second=4,
+        window_size=112,
+        out_hidden_size=3584,
+        fullatt_block_indexes=[7, 15, 23, 31],
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.depth = depth
+        self.hidden_size = hidden_size
+        self.hidden_act = hidden_act
+        self.intermediate_size = intermediate_size
+        self.num_heads = num_heads
+        self.in_channels = in_channels
+        self.patch_size = patch_size
+        self.spatial_merge_size = spatial_merge_size
+        self.temporal_patch_size = temporal_patch_size
+        self.tokens_per_second = tokens_per_second
+        self.window_size = window_size
+        self.fullatt_block_indexes = fullatt_block_indexes
+        self.out_hidden_size = out_hidden_size
+
+class DLConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`Qwen2_5_VLModel`]. It is used to instantiate a
+    Qwen2-VL model according to the specified arguments, defining the model architecture. Instantiating a configuration
+    with the defaults will yield a similar configuration to that of
+    Qwen2-VL-7B-Instruct [Qwen/Qwen2-VL-7B-Instruct](https://huggingface.co/Qwen/Qwen2-VL-7B-Instruct).
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+
+    Args:
+        vocab_size (`int`, *optional*, defaults to 152064):
+            Vocabulary size of the Qwen2_5_VL model. Defines the number of different tokens that can be represented by the
+            `inputs_ids` passed when calling [`Qwen2_5_VLModel`]
+        hidden_size (`int`, *optional*, defaults to 8192):
+            Dimension of the hidden representations.
+        intermediate_size (`int`, *optional*, defaults to 29568):
+            Dimension of the MLP representations.
+        num_hidden_layers (`int`, *optional*, defaults to 80):
+            Number of hidden layers in the Transformer encoder.
+        num_attention_heads (`int`, *optional*, defaults to 64):
+            Number of attention heads for each attention layer in the Transformer encoder.
+        num_key_value_heads (`int`, *optional*, defaults to 8):
+            This is the number of key_value heads that should be used to implement Grouped Query Attention. If
+            `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
+            `num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When
+            converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
+            by meanpooling all the original heads within that group. For more details checkout [this
+            paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to `32`.
+        hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
+            The non-linear activation function (function or string) in the decoder.
+        max_position_embeddings (`int`, *optional*, defaults to 32768):
+            The maximum sequence length that this model might ever be used with.
+        initializer_range (`float`, *optional*, defaults to 0.02):
+            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+        rms_norm_eps (`float`, *optional*, defaults to 1e-05):
+            The epsilon used by the rms normalization layers.
+        use_cache (`bool`, *optional*, defaults to `True`):
+            Whether or not the model should return the last key/values attentions (not used by all models). Only
+            relevant if `config.is_decoder=True`.
+        tie_word_embeddings (`bool`, *optional*, defaults to `False`):
+            Whether the model's input and output word embeddings should be tied.
+        rope_theta (`float`, *optional*, defaults to 1000000.0):
+            The base period of the RoPE embeddings.
+        use_sliding_window (`bool`, *optional*, defaults to `False`):
+            Whether to use sliding window attention.
+        sliding_window (`int`, *optional*, defaults to 4096):
+            Sliding window attention (SWA) window size. If not specified, will default to `4096`.
+        max_window_layers (`int`, *optional*, defaults to 80):
+            The number of layers that use SWA (Sliding Window Attention). The bottom layers use SWA while the top use full attention.
+        attention_dropout (`float`, *optional*, defaults to 0.0):
+            The dropout ratio for the attention probabilities.
+        vision_config (`Dict`, *optional*):
+            The config for the visual encoder initialization.
+        rope_scaling (`Dict`, *optional*):
+            Dictionary containing the scaling configuration for the RoPE embeddings. NOTE: if you apply new rope type
+            and you expect the model to work on longer `max_position_embeddings`, we recommend you to update this value
+            accordingly.
+            Expected contents:
+                `rope_type` (`str`):
+                    The sub-variant of RoPE to use. Can be one of ['default', 'linear', 'dynamic', 'yarn', 'longrope',
+                    'llama3'], with 'default' being the original RoPE implementation.
+                `factor` (`float`, *optional*):
+                    Used with all rope types except 'default'. The scaling factor to apply to the RoPE embeddings. In
+                    most scaling types, a `factor` of x will enable the model to handle sequences of length x *
+                    original maximum pre-trained length.
+                `original_max_position_embeddings` (`int`, *optional*):
+                    Used with 'dynamic', 'longrope' and 'llama3'. The original max position embeddings used during
+                    pretraining.
+                `attention_factor` (`float`, *optional*):
+                    Used with 'yarn' and 'longrope'. The scaling factor to be applied on the attention
+                    computation. If unspecified, it defaults to value recommended by the implementation, using the
+                    `factor` field to infer the suggested value.
+                `beta_fast` (`float`, *optional*):
+                    Only used with 'yarn'. Parameter to set the boundary for extrapolation (only) in the linear
+                    ramp function. If unspecified, it defaults to 32.
+                `beta_slow` (`float`, *optional*):
+                    Only used with 'yarn'. Parameter to set the boundary for interpolation (only) in the linear
+                    ramp function. If unspecified, it defaults to 1.
+                `short_factor` (`List[float]`, *optional*):
+                    Only used with 'longrope'. The scaling factor to be applied to short contexts (<
+                    `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
+                    size divided by the number of attention heads divided by 2
+                `long_factor` (`List[float]`, *optional*):
+                    Only used with 'longrope'. The scaling factor to be applied to long contexts (<
+                    `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
+                    size divided by the number of attention heads divided by 2
+                `low_freq_factor` (`float`, *optional*):
+                    Only used with 'llama3'. Scaling factor applied to low frequency components of the RoPE
+                `high_freq_factor` (`float`, *optional*):
+                    Only used with 'llama3'. Scaling factor applied to high frequency components of the RoPE
+
+    ```python
+    >>> from transformers import Qwen2_5_VLForConditionalGeneration, Qwen2_5_VLConfig
+
+    >>> # Initializing a Qwen2_5_VL style configuration
+    >>> configuration = Qwen2_5_VLConfig()
+
+    >>> # Initializing a model from the Qwen2-VL-7B style configuration
+    >>> model = Qwen2_5_VLForConditionalGeneration(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "dl"
+    sub_configs = {"dna_config": DLDNAConfig}
+    keys_to_ignore_at_inference = ["past_key_values"]
+    # Default tensor parallel plan for base model `Qwen2_5_VL`
+    base_model_tp_plan = {
+        "layers.*.self_attn.q_proj": "colwise",
+        "layers.*.self_attn.k_proj": "colwise",
+        "layers.*.self_attn.v_proj": "colwise",
+        "layers.*.self_attn.o_proj": "rowwise",
+        "layers.*.mlp.gate_proj": "colwise",
+        "layers.*.mlp.up_proj": "colwise",
+        "layers.*.mlp.down_proj": "rowwise",
+    }
+    base_model_pp_plan = {
+        "embed_tokens": (["input_ids"], ["inputs_embeds"]),
+        "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
+        "norm": (["hidden_states"], ["hidden_states"]),
+    }
+
+    def __init__(
+        self,
+        vocab_size=152064,
+        hidden_size=8192,
+        intermediate_size=29568,
+        num_hidden_layers=80,
+        num_attention_heads=64,
+        num_key_value_heads=8,
+        hidden_act="silu",
+        max_position_embeddings=32768,
+        initializer_range=0.02,
+        rms_norm_eps=1e-05,
+        use_cache=True,
+        tie_word_embeddings=False,
+        rope_theta=1000000.0,
+        use_sliding_window=False,
+        sliding_window=4096,
+        max_window_layers=80,
+        attention_dropout=0.0,
+        vision_config=None,
+        rope_scaling=None,
+        image_token_id=None,
+        **kwargs,
+    ):
+        if isinstance(vision_config, dict):
+            self.vision_config = self.sub_configs["vision_config"](**vision_config)
+        elif vision_config is None:
+            self.vision_config = self.sub_configs["vision_config"]()
+
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.use_sliding_window = use_sliding_window
+        self.sliding_window = sliding_window
+        self.max_window_layers = max_window_layers
+
+        # for backward compatibility
+        if num_key_value_heads is None:
+            num_key_value_heads = num_attention_heads
+
+        self.num_key_value_heads = num_key_value_heads
+        self.hidden_act = hidden_act
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.use_cache = use_cache
+        self.rope_theta = rope_theta
+        self.attention_dropout = attention_dropout
+        self.rope_scaling = rope_scaling
+
+        self.dna_token_id = image_token_id
+
+        # Validate the correctness of rotary position embeddings parameters
+        # BC: if there is a 'type' field, move it to 'rope_type'.
+        # and change type from 'mrope' to 'default' because `mrope` does default RoPE calculations
+        # one can set it to "linear"/"dynamic" etc. to have scaled RoPE
+        # TODO: @raushan update config in the hub
+        if self.rope_scaling is not None and "type" in self.rope_scaling:
+            if self.rope_scaling["type"] == "mrope":
+                self.rope_scaling["type"] = "default"
+            self.rope_scaling["rope_type"] = self.rope_scaling["type"]
+        rope_config_validation(self, ignore_keys={"mrope_section"})
+
+        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
+
+__all__ = ["DLConfig"]

BioReason/bioreason/models/dl/processing_dl.py

@@ -0,0 +1,275 @@
+from typing import List, Optional, Union, Dict, Any, Tuple
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+from transformers import AutoTokenizer
+from transformers.processing_utils import (
+    CommonKwargs,
+    ProcessingKwargs,
+    ProcessorMixin,
+    Unpack,
+)
+from transformers.feature_extraction_utils import BatchFeature
+from transformers.tokenization_utils_base import PreTokenizedInput, TextInput
+from transformers.utils import logging
+
+from bioreason.utils.dna_utils import DNAInput
+
+class DLDNAKwargs(CommonKwargs):
+    """Keyword arguments specific to DNA processing"""
+    max_length_text: Optional[int]
+    max_length_dna: Optional[int]
+
+
+class DLProcessorKwargs(ProcessingKwargs, total=False):
+    """Processing keyword arguments for the DL processor"""
+    dna_kwargs: DLDNAKwargs
+    _defaults = {
+        "text_kwargs": {
+            "padding": False,
+        },
+    }
+
+class DLProcessor(ProcessorMixin):
+    r"""
+    Constructs a DL processor which wraps a NucleotideTransformer DNA processor and a Qwen2_5 tokenizer into a single processor.
+    This processor handles both text and DNA sequence processing to prepare inputs for the DNALLMModel.
+    
+    Args:
+        tokenizer (PreTrainedTokenizerBase, *optional*):
+            The text tokenizer used for processing text inputs.
+        dna_tokenizer (PreTrainedTokenizerBase, *optional*):
+            The DNA tokenizer used for processing DNA sequences.
+        chat_template (`str`, *optional*): 
+            A Jinja template for chat formatting. If None, will use the tokenizer's template.
+    """
+
+    attributes = ["tokenizer", "dna_tokenizer"]
+    valid_kwargs = ["model", "chat_template"]
+    tokenizer_class = (
+        "Qwen2Tokenizer", "Qwen2TokenizerFast",
+        "GPT2TokenizerFast",
+    )
+    dna_tokenizer_class = ("EsmTokenizer", "Evo2Tokenizer")
+
+    def __init__(
+        self, tokenizer=None, dna_tokenizer=None, chat_template=None, **kwargs
+    ):
+        """
+        Initialize the processor with text and DNA tokenizers.
+        
+        Args:
+            tokenizer: Text tokenizer (usually from a language model)
+            dna_tokenizer: DNA tokenizer (usually from a DNA model)
+            chat_template: Template for formatting chat conversations
+            **kwargs: Additional arguments
+        """
+        self.tokenizer = tokenizer
+        self.dna_tokenizer = dna_tokenizer
+
+        self.dna_token = (
+            "<|dna_pad|>"
+            if not hasattr(self.tokenizer, "dna_token")
+            else self.tokenizer.dna_token
+        )
+    
+        # Get chat template from tokenizer if not provided
+        if chat_template is None and hasattr(self.tokenizer, "chat_template"):
+            chat_template = self.tokenizer.chat_template
+        super().__init__(tokenizer, dna_tokenizer, chat_template=chat_template)
+      
+        # The GRPO trainer might expect this to be set
+        if not hasattr(self.tokenizer, 'pad_token') or self.tokenizer.pad_token is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
+
+    def tokenize_dna_sequences(
+        self, 
+        batch_dna_sequences: List[List[str]], 
+        max_length: int = 2048,
+        return_tensors: str = "pt",
+        device: str = "cuda",
+    ) -> Dict[str, Any]:
+        """
+        Tokenize a batch of DNA sequences.
+        
+        Args:
+            batch_dna_sequences: List of lists of DNA sequences per batch item
+            max_length: Maximum allowed length for DNA sequences
+            return_tensors: Return format for tensors ("pt" for PyTorch)
+            device: Device to place tensors on
+            
+        Returns:
+            Dict containing:
+                - dna_tokenized: The tokenized DNA sequences 
+                - batch_idx_map: Mapping of which sequences belong to which batch item
+        """
+        # Create a mapping to track which sequences belong to which batch item
+        batch_idx_map = []
+        all_sequences = []
+
+        # Flatten all sequences with batch tracking
+        for batch_idx, dna_sequences in enumerate(batch_dna_sequences):
+            for seq in dna_sequences:
+                all_sequences.append(seq)
+                batch_idx_map.append(batch_idx)
+
+        # If no sequences in the entire batch, return empty dict
+        if not all_sequences:
+            return {"dna_tokenized": None, "batch_idx_map": []}
+
+        # Tokenize all sequences at once
+        dna_tokenized = self.dna_tokenizer(
+            all_sequences,
+            padding=True,
+            truncation=True,
+            max_length=max_length,
+            return_tensors=return_tensors,
+            return_attention_mask=True,
+        )
+            
+        return {"dna_tokenized": dna_tokenized, "batch_idx_map": batch_idx_map}
+
+    def __call__(
+        self,
+        batch_dna_sequences: Optional[List[List[str]]] = None,
+        text: Optional[
+            Union[
+                TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]
+            ]
+        ] = None,
+        max_length_text: int = 512,
+        max_length_dna: int = 2048,
+        return_tensors: str = "pt",
+        device: str = "cuda",
+        **kwargs: Unpack[DLProcessorKwargs],
+    ) -> BatchFeature:
+        """
+        Process text and DNA sequences for model input.
+        
+        Args:
+            batch_dna_sequences: List of lists of DNA sequences per batch item
+            text: Input text or list of texts
+            max_length_text: Maximum length for text sequences
+            max_length_dna: Maximum length for DNA sequences
+            return_tensors: Return format for tensors
+            device: Device to place tensors on
+            **kwargs: Additional processor keyword arguments
+            
+        Returns:
+            BatchFeature with tokenized inputs for the model
+        """
+        output_kwargs = self._merge_kwargs(
+            DLProcessorKwargs,
+            tokenizer_init_kwargs=self.tokenizer.init_kwargs,
+            **kwargs,
+        )
+
+        # Ensure text is a list
+        if not isinstance(text, list):
+            text = [text]
+
+        # flattened_dna_sequences = [dna_sequence for dna_sequences in batch_dna_sequences for dna_sequence in dna_sequences]
+        dna_inputs = {}
+        if batch_dna_sequences is not None:
+            # Tokenize DNA sequences
+            dna_processing_result = self.tokenize_dna_sequences(
+                batch_dna_sequences,
+                max_length=max_length_dna,
+                return_tensors=return_tensors,
+                device=device,
+            )
+            
+            # Replace DNA tokens in text if needed
+            index = 0
+            for i in range(len(text)):
+                while self.dna_token in text[i]:
+                    num_dna_tokens = (dna_processing_result['dna_tokenized']['input_ids'][index] != 1).sum().item()
+                    text[i] = text[i].replace(
+                        self.dna_token, "<|placeholder|>" * num_dna_tokens, 1
+                    )
+                    index += 1
+                text[i] = text[i].replace("<|placeholder|>", self.dna_token)
+            
+            
+            
+            # Add batch info to the output
+            dna_inputs = {
+                # "batch_dna_sequences": batch_dna_sequences,
+                "dna_tokenized": dna_processing_result["dna_tokenized"],
+                "batch_idx_map": dna_processing_result["batch_idx_map"],
+            }
+
+        # Tokenize text
+        text_kwargs = output_kwargs.get("text_kwargs", {})
+        
+        if 'padding' in text_kwargs:
+            del text_kwargs['padding']
+        
+        # print("__call__ (processor):", text)
+        text_inputs = self.tokenizer(
+            text, 
+            max_length=max_length_text + 2 * max_length_dna,
+            return_tensors=return_tensors,
+            padding=True,
+            truncation=True,
+            **text_kwargs,
+        )
+        
+        # The BatchFeature should have all required fields for the model's forward pass
+        return BatchFeature(data={**text_inputs, **dna_inputs})
+
+    def batch_decode(self, *args, **kwargs) -> List[str]:
+        """
+        This method forwards all its arguments to the tokenizer's batch_decode.
+        
+        Returns:
+            List of decoded strings
+        """
+        return self.tokenizer.batch_decode(*args, **kwargs)
+
+    def decode(self, *args, **kwargs) -> str:
+        """
+        This method forwards all its arguments to the tokenizer's decode.
+        
+        Returns:
+            Decoded string
+        """
+        return self.tokenizer.decode(*args, **kwargs)
+
+    def post_process_dna_to_text(
+        self,
+        generated_outputs: torch.Tensor,
+        skip_special_tokens: bool = True,
+        **kwargs,
+    ) -> List[str]:
+        """
+        Post-process the model output to decode the text.
+        
+        Args:
+            generated_outputs: The token IDs generated by the model
+            skip_special_tokens: Whether to skip special tokens in the output
+            **kwargs: Additional arguments for the decoder
+            
+        Returns:
+            List of decoded strings
+        """
+        return self.tokenizer.batch_decode(
+            generated_outputs,
+            skip_special_tokens=skip_special_tokens,
+            **kwargs,
+        )
+
+    @property
+    def model_input_names(self) -> List[str]:
+        """
+        Get the input names expected by the model.
+        
+        Returns:
+            List of input names
+        """
+        tokenizer_input_names = self.tokenizer.model_input_names
+        dna_input_names = ["dna_tokenized", "batch_idx_map"]
+        
+        return list(dict.fromkeys(tokenizer_input_names + dna_input_names))

BioReason/bioreason/models/dna_llm.py

@@ -0,0 +1,306 @@
+import os
+from argparse import ArgumentParser
+import torch
+import torch.nn as nn
+from transformers import (
+    AutoTokenizer,
+    AutoModelForCausalLM,
+    AutoModelForMaskedLM,
+)
+
+from typing import Optional, List, Dict, Any, Union, Tuple
+
+from bioreason.utils.dna_utils import DNAInput
+from bioreason.models.dl.processing_dl import DLProcessor
+from bioreason.models.dl.chat_template_dl import CHAT_TEMPLATE
+from bioreason.models.evo2_tokenizer import Evo2Tokenizer
+
+class DNALLMModel(nn.Module):
+    """
+    A combined model that processes both DNA sequences and text inputs.
+
+    The model uses a DNA encoder (like NucleotideTransformer) to extract features from DNA sequences
+    and a text model (LLM) to process text inputs and generate responses. The DNA features are
+    projected to the text model's embedding space and prepended to the text embeddings.
+    """
+
+    def __init__(
+        self,
+        text_model_name: str,
+        dna_model_name: str,
+        cache_dir: Optional[str] = None,
+        max_length_dna: int = 2048,
+        max_length_text: int = 512,
+        text_model_finetune: bool = True,
+        dna_model_finetune: bool = True,
+        dna_is_evo2: bool = False,
+        dna_embedding_layer: str = None
+    ):
+        """
+        Initialize the DNALLMModel.
+
+        Args:
+            text_model_name: Name of the text model to be used.
+            dna_model_name: Name of the DNA model to be used.
+            cache_dir: Directory to cache the models.
+            max_length_dna: Maximum length of DNA sequences. Defaults to 2048.
+            max_length_text: Maximum length of text sequences. Defaults to 512.
+            text_model_finetune: Whether to finetune the text model. Defaults to True.
+            dna_model_finetune: Whether to finetune the DNA model. Defaults to True.
+            dna_is_evo2: Whether the DNA model is Evo2. Defaults to False.
+            dna_embedding_layer: Name of the layer to use for the Evo2 model. Defaults to None.
+        """
+        super().__init__()
+
+        self.text_model_finetune = text_model_finetune
+        self.dna_model_finetune = dna_model_finetune
+        self.max_length_dna = max_length_dna
+        self.max_length_text = max_length_text
+        self.dna_is_evo2 = dna_is_evo2
+        self.dna_embedding_layer = dna_embedding_layer
+
+
+        # Load the text model and tokenizer
+        self.text_model = AutoModelForCausalLM.from_pretrained(
+            text_model_name, cache_dir=cache_dir, trust_remote_code=True
+        )
+        self.text_tokenizer = AutoTokenizer.from_pretrained(text_model_name, trust_remote_code=True)
+        self.text_config = self.text_model.config
+        self.text_tokenizer.chat_template = CHAT_TEMPLATE
+        self.text_tokenizer.pad_token = self.text_tokenizer.eos_token
+
+        new_tokens = ["<|dna_start|>", "<|dna_pad|>", "<|dna_end|>"]
+        self.text_tokenizer.add_special_tokens({"additional_special_tokens": new_tokens})
+        self.dna_token_id = self.text_tokenizer.convert_tokens_to_ids("<|dna_pad|>")
+
+
+        # Load the DNA model and tokenizer
+        if not self.dna_is_evo2:
+            self.dna_model = AutoModelForMaskedLM.from_pretrained(
+                dna_model_name, cache_dir=cache_dir, trust_remote_code=True
+            )
+            self.dna_tokenizer = AutoTokenizer.from_pretrained(dna_model_name, trust_remote_code=True)
+            self.dna_config = self.dna_model.config
+
+        else:
+            from evo2 import Evo2
+            self.dna_model = Evo2(dna_model_name)
+            self.dna_tokenizer = Evo2Tokenizer(self.dna_model.tokenizer)
+            self.dna_config = self.dna_model.model.config
+            self.dna_embedding_layer = self.dna_embedding_layer
+
+        # Get model dimensions
+        self.text_hidden_size = self.text_config.hidden_size
+        self.dna_hidden_size = self.dna_config.hidden_size
+
+        # Create projection layer to map DNA embeddings to text model's embedding space
+        self.dna_projection = nn.Linear(self.dna_hidden_size, self.text_hidden_size)
+
+        # Create processor for handling inputs
+        self.processor = DLProcessor(tokenizer=self.text_tokenizer, dna_tokenizer=self.dna_tokenizer)
+
+    
+    def process_dna_embeddings(
+        self,
+        dna_tokenized: Dict[str, torch.Tensor],
+        batch_idx_map: List[int],
+        batch_size: int,
+    ) -> List[torch.Tensor]:
+        """
+        Process DNA sequences to obtain embeddings.
+
+        Args:
+            dna_tokenized: Tokenized DNA sequences
+            batch_idx_map: Mapping of each sequence to its batch item
+            batch_size: Number of items in the batch
+
+        Returns:
+            List of tensor embeddings for each batch item
+        """
+        # Process all sequences to get DNA representations
+        with torch.no_grad():
+            # Handle different model types based on dna_is_evo2 attribute
+            if self.dna_is_evo2 and self.dna_embedding_layer is not None:  # Evo2 model
+                # Get embeddings from the specific layer in Evo2
+                hidden_states_list = []
+                
+                for seq_idx in range(len(dna_tokenized["input_ids"])):
+                    # Extract single sequence
+                    input_ids = dna_tokenized["input_ids"][seq_idx:seq_idx+1]
+                    
+                    # Call Evo2 with return_embeddings=True
+                    _, embeddings = self.dna_model(
+                        input_ids,
+                        return_embeddings=True,
+                        layer_names=[self.dna_embedding_layer]
+                    )
+                    
+                    # Get embeddings for the specified layer
+                    seq_embeddings = embeddings[self.dna_embedding_layer].squeeze(0)
+                    hidden_states_list.append(seq_embeddings)
+                
+                # Stack to get same format as non-Evo2 output
+                if hidden_states_list:
+                    hidden_states = torch.stack(hidden_states_list)
+                else:
+                    return [torch.zeros((0, self.text_hidden_size)) for _ in range(batch_size)]
+                    
+            else:  # Standard HuggingFace model
+                # Use existing code path for HF models
+                outputs = self.dna_model(
+                    input_ids=dna_tokenized["input_ids"],
+                    attention_mask=dna_tokenized["attention_mask"],
+                    output_hidden_states=True,
+                )
+                # Get the last hidden state
+                hidden_states = outputs.hidden_states[-1]  # shape: [n_seqs, seq_len, hidden_dim]
+
+        # Project all embeddings at once
+        hidden_states = hidden_states.to(device=self.dna_projection.weight.device, dtype=self.dna_projection.weight.dtype)
+        projected_states = self.dna_projection(hidden_states)
+
+        # Group embeddings by batch item
+        result = [[] for _ in range(batch_size)]
+
+        # For each sequence, get its embeddings and add to appropriate batch result
+        for seq_idx, batch_idx in enumerate(batch_idx_map):
+            # Get only the valid (non-padding) tokens
+            valid_length = dna_tokenized["attention_mask"][seq_idx].sum().item()
+            seq_embedding = projected_states[seq_idx, :valid_length]
+            result[batch_idx].append(seq_embedding)
+
+        # Concatenate embeddings for each batch item
+        for i in range(batch_size):
+            if result[i]:
+                result[i] = torch.cat(result[i], dim=0)
+            else:
+                result[i] = torch.zeros((0, self.text_hidden_size))
+
+        return result
+
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        dna_tokenized: Optional[Dict[str, torch.Tensor]] = None,
+        batch_idx_map: Optional[List[int]] = None,
+        labels: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        """
+        Generate text based on DNA and text inputs.
+
+        Args:
+            input_ids: Input IDs (used if provided directly)
+            attention_mask: Attention mask (used if provided directly)
+            dna_tokenized: Tokenized DNA sequences (used if provided directly)
+            batch_idx_map: Batch mapping for DNA sequences (used if provided directly)
+            labels: Labels for supervised fine-tuning (used if provided directly)
+            **kwargs: Additional arguments for generation
+
+        Returns:
+            Outputs from the text model
+        """
+        # Ensure required inputs are available
+        if input_ids is None or attention_mask is None:
+            raise ValueError("Either 'inputs' or 'input_ids'/'attention_mask' must be provided")
+
+        batch_size = input_ids.shape[0]
+
+        # Get text embeddings from the model's embedding layer
+        text_inputs_embeds = self.text_model.get_input_embeddings()(input_ids)
+
+        if dna_tokenized is not None and batch_idx_map:
+            batch_dna_embeds = self.process_dna_embeddings(dna_tokenized, batch_idx_map, batch_size)
+
+            mask = input_ids == self.dna_token_id
+
+            n_dna_tokens = mask.sum().item()
+            dna_embeds_flat = torch.cat(batch_dna_embeds, dim=0)
+            n_dna_features = dna_embeds_flat.shape[0]
+
+            if n_dna_features != n_dna_tokens:
+                raise ValueError(
+                    f"DNA features and DNA tokens do not match: features {n_dna_features}, tokens: {n_dna_tokens}"
+                )
+
+            # Ensure DNA embeddings have the same dtype as the text embeddings
+            dna_embeds_flat = dna_embeds_flat.to(dtype=text_inputs_embeds.dtype)
+            text_inputs_embeds[mask] = dna_embeds_flat
+
+        # Handle labels if provided (for training)
+        if labels is not None:
+            # TODO: Implement this
+            pass
+
+        # Forward pass through the text model (loss is computed if labels is provided)
+        outputs = self.text_model(
+            inputs_embeds=text_inputs_embeds,
+            attention_mask=attention_mask,
+            labels=labels,
+            **kwargs,
+        )
+
+        return outputs
+
+    def generate(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        dna_tokenized: Optional[Dict[str, torch.Tensor]] = None,
+        batch_idx_map: Optional[List[int]] = None,
+        **generation_kwargs,
+    ) -> Union[torch.Tensor, List[str]]:
+        """
+        Generate text based on DNA and text inputs.
+
+        Args:
+            inputs: The preprocessed inputs from the processor (preferred method)
+            batch_dna_sequences: List of lists of DNA sequences per batch item (legacy method)
+            input_texts: List of input texts (legacy method)
+            input_ids: Input IDs (used if provided directly)
+            attention_mask: Attention mask (used if provided directly)
+            dna_tokenized: Tokenized DNA sequences (used if provided directly)
+            batch_idx_map: Batch mapping for DNA sequences (used if provided directly)
+            **generation_kwargs: Additional arguments for generation
+
+        Returns:
+            Generated token IDs which can be decoded using the processor
+        """
+        # Ensure required inputs are available
+        if input_ids is None or attention_mask is None:
+            raise ValueError("Either 'inputs' or 'input_ids'/'attention_mask' must be provided")
+
+        batch_size = input_ids.shape[0]
+
+        # Get text embeddings from the model's embedding layer
+        text_inputs_embeds = self.text_model.get_input_embeddings()(input_ids)
+
+        if dna_tokenized is not None and batch_idx_map:
+            batch_dna_embeds = self.process_dna_embeddings(dna_tokenized, batch_idx_map, batch_size)
+
+            mask = input_ids == self.dna_token_id
+
+            n_dna_tokens = mask.sum().item()
+            dna_embeds_flat = torch.cat(batch_dna_embeds, dim=0)
+            n_dna_features = dna_embeds_flat.shape[0]
+
+            if n_dna_features != n_dna_tokens:
+                raise ValueError(
+                    f"DNA features and DNA tokens do not match: features {n_dna_features}, tokens: {n_dna_tokens}"
+                )
+
+            # Ensure DNA embeddings have the same dtype as the text embeddings
+            dna_embeds_flat = dna_embeds_flat.to(dtype=text_inputs_embeds.dtype)
+            text_inputs_embeds[mask] = dna_embeds_flat
+
+        # Generation parameters may need adjustment based on model type
+        with torch.no_grad():
+            outputs = self.text_model.generate(
+                inputs_embeds=text_inputs_embeds,
+                attention_mask=attention_mask,
+                use_cache=True,
+                **generation_kwargs,
+            )
+
+        return outputs

BioReason/bioreason/models/dna_only.py

@@ -0,0 +1,203 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Dict
+from transformers import AutoModelForMaskedLM, AutoTokenizer
+
+
+class SelfAttentionPooling(nn.Module):
+    def __init__(self, hidden_size, num_heads=8):
+        super().__init__()
+        # Use PyTorch's built-in multi-head attention
+        self.attention = nn.MultiheadAttention(
+            embed_dim=hidden_size,
+            num_heads=num_heads,
+            batch_first=True
+        )
+        # Learnable query vector
+        self.query = nn.Parameter(torch.randn(1, 1, hidden_size))
+        
+    def forward(self, embeddings, attention_mask=None):
+        # Expand query to batch size
+        batch_size = embeddings.size(0)
+        query = self.query.expand(batch_size, -1, -1)
+        
+        # Create key padding mask from attention mask if provided
+        key_padding_mask = None
+        if attention_mask is not None:
+            key_padding_mask = attention_mask == 0  # Convert to boolean mask where True means ignore
+        
+        # Apply attention: query attends to embeddings
+        context, _ = self.attention(
+            query=query,                  # [batch_size, 1, hidden_size]
+            key=embeddings,               # [batch_size, seq_len, hidden_size]
+            value=embeddings,             # [batch_size, seq_len, hidden_size]
+            key_padding_mask=key_padding_mask
+        )
+        
+        # Squeeze out the singleton dimension
+        return context.squeeze(1)         # [batch_size, hidden_size]
+
+
+class DNAClassifierModel(nn.Module):
+    """
+    A simple classifier that uses a DNA model with a classification head.
+    """
+
+    def __init__(
+        self,
+        dna_model_name: str,
+        cache_dir: str = None,
+        max_length_dna: int = 4096,
+        num_classes: int = 2,  # Binary classification by default
+        dna_is_evo2: bool = False,
+        dna_embedding_layer: str = None,
+        train_just_classifier: bool = True
+    ):
+        """
+        Initialize the DNAClassifierModel.
+
+        Args:
+            dna_model_name (str): Name of the DNA model to use
+            cache_dir (str): Directory to cache models
+            max_length_dna (int): Maximum sequence length
+            num_classes (int): Number of output classes
+            dna_is_evo2: Whether the DNA model is Evo2. Defaults to False
+            dna_embedding_layer: Name of the layer to use for the Evo2 model. Defaults to None
+            train_just_classifier: Whether to train just the classifier. Defaults to True
+        """
+        super().__init__()
+
+        self.dna_model_name = dna_model_name
+        self.cache_dir = cache_dir
+        self.max_length_dna = max_length_dna
+        self.num_classes = num_classes
+        self.dna_is_evo2 = dna_is_evo2
+        self.dna_embedding_layer = dna_embedding_layer
+        self.train_just_classifier = train_just_classifier
+
+        # Load the DNA model and tokenizer
+        if not self.dna_is_evo2:
+            self.dna_model = AutoModelForMaskedLM.from_pretrained(
+                dna_model_name, cache_dir=cache_dir, trust_remote_code=True
+            )
+            self.dna_tokenizer = AutoTokenizer.from_pretrained(dna_model_name, trust_remote_code=True)
+            self.dna_config = self.dna_model.config
+
+        else:
+            from evo2 import Evo2
+            from bioreason.models.evo2_tokenizer import Evo2Tokenizer
+            self.dna_model = Evo2(dna_model_name)
+            self.dna_tokenizer = Evo2Tokenizer(self.dna_model.tokenizer)
+            self.dna_config = self.dna_model.model.config
+            self.dna_embedding_layer = self.dna_embedding_layer
+
+        # Get hidden size from model config
+        self.hidden_size = self.dna_config.hidden_size
+
+        # Add the self-attention pooling module
+        self.pooler = SelfAttentionPooling(self.hidden_size)
+
+        # Create classification head that takes concatenated embeddings from both sequences
+        self.classifier = nn.Sequential(
+            nn.Linear(self.hidden_size * 2, self.hidden_size),
+            nn.ReLU(),
+            nn.Dropout(0.1),
+            nn.Linear(self.hidden_size, num_classes),
+        )
+
+        self.max_length_dna = max_length_dna
+
+    def get_dna_embedding(self, input_ids: torch.Tensor, attention_mask: torch.Tensor):
+        """
+        Get DNA embedding for a single DNA sequence using self-attention pooling.
+
+        Args:
+            input_ids: DNA tokenized sequence
+            attention_mask: DNA tokenized sequence attention mask
+
+        Returns:
+            torch.Tensor: Tensor containing the self-attention pooled DNA embedding
+        """
+        # Add batch dimension if not present
+        if input_ids.dim() == 1:
+            input_ids = input_ids.unsqueeze(0)  # [1, seq_len]
+        
+        # Handle attention mask - create if not provided or add batch dimension
+        if attention_mask is None:
+            attention_mask = torch.ones_like(input_ids)
+        elif attention_mask.dim() == 1:
+            attention_mask = attention_mask.unsqueeze(0)  # [1, seq_len]
+        
+        # Get embeddings from DNA model
+        with torch.set_grad_enabled(not self.train_just_classifier):  # Enable gradients for fine-tuning
+
+            if self.dna_is_evo2 and self.dna_embedding_layer is not None:  # Evo2 model
+                # Get embeddings from the specific layer in Evo2
+                _, embeddings = self.dna_model(
+                    input_ids,
+                    return_embeddings=True,
+                    layer_names=[self.dna_embedding_layer]
+                )
+                
+                # Get embeddings for the specified layer
+                hidden_states = embeddings[self.dna_embedding_layer]
+            
+            else:
+                # Get embeddings from the last hidden state
+                outputs = self.dna_model(
+                    input_ids,
+                    attention_mask=attention_mask,
+                    output_hidden_states=True,
+                )
+
+                # Get the last hidden state
+                hidden_states = outputs.hidden_states[-1]
+        
+        # Apply self-attention pooling to get a weighted representation
+        sequence_embedding = self.pooler(hidden_states, attention_mask)
+        return sequence_embedding.squeeze(0)
+
+    def forward(
+        self, ref_ids=None, alt_ids=None, ref_attention_mask=None, alt_attention_mask=None
+    ):
+        """
+        Forward pass of the model.
+
+        Args:
+            ref_ids: Reference sequence token IDsself.dna_model
+            alt_ids: Alternate sequence token IDsself.dna_model
+            ref_attention_mask: Reference sequence attention maskself.dna_model
+            alt_attention_mask: Alternate sequence attention maskself.dna_model
+
+        Returns:
+            torch.Tensor: Classification logits
+        """
+        batch_size = ref_ids.shape[0] if ref_ids is not None else alt_ids.shape[0]
+
+        if batch_size is None:
+            raise ValueError("Either token IDs must be provided")
+
+        ref_embeddings = []
+        alt_embeddings = []
+
+        # Process each example in the batch
+        for i in range(batch_size):
+
+            # Get sequence embeddings
+            ref_embed = self.get_dna_embedding(ref_ids[i], ref_attention_mask[i])
+            alt_embed = self.get_dna_embedding(alt_ids[i], alt_attention_mask[i])
+            ref_embeddings.append(ref_embed)
+            alt_embeddings.append(alt_embed)
+
+        # Stack embeddings
+        ref_embeddings = torch.stack(ref_embeddings)
+        alt_embeddings = torch.stack(alt_embeddings)
+
+        # Concatenate ref and alt embeddings
+        combined_embeddings = torch.cat([ref_embeddings, alt_embeddings], dim=1)
+
+        # Pass through classifier
+        logits = self.classifier(combined_embeddings)
+
+        return logits

BioReason/bioreason/models/esm_tokenizer.py

@@ -0,0 +1,330 @@
+from transformers.tokenization_utils import PreTrainedTokenizer
+from transformers.utils import logging
+from transformers import AutoTokenizer
+from transformers.tokenization_utils_base import BatchEncoding
+import torch
+import numpy as np
+from typing import List, Dict, Optional, Union, Tuple
+
+logger = logging.get_logger(__name__)
+
+class EsmTokenizer(PreTrainedTokenizer):
+    """
+    Tokenizer for ESM models - wraps the ESM tokenizer to be compatible with HuggingFace interfaces.
+    This tokenizer handles protein sequences (amino acid sequences).
+    """
+    vocab_files_names = {}  # ESM tokenizer doesn't require vocab files
+    model_input_names = ["input_ids", "attention_mask"]
+    
+    # Standard amino acid alphabet used by ESM
+    AMINO_ACIDS = "ACDEFGHIKLMNPQRSTVWY"
+    
+    def __init__(
+        self,
+        esm_model_name: str = "facebook/esm2_t33_650M_UR50D",
+        bos_token="<cls>",
+        eos_token="<eos>",
+        pad_token="<pad>",
+        unk_token="<unk>",
+        mask_token="<mask>",
+        **kwargs
+    ):
+        """
+        Initialize the ESM Tokenizer.
+        
+        Args:
+            esm_model_name: Name of the ESM model to load the tokenizer from
+            bos_token: Beginning of sequence token (CLS token in ESM)
+            eos_token: End of sequence token
+            pad_token: Padding token
+            unk_token: Unknown token
+            mask_token: Mask token for masked language modeling
+        """
+        # Load the actual ESM tokenizer
+        try:
+            self.esm_tokenizer = AutoTokenizer.from_pretrained(esm_model_name, trust_remote_code=True)
+        except:
+            # Fallback to manual tokenizer if auto loading fails
+            self.esm_tokenizer = None
+            self._create_manual_tokenizer()
+        
+        # Set special tokens
+        self._pad_token = pad_token
+        self._eos_token = eos_token
+        self._bos_token = bos_token
+        self._unk_token = unk_token
+        self._mask_token = mask_token
+        
+        # Initialize with special tokens
+        super().__init__(
+            bos_token=bos_token,
+            eos_token=eos_token,
+            pad_token=pad_token,
+            unk_token=unk_token,
+            mask_token=mask_token,
+            **kwargs
+        )
+        
+        # Set token IDs
+        if self.esm_tokenizer is not None:
+            self.pad_token_id = self.esm_tokenizer.pad_token_id
+            self.eos_token_id = self.esm_tokenizer.eos_token_id
+            self.bos_token_id = getattr(self.esm_tokenizer, 'cls_token_id', 0)
+            self.unk_token_id = self.esm_tokenizer.unk_token_id
+            self.mask_token_id = getattr(self.esm_tokenizer, 'mask_token_id', 32)
+        else:
+            # Manual token IDs for fallback
+            self.pad_token_id = 1
+            self.eos_token_id = 2
+            self.bos_token_id = 0  # CLS token
+            self.unk_token_id = 3
+            self.mask_token_id = 32
+    
+    def _create_manual_tokenizer(self):
+        """Create a manual tokenizer mapping if ESM tokenizer loading fails."""
+        # Create vocabulary mapping
+        special_tokens = ["<cls>", "<pad>", "<eos>", "<unk>"]
+        amino_acids = list(self.AMINO_ACIDS)
+        
+        self.token_to_id = {}
+        self.id_to_token = {}
+        
+        # Add special tokens first
+        for i, token in enumerate(special_tokens):
+            self.token_to_id[token] = i
+            self.id_to_token[i] = token
+        
+        # Add amino acids
+        for i, aa in enumerate(amino_acids):
+            token_id = i + len(special_tokens)
+            self.token_to_id[aa] = token_id
+            self.id_to_token[token_id] = aa
+        
+        # Add mask token
+        mask_id = 32
+        self.token_to_id["<mask>"] = mask_id
+        self.id_to_token[mask_id] = "<mask>"
+        
+        self._vocab_size = max(self.id_to_token.keys()) + 1
+        
+    @property
+    def vocab_size(self) -> int:
+        """Return the vocab size of the tokenizer."""
+        if self.esm_tokenizer is not None:
+            return self.esm_tokenizer.vocab_size
+        else:
+            return self._vocab_size
+    
+    def get_vocab(self) -> Dict:
+        """Return vocab as a dictionary."""
+        if self.esm_tokenizer is not None:
+            return self.esm_tokenizer.get_vocab()
+        else:
+            return self.token_to_id.copy()
+    
+    def _tokenize(self, text: str) -> List[str]:
+        """Tokenize a protein sequence string."""
+        if self.esm_tokenizer is not None:
+            return self.esm_tokenizer.tokenize(text)
+        else:
+            # Manual tokenization - split into individual amino acids
+            tokens = []
+            for char in text.upper():
+                if char in self.AMINO_ACIDS:
+                    tokens.append(char)
+                else:
+                    tokens.append(self._unk_token)
+            return tokens
+    
+    def _convert_token_to_id(self, token: str) -> int:
+        """Convert a token to an id."""
+        if self.esm_tokenizer is not None:
+            return self.esm_tokenizer.convert_tokens_to_ids(token)
+        else:
+            return self.token_to_id.get(token, self.unk_token_id)
+    
+    def _convert_id_to_token(self, index: int) -> str:
+        """Convert an id to a token."""
+        if self.esm_tokenizer is not None:
+            return self.esm_tokenizer.convert_ids_to_tokens(index)
+        else:
+            return self.id_to_token.get(index, self._unk_token)
+    
+    def convert_tokens_to_string(self, tokens: List[str]) -> str:
+        """Convert a sequence of tokens to a single string."""
+        # Filter out special tokens and join
+        filtered_tokens = []
+        for token in tokens:
+            if token not in [self._bos_token, self._eos_token, self._pad_token]:
+                filtered_tokens.append(token)
+        return "".join(filtered_tokens)
+    
+    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
+        """ESM tokenizer doesn't need vocabulary saving, return empty tuple."""
+        return ()
+    
+    def __call__(
+        self,
+        text: Union[str, List[str]],
+        text_pair: Optional[Union[str, List[str]]] = None,
+        padding: Union[bool, str] = False,
+        truncation: Union[bool, str] = False,
+        max_length: Optional[int] = None,
+        return_tensors: Optional[str] = None,
+        return_token_type_ids: Optional[bool] = None,
+        return_attention_mask: Optional[bool] = True,
+        add_special_tokens: bool = True,
+        **kwargs
+    ) -> BatchEncoding:
+        """
+        Main tokenization method that handles batching and converts to tensors.
+        """
+        # Use ESM tokenizer if available
+        if self.esm_tokenizer is not None:
+            return self.esm_tokenizer(
+                text=text,
+                text_pair=text_pair,
+                padding=padding,
+                truncation=truncation,
+                max_length=max_length,
+                return_tensors=return_tensors,
+                return_token_type_ids=return_token_type_ids,
+                return_attention_mask=return_attention_mask,
+                add_special_tokens=add_special_tokens,
+                **kwargs
+            )
+        
+        # Manual tokenization fallback
+        # Handle single string vs list of strings
+        if isinstance(text, str):
+            text = [text]
+            
+        # Tokenize all sequences
+        input_ids_list = []
+        for seq in text:
+            # Clean sequence (remove spaces, convert to uppercase)
+            seq = seq.replace(" ", "").upper()
+            
+            # Tokenize sequence
+            tokens = self._tokenize(seq)
+            token_ids = [self._convert_token_to_id(token) for token in tokens]
+            
+            # Add special tokens if requested
+            if add_special_tokens:
+                token_ids = [self.bos_token_id] + token_ids + [self.eos_token_id]
+            
+            # Truncate if needed
+            if truncation and max_length and len(token_ids) > max_length:
+                if add_special_tokens:
+                    # Keep BOS, truncate middle, keep EOS
+                    token_ids = [token_ids[0]] + token_ids[1:max_length-1] + [token_ids[-1]]
+                else:
+                    token_ids = token_ids[:max_length]
+                
+            input_ids_list.append(token_ids)
+        
+        # Apply padding if needed
+        if padding:
+            if max_length:
+                max_len = max_length
+            else:
+                max_len = max(len(ids) for ids in input_ids_list)
+            
+            # Create padded sequences and attention masks
+            padded_input_ids = []
+            attention_mask = []
+            
+            for ids in input_ids_list:
+                # Apply right padding (pad on the right for protein sequences)
+                padding_length = max_len - len(ids)
+                padded_ids = ids + [self.pad_token_id] * padding_length
+                mask = [1] * len(ids) + [0] * padding_length
+                
+                padded_input_ids.append(padded_ids)
+                attention_mask.append(mask)
+                
+            input_ids_list = padded_input_ids
+        else:
+            # Create attention mask without padding
+            attention_mask = [[1] * len(ids) for ids in input_ids_list]
+        
+        # Create result dictionary
+        result = {"input_ids": input_ids_list}
+        if return_attention_mask:
+            result["attention_mask"] = attention_mask
+            
+        # Convert to tensors if requested
+        if return_tensors == "pt":
+            result = {k: torch.tensor(v) for k, v in result.items()}
+            
+        # Return a BatchEncoding object
+        return BatchEncoding(
+            data=result,
+            tensor_type=return_tensors,
+            prepend_batch_axis=False,
+            encoding=None
+        )
+    
+    def batch_decode(
+        self, 
+        sequences: Union[List[int], List[List[int]], torch.Tensor], 
+        skip_special_tokens: bool = True, 
+        **kwargs
+    ) -> List[str]:
+        """
+        Decode a batch of token ids to strings.
+        """
+        if self.esm_tokenizer is not None:
+            return self.esm_tokenizer.batch_decode(sequences, skip_special_tokens=skip_special_tokens, **kwargs)
+        
+        if isinstance(sequences, torch.Tensor):
+            sequences = sequences.tolist()
+        
+        results = []
+        for seq in sequences:
+            tokens = [self._convert_id_to_token(token_id) for token_id in seq]
+            if skip_special_tokens:
+                tokens = [token for token in tokens if token not in [
+                    self._bos_token, self._eos_token, self._pad_token, self._unk_token
+                ]]
+            results.append("".join(tokens))
+        
+        return results
+    
+    def decode(
+        self, 
+        token_ids: Union[int, List[int], torch.Tensor], 
+        skip_special_tokens: bool = True, 
+        **kwargs
+    ) -> str:
+        """
+        Decode a single sequence of token ids to a string.
+        """
+        if self.esm_tokenizer is not None:
+            return self.esm_tokenizer.decode(token_ids, skip_special_tokens=skip_special_tokens, **kwargs)
+        
+        if isinstance(token_ids, torch.Tensor):
+            token_ids = token_ids.tolist()
+            
+        # Handle both single sequence and batch with one item
+        if not isinstance(token_ids, list) or not token_ids or not isinstance(token_ids[0], (list, torch.Tensor)):
+            # Single sequence
+            tokens = [self._convert_id_to_token(token_id) for token_id in token_ids]
+            if skip_special_tokens:
+                tokens = [token for token in tokens if token not in [
+                    self._bos_token, self._eos_token, self._pad_token, self._unk_token
+                ]]
+            return "".join(tokens)
+        
+        # Batch with one item
+        return self.batch_decode(token_ids, skip_special_tokens, **kwargs)[0]
+
+
+def register_esm_tokenizer():
+    """Register the EsmTokenizer with HuggingFace's AutoTokenizer."""
+    AutoTokenizer.register("esm", EsmTokenizer)
+    print("EsmTokenizer registered with AutoTokenizer")
+
+
+if __name__ == "__main__":
+    register_esm_tokenizer()

BioReason/bioreason/models/evo2_tokenizer.py

@@ -0,0 +1,219 @@
+from transformers.tokenization_utils import PreTrainedTokenizer
+from transformers.utils import logging
+from transformers import AutoTokenizer
+from transformers.tokenization_utils_base import BatchEncoding
+import torch
+import numpy as np
+from typing import List, Dict, Optional, Union, Tuple
+
+# Register the tokenizer with AutoTokenizer
+from transformers.models.auto import AutoTokenizer
+from transformers.models.auto.tokenization_auto import TOKENIZER_MAPPING
+from transformers.models.auto.configuration_auto import CONFIG_MAPPING
+
+logger = logging.get_logger(__name__)
+
+class Evo2Tokenizer(PreTrainedTokenizer):
+    """
+    Tokenizer for Evo2 models - wraps the CharLevelTokenizer to be compatible with HuggingFace.
+    """
+    vocab_files_names = {}  # No vocab files needed
+    model_input_names = ["input_ids", "attention_mask"]
+    
+    def __init__(
+        self,
+        evo2_tokenizer,
+        bos_token="<s>",
+        eos_token="</s>",
+        pad_token="<pad>",
+        unk_token="<unk>",
+        **kwargs
+    ):
+        """
+        Initialize the Evo2Tokenizer.
+        
+        Args:
+            evo2_tokenizer: The Evo2 CharLevelTokenizer to wrap
+            bos_token: Beginning of sequence token
+            eos_token: End of sequence token
+            pad_token: Padding token
+            unk_token: Unknown token
+        """
+        self.evo2_tokenizer = evo2_tokenizer
+        
+        # Map special tokens to Evo2 tokenizer's special token IDs
+        self._pad_token = pad_token
+        self._eos_token = eos_token
+        self._bos_token = bos_token
+        self._unk_token = unk_token
+        
+        # Initialize with special tokens
+        super().__init__(
+            bos_token=bos_token,
+            eos_token=eos_token,
+            pad_token=pad_token,
+            unk_token=unk_token,
+            **kwargs
+        )
+        
+        # Set token IDs from Evo2 tokenizer
+        self.pad_token_id = self.evo2_tokenizer.pad_id
+        self.eos_token_id = self.evo2_tokenizer.eos_id
+        
+    @property
+    def vocab_size(self) -> int:
+        """Return the vocab size of the tokenizer."""
+        return self.evo2_tokenizer.vocab_size
+    
+    def get_vocab(self) -> Dict:
+        """Return vocab as a dictionary."""
+        # Evo2 CharLevelTokenizer doesn't have a traditional vocab dict
+        # Create a simple mapping of ASCII codes to tokens
+        return {chr(i): i for i in range(self.vocab_size)}
+    
+    def _tokenize(self, text: str) -> List[int]:
+        """Tokenize a string using the Evo2 tokenizer."""
+        return [chr(int(token)) for token in self.evo2_tokenizer.tokenize(text)]
+    
+    def _convert_token_to_id(self, token: str) -> int:
+        """Convert a token to an id using the Evo2 tokenizer."""
+        # Since tokens are just characters, convert to their ASCII value
+        return ord(token)
+    
+    def _convert_id_to_token(self, index: int) -> str:
+        """Convert an id to a token using the Evo2 tokenizer."""
+        # Convert ASCII value back to character
+        return chr(index)
+    
+    def convert_tokens_to_string(self, tokens: List[str]) -> str:
+        """Convert a sequence of tokens to a single string."""
+        return "".join(tokens)
+    
+    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
+        """No vocabulary to save for Evo2Tokenizer, so just return an empty tuple."""
+        return ()
+    
+    def __call__(
+        self,
+        text: Union[str, List[str]],
+        text_pair: Optional[Union[str, List[str]]] = None,
+        padding: Union[bool, str] = False,
+        truncation: Union[bool, str] = False,
+        max_length: Optional[int] = None,
+        return_tensors: Optional[str] = None,
+        return_token_type_ids: Optional[bool] = None,
+        return_attention_mask: Optional[bool] = True,
+        **kwargs
+    ) -> Dict[str, torch.Tensor]:
+        """
+        Main tokenization method that handles batching and converts to tensors.
+        """
+        # Handle single string vs list of strings
+        if isinstance(text, str):
+            text = [text]
+            
+        # Tokenize all sequences - note: tokenizer only accepts strings, not lists
+        input_ids_list = []
+        for seq in text:
+            # Tokenize and convert numpy.uint8 to Python integers
+            tokens = [int(token) for token in self.evo2_tokenizer.tokenize(seq)]
+            
+            # Truncate if needed
+            if truncation and max_length and len(tokens) > max_length:
+                tokens = tokens[:max_length]
+                
+            input_ids_list.append(tokens)
+        
+        # Apply padding if needed
+        if padding:
+            if False:#max_length:
+                max_len = max_length
+            else:
+                max_len = max(len(ids) for ids in input_ids_list)
+            
+            # Create padded sequences and attention masks
+            padded_input_ids = []
+            attention_mask = []
+            
+            for ids in input_ids_list:
+                # Apply left padding (pad on the left)
+                padding_length = max_len - len(ids)
+                padded_ids = [self.pad_token_id] * padding_length + ids
+                mask = [0] * padding_length + [1] * len(ids)
+                
+                padded_input_ids.append(padded_ids)
+                attention_mask.append(mask)
+                
+            input_ids_list = padded_input_ids
+        else:
+            # Create attention mask without padding
+            attention_mask = [[1] * len(ids) for ids in input_ids_list]
+        
+        # Create result dictionary
+        result = {"input_ids": input_ids_list}
+        if return_attention_mask:
+            result["attention_mask"] = attention_mask
+            
+        # Convert to tensors if requested
+        if return_tensors == "pt":
+            result = {k: torch.tensor(v) for k, v in result.items()}
+            
+        # Return a BatchEncoding object rather than a plain dictionary
+        return BatchEncoding(
+            data=result,
+            tensor_type=return_tensors,
+            prepend_batch_axis=False,  # Already handled in our tensor creation
+            encoding=None  # No encoding info from Evo2's tokenizer
+        )
+    
+    def batch_decode(
+        self, 
+        sequences: Union[List[int], List[List[int]], torch.Tensor], 
+        skip_special_tokens: bool = False, 
+        **kwargs
+    ) -> List[str]:
+        """
+        Decode a batch of token ids to strings.
+        """
+        if isinstance(sequences, torch.Tensor):
+            sequences = sequences.tolist()
+        
+        return self.evo2_tokenizer.detokenize_batch(sequences)
+    
+    def decode(
+        self, 
+        token_ids: Union[int, List[int], torch.Tensor], 
+        skip_special_tokens: bool = False, 
+        **kwargs
+    ) -> str:
+        """
+        Decode a single sequence of token ids to a string.
+        """
+        if isinstance(token_ids, torch.Tensor):
+            token_ids = token_ids.tolist()
+            
+        # Single sequence
+        if not isinstance(token_ids, list) or not token_ids or not isinstance(token_ids[0], (list, torch.Tensor)):
+            return self.evo2_tokenizer.detokenize(token_ids)
+        
+        # Batch with one item
+        return self.batch_decode(token_ids, skip_special_tokens, **kwargs)[0]
+
+
+# Register the tokenizer - you'll need to do this when your script loads
+# You might want to put this in your __init__.py file
+def register_evo2_tokenizer():
+    """Register the Evo2Tokenizer with HuggingFace's AutoTokenizer."""
+    
+    # This will register the tokenizer so AutoTokenizer.from_pretrained knows about it
+    AutoTokenizer.register("evo2", Evo2Tokenizer)
+    
+    # If you have a config class, you would also register that
+    # from transformers.models.auto import AutoConfig
+    # AutoConfig.register("evo2", Evo2Config)
+    
+    print("Evo2Tokenizer registered with AutoTokenizer")
+
+
+if __name__ == "__main__":
+    register_evo2_tokenizer()

BioReason/bioreason/models/pl/chat_template_pl.py

@@ -0,0 +1 @@
+CHAT_TEMPLATE = "{%- set protein_count = namespace(value=0) %}{%- if tools %}\n    {{- '<|im_start|>system\\n' }}\n    {%- if messages[0].role == 'system' %}\n        {{- messages[0].content + '\\n\\n' }}\n    {%- endif %}\n    {{- \"# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n    {%- for tool in tools %}\n        {{- \"\\n\" }}\n        {{- tool | tojson }}\n    {%- endfor %}\n    {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}\n{%- else %}\n    {%- if messages[0].role == 'system' %}\n        {{- '<|im_start|>system\\n' + messages[0].content + '<|im_end|>\\n' }}\n    {%- endif %}\n{%- endif %}\n{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}\n{%- for message in messages[::-1] %}\n    {%- set index = (messages|length - 1) - loop.index0 %}\n    {%- if ns.multi_step_tool and message.role == \"user\" and not(message.content is string and message.content.startswith('<tool_response>') and message.content.endswith('</tool_response>')) %}\n        {%- set ns.multi_step_tool = false %}\n        {%- set ns.last_query_index = index %}\n    {%- endif %}\n{%- endfor %}\n{%- for message in messages %}\n    {%- if (message.role == \"user\") or (message.role == \"system\" and not loop.first) %}\n        {{- '<|im_start|>' + message.role + '\\n' }} {%- if message.content is string %}{{- message.content + '<|im_end|>' + '\\n' }}{%- else %}{%- for content in message.content %}{%- if content.type == 'protein' or 'protein' in content %}{%- set protein_count.value = protein_count.value + 1 %}{%- if add_protein_id %}Protein Sequence {{- protein_count.value }}: {%- endif %}<|protein_start|><|protein_pad|><|protein_end|>{%- elif 'text' in content %}{{- content.text }}{%- endif %}{%- endfor %}{{- '<|im_end|>' + '\\n' }}{%- endif %}{%- elif message.role == \"assistant\" %}\n        {%- set content = message.content[0].text %}\n        {%- set reasoning_content = '' %}\n        {%- if message.reasoning_content is defined and message.reasoning_content is not none %}\n            {%- set reasoning_content = message.reasoning_content %}\n        {%- else %}\n            {%- if '</think>' in message.content %}\n                {%- set content = message.content[0].text.split('</think>')[-1].lstrip('\\n') %}\n                {%- set reasoning_content = message.content[0].text.split('</think>')[0].rstrip('\\n').split('<think>')[-1].lstrip('\\n') %}\n            {%- endif %}\n        {%- endif %}\n        {%- if loop.index0 > ns.last_query_index %}\n            {%- if loop.last or (not loop.last and reasoning_content) %}\n                {{- '<|im_start|>' + message.role + '\\n<think>\\n' + reasoning_content.strip('\\n') + '\\n</think>\\n\\n' + content.lstrip('\\n') }}\n            {%- else %}\n                {{- '<|im_start|>' + message.role + '\\n' + content }}\n            {%- endif %}\n        {%- else %}\n            {{- '<|im_start|>' + message.role + '\\n' + content }}\n        {%- endif %}\n        {%- if message.tool_calls %}\n            {%- for tool_call in message.tool_calls %}\n                {%- if (loop.first and content) or (not loop.first) %}\n                    {{- '\\n' }}\n                {%- endif %}\n                {%- if tool_call.function %}\n                    {%- set tool_call = tool_call.function %}\n                {%- endif %}\n                {{- '<tool_call>\\n{\"name\": \"' }}\n                {{- tool_call.name }}\n                {{- '\", \"arguments\": ' }}\n                {%- if tool_call.arguments is string %}\n                    {{- tool_call.arguments }}\n                {%- else %}\n                    {{- tool_call.arguments | tojson }}\n                {%- endif %}\n                {{- '}\\n</tool_call>' }}\n            {%- endfor %}\n        {%- endif %}\n        {{- '<|im_end|>\\n' }}\n    {%- elif message.role == \"tool\" %}\n        {%- if loop.first or (messages[loop.index0 - 1].role != \"tool\") %}\n            {{- '<|im_start|>user' }}\n        {%- endif %}\n        {{- '\\n<tool_response>\\n' }}\n        {{- message.content }}\n        {{- '\\n</tool_response>' }}\n        {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n            {{- '<|im_end|>\\n' }}\n        {%- endif %}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {{- '<|im_start|>assistant\\n' }}\n    {%- if enable_thinking is defined and enable_thinking is false %}\n        {{- '<think>\\n\\n</think>\\n\\n' }}\n    {%- endif %}\n{%- endif %}"

BioReason/bioreason/models/pl/configuration_pl.py

@@ -0,0 +1,234 @@
+from transformers import PretrainedConfig
+
+class ProteinLLMESMConfig(PretrainedConfig):
+    model_type = "protein_llm"
+    base_config_key = "esm_config"
+
+    def __init__(
+        self,
+        # ESM2 related configurations
+        esm_hidden_size=1280,
+        esm_num_layers=33,
+        esm_num_attention_heads=20,
+        esm_vocab_size=33,
+        esm_max_position_embeddings=1026,
+        esm_layer_norm_eps=1e-5,
+        esm_hidden_dropout_prob=0.1,
+        esm_attention_probs_dropout_prob=0.1,
+        esm_intermediate_size=5120,
+        esm_hidden_act="gelu",
+        esm_initializer_range=0.02,
+        esm_layer_norm_eps=1e-5,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        # ESM2 configurations
+        self.esm_hidden_size = esm_hidden_size
+        self.esm_num_layers = esm_num_layers
+        self.esm_num_attention_heads = esm_num_attention_heads
+        self.esm_vocab_size = esm_vocab_size
+        self.esm_max_position_embeddings = esm_max_position_embeddings
+        self.esm_layer_norm_eps = esm_layer_norm_eps
+        self.esm_hidden_dropout_prob = esm_hidden_dropout_prob
+        self.esm_attention_probs_dropout_prob = esm_attention_probs_dropout_prob
+        self.esm_intermediate_size = esm_intermediate_size
+        self.esm_hidden_act = esm_hidden_act
+        self.esm_initializer_range = esm_initializer_range
+
+class ProteinLLMQFormerConfig(PretrainedConfig):
+    model_type = "protein_llm"
+    base_config_key = "qformer_config"
+
+    def __init__(
+        self,
+        # Q-Former configurations
+        qformer_hidden_size=768,
+        qformer_num_hidden_layers=12,
+        qformer_num_attention_heads=12,
+        qformer_intermediate_size=3072,
+        qformer_hidden_act="gelu",
+        qformer_hidden_dropout_prob=0.1,
+        qformer_attention_probs_dropout_prob=0.1,
+        qformer_max_position_embeddings=512,
+        qformer_layer_norm_eps=1e-12,
+        qformer_initializer_range=0.02,
+        qformer_vocab_size=30522,
+        qformer_pad_token_id=0,
+        qformer_position_embedding_type="absolute",
+        qformer_use_cache=True,
+        # Query tokens
+        num_query_tokens=32,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        # Q-Former configurations
+        self.qformer_hidden_size = qformer_hidden_size
+        self.qformer_num_hidden_layers = qformer_num_hidden_layers
+        self.qformer_num_attention_heads = qformer_num_attention_heads
+        self.qformer_intermediate_size = qformer_intermediate_size
+        self.qformer_hidden_act = qformer_hidden_act
+        self.qformer_hidden_dropout_prob = qformer_hidden_dropout_prob
+        self.qformer_attention_probs_dropout_prob = qformer_attention_probs_dropout_prob
+        self.qformer_max_position_embeddings = qformer_max_position_embeddings
+        self.qformer_layer_norm_eps = qformer_layer_norm_eps
+        self.qformer_initializer_range = qformer_initializer_range
+        self.qformer_vocab_size = qformer_vocab_size
+        self.qformer_pad_token_id = qformer_pad_token_id
+        self.qformer_position_embedding_type = qformer_position_embedding_type
+        self.qformer_use_cache = qformer_use_cache
+        self.num_query_tokens = num_query_tokens
+
+class ProteinLLMConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`ProteinLLMModel`]. It is used to instantiate a
+    Protein-LLM model according to the specified arguments, defining the model architecture. The model combines
+    ESM2 protein encoder, Q-Former, and a language model for protein understanding and generation.
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+    Args:
+        vocab_size (`int`, *optional*, defaults to 152064):
+            Vocabulary size of the language model. Defines the number of different tokens that can be represented by the
+            `inputs_ids` passed when calling the model.
+        hidden_size (`int`, *optional*, defaults to 8192):
+            Dimension of the hidden representations in the language model.
+        intermediate_size (`int`, *optional*, defaults to 29568):
+            Dimension of the MLP representations in the language model.
+        num_hidden_layers (`int`, *optional*, defaults to 80):
+            Number of hidden layers in the Transformer encoder of the language model.
+        num_attention_heads (`int`, *optional*, defaults to 64):
+            Number of attention heads for each attention layer in the Transformer encoder.
+        num_key_value_heads (`int`, *optional*, defaults to 8):
+            This is the number of key_value heads that should be used to implement Grouped Query Attention.
+        hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
+            The non-linear activation function (function or string) in the decoder.
+        max_position_embeddings (`int`, *optional*, defaults to 32768):
+            The maximum sequence length that this model might ever be used with.
+        initializer_range (`float`, *optional*, defaults to 0.02):
+            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+        rms_norm_eps (`float`, *optional*, defaults to 1e-05):
+            The epsilon used by the rms normalization layers.
+        use_cache (`bool`, *optional*, defaults to `True`):
+            Whether or not the model should return the last key/values attentions.
+        tie_word_embeddings (`bool`, *optional*, defaults to `False`):
+            Whether the model's input and output word embeddings should be tied.
+        rope_theta (`float`, *optional*, defaults to 1000000.0):
+            The base period of the RoPE embeddings.
+        use_sliding_window (`bool`, *optional*, defaults to `False`):
+            Whether to use sliding window attention.
+        sliding_window (`int`, *optional*, defaults to 4096):
+            Sliding window attention (SWA) window size.
+        max_window_layers (`int`, *optional*, defaults to 80):
+            The number of layers that use SWA.
+        attention_dropout (`float`, *optional*, defaults to 0.0):
+            The dropout ratio for the attention probabilities.
+        esm_config (`Dict`, *optional*):
+            The config for the ESM2 protein encoder initialization.
+        qformer_config (`Dict`, *optional*):
+            The config for the Q-Former initialization.
+        rope_scaling (`Dict`, *optional*):
+            Dictionary containing the scaling configuration for the RoPE embeddings.
+    """
+
+    model_type = "protein_llm"
+    sub_configs = {
+        "esm_config": ProteinLLMESMConfig,
+        "qformer_config": ProteinLLMQFormerConfig
+    }
+    keys_to_ignore_at_inference = ["past_key_values"]
+    
+    # Default tensor parallel plan for base model
+    base_model_tp_plan = {
+        "layers.*.self_attn.q_proj": "colwise",
+        "layers.*.self_attn.k_proj": "colwise",
+        "layers.*.self_attn.v_proj": "colwise",
+        "layers.*.self_attn.o_proj": "rowwise",
+        "layers.*.mlp.gate_proj": "colwise",
+        "layers.*.mlp.up_proj": "colwise",
+        "layers.*.mlp.down_proj": "rowwise",
+    }
+    base_model_pp_plan = {
+        "embed_tokens": (["input_ids"], ["inputs_embeds"]),
+        "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
+        "norm": (["hidden_states"], ["hidden_states"]),
+    }
+
+    def __init__(
+        self,
+        vocab_size=152064,
+        hidden_size=8192,
+        intermediate_size=29568,
+        num_hidden_layers=80,
+        num_attention_heads=64,
+        num_key_value_heads=8,
+        hidden_act="silu",
+        max_position_embeddings=32768,
+        initializer_range=0.02,
+        rms_norm_eps=1e-05,
+        use_cache=True,
+        tie_word_embeddings=False,
+        rope_theta=1000000.0,
+        use_sliding_window=False,
+        sliding_window=4096,
+        max_window_layers=80,
+        attention_dropout=0.0,
+        esm_config=None,
+        qformer_config=None,
+        rope_scaling=None,
+        protein_token_id=None,
+        **kwargs,
+    ):
+        # Initialize ESM config
+        if isinstance(esm_config, dict):
+            self.esm_config = self.sub_configs["esm_config"](**esm_config)
+        elif esm_config is None:
+            self.esm_config = self.sub_configs["esm_config"]()
+        else:
+            self.esm_config = esm_config
+
+        # Initialize Q-Former config
+        if isinstance(qformer_config, dict):
+            self.qformer_config = self.sub_configs["qformer_config"](**qformer_config)
+        elif qformer_config is None:
+            self.qformer_config = self.sub_configs["qformer_config"]()
+        else:
+            self.qformer_config = qformer_config
+
+        # Language model configurations
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.use_sliding_window = use_sliding_window
+        self.sliding_window = sliding_window
+        self.max_window_layers = max_window_layers
+
+        # for backward compatibility
+        if num_key_value_heads is None:
+            num_key_value_heads = num_attention_heads
+
+        self.num_key_value_heads = num_key_value_heads
+        self.hidden_act = hidden_act
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.use_cache = use_cache
+        self.rope_theta = rope_theta
+        self.attention_dropout = attention_dropout
+        self.rope_scaling = rope_scaling
+
+        self.protein_token_id = protein_token_id
+
+        # Validate the correctness of rotary position embeddings parameters
+        if self.rope_scaling is not None and "type" in self.rope_scaling:
+            if self.rope_scaling["type"] == "mrope":
+                self.rope_scaling["type"] = "default"
+            self.rope_scaling["rope_type"] = self.rope_scaling["type"]
+
+        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
+
+__all__ = ["ProteinLLMConfig"]

BioReason/bioreason/models/pl/processing_pl.py

@@ -0,0 +1,276 @@
+from typing import List, Optional, Union, Dict, Any, Tuple
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+from transformers import AutoTokenizer
+from transformers.processing_utils import (
+    CommonKwargs,
+    ProcessingKwargs,
+    ProcessorMixin,
+    Unpack,
+)
+from transformers.feature_extraction_utils import BatchFeature
+from transformers.tokenization_utils_base import PreTokenizedInput, TextInput
+from transformers.utils import logging
+
+from bioreason.utils.protein_utils import ProteinInput
+
+class ProteinLLMProteinKwargs(CommonKwargs):
+    """Keyword arguments specific to protein sequence processing"""
+    max_length_text: Optional[int]
+    max_length_protein: Optional[int]
+
+
+class ProteinLLMProcessorKwargs(ProcessingKwargs, total=False):
+    """Processing keyword arguments for the ProteinLLM processor"""
+    protein_kwargs: ProteinLLMProteinKwargs
+    _defaults = {
+        "text_kwargs": {
+            "padding": False,
+        },
+    }
+
+class ProteinLLMProcessor(ProcessorMixin):
+    r"""
+    Constructs a ProteinLLM processor which wraps an ESM protein tokenizer and a language model tokenizer into a single processor.
+    This processor handles both text and protein sequence processing to prepare inputs for the ProteinLLMModel.
+    
+    Args:
+        tokenizer (PreTrainedTokenizerBase, *optional*):
+            The text tokenizer used for processing text inputs.
+        protein_tokenizer (PreTrainedTokenizerBase, *optional*):
+            The protein tokenizer (ESM) used for processing protein sequences.
+        chat_template (`str`, *optional*): 
+            A Jinja template for chat formatting. If None, will use the tokenizer's template.
+    """
+
+    attributes = ["tokenizer", "protein_tokenizer"]
+    valid_kwargs = ["model", "chat_template"]
+    tokenizer_class = (
+        "Qwen2Tokenizer", "Qwen2TokenizerFast",
+        "GPT2TokenizerFast", "LlamaTokenizer", "LlamaTokenizerFast",
+    )
+    protein_tokenizer_class = ("EsmTokenizer",)
+
+    def __init__(
+        self, tokenizer=None, protein_tokenizer=None, chat_template=None, **kwargs
+    ):
+        """
+        Initialize the processor with text and protein tokenizers.
+        
+        Args:
+            tokenizer: Text tokenizer (usually from a language model)
+            protein_tokenizer: Protein tokenizer (usually ESM tokenizer)
+            chat_template: Template for formatting chat conversations
+            **kwargs: Additional arguments
+        """
+        self.tokenizer = tokenizer
+        self.protein_tokenizer = protein_tokenizer
+
+        self.protein_token = (
+            "<|protein_pad|>"
+            if not hasattr(self.tokenizer, "protein_token")
+            else self.tokenizer.protein_token
+        )
+    
+        # Get chat template from tokenizer if not provided
+        if chat_template is None and hasattr(self.tokenizer, "chat_template"):
+            chat_template = self.tokenizer.chat_template
+        super().__init__(tokenizer, protein_tokenizer, chat_template=chat_template)
+      
+        # The GRPO trainer might expect this to be set
+        if not hasattr(self.tokenizer, 'pad_token') or self.tokenizer.pad_token is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
+
+    def tokenize_protein_sequences(
+        self, 
+        batch_protein_sequences: List[List[str]], 
+        max_length: int = 1024,
+        return_tensors: str = "pt",
+        device: str = "cuda",
+    ) -> Dict[str, Any]:
+        """
+        Tokenize a batch of protein sequences.
+        
+        Args:
+            batch_protein_sequences: List of lists of protein sequences per batch item
+            max_length: Maximum allowed length for protein sequences
+            return_tensors: Return format for tensors ("pt" for PyTorch)
+            device: Device to place tensors on
+            
+        Returns:
+            Dict containing:
+                - protein_tokenized: The tokenized protein sequences 
+                - batch_idx_map: Mapping of which sequences belong to which batch item
+        """
+        # Create a mapping to track which sequences belong to which batch item
+        batch_idx_map = []
+        all_sequences = []
+
+        # Flatten all sequences with batch tracking
+        for batch_idx, protein_sequences in enumerate(batch_protein_sequences):
+            for seq in protein_sequences:
+                all_sequences.append(seq)
+                batch_idx_map.append(batch_idx)
+
+        # If no sequences in the entire batch, return empty dict
+        if not all_sequences:
+            return {"protein_tokenized": None, "batch_idx_map": []}
+
+        # Tokenize all sequences at once
+        protein_tokenized = self.protein_tokenizer(
+            all_sequences,
+            padding=True,
+            truncation=True,
+            max_length=max_length,
+            return_tensors=return_tensors,
+            return_attention_mask=True,
+            add_special_tokens=True,
+        )
+            
+        return {"protein_tokenized": protein_tokenized, "batch_idx_map": batch_idx_map}
+
+    def __call__(
+        self,
+        batch_protein_sequences: Optional[List[List[str]]] = None,
+        text: Optional[
+            Union[
+                TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]
+            ]
+        ] = None,
+        max_length_text: int = 512,
+        max_length_protein: int = 1024,
+        return_tensors: str = "pt",
+        device: str = "cuda",
+        **kwargs: Unpack[ProteinLLMProcessorKwargs],
+    ) -> BatchFeature:
+        """
+        Process text and protein sequences for model input.
+        
+        Args:
+            batch_protein_sequences: List of lists of protein sequences per batch item
+            text: Input text or list of texts
+            max_length_text: Maximum length for text sequences
+            max_length_protein: Maximum length for protein sequences
+            return_tensors: Return format for tensors
+            device: Device to place tensors on
+            **kwargs: Additional processor keyword arguments
+            
+        Returns:
+            BatchFeature with tokenized inputs for the model
+        """
+        output_kwargs = self._merge_kwargs(
+            ProteinLLMProcessorKwargs,
+            tokenizer_init_kwargs=self.tokenizer.init_kwargs,
+            **kwargs,
+        )
+
+        # Ensure text is a list
+        if not isinstance(text, list):
+            text = [text]
+
+        protein_inputs = {}
+        if batch_protein_sequences is not None:
+            # Tokenize protein sequences
+            protein_processing_result = self.tokenize_protein_sequences(
+                batch_protein_sequences,
+                max_length=max_length_protein,
+                return_tensors=return_tensors,
+                device=device,
+            )
+            
+            # Replace protein tokens in text if needed
+            index = 0
+            for i in range(len(text)):
+                while self.protein_token in text[i]:
+                    # For ESM tokenizer, calculate actual tokens excluding special tokens
+                    protein_token_ids = protein_processing_result['protein_tokenized']['input_ids'][index]
+                    # Exclude BOS and EOS tokens (typically first and last tokens in ESM)
+                    num_protein_tokens = (protein_token_ids != self.protein_tokenizer.pad_token_id).sum().item() - 2
+                    num_protein_tokens = max(1, num_protein_tokens)  # Ensure at least 1 token
+                    
+                    text[i] = text[i].replace(
+                        self.protein_token, "<|placeholder|>" * num_protein_tokens, 1
+                    )
+                    index += 1
+                text[i] = text[i].replace("<|placeholder|>", self.protein_token)
+            
+            # Add batch info to the output
+            protein_inputs = {
+                "protein_tokenized": protein_processing_result["protein_tokenized"],
+                "batch_idx_map": protein_processing_result["batch_idx_map"],
+            }
+
+        # Tokenize text
+        text_kwargs = output_kwargs.get("text_kwargs", {})
+        
+        if 'padding' in text_kwargs:
+            del text_kwargs['padding']
+        
+        text_inputs = self.tokenizer(
+            text, 
+            max_length=max_length_text + 2 * max_length_protein,
+            return_tensors=return_tensors,
+            padding=True,
+            truncation=True,
+            **text_kwargs,
+        )
+        
+        # The BatchFeature should have all required fields for the model's forward pass
+        return BatchFeature(data={**text_inputs, **protein_inputs})
+
+    def batch_decode(self, *args, **kwargs) -> List[str]:
+        """
+        This method forwards all its arguments to the tokenizer's batch_decode.
+        
+        Returns:
+            List of decoded strings
+        """
+        return self.tokenizer.batch_decode(*args, **kwargs)
+
+    def decode(self, *args, **kwargs) -> str:
+        """
+        This method forwards all its arguments to the tokenizer's decode.
+        
+        Returns:
+            Decoded string
+        """
+        return self.tokenizer.decode(*args, **kwargs)
+
+    def post_process_protein_to_text(
+        self,
+        generated_outputs: torch.Tensor,
+        skip_special_tokens: bool = True,
+        **kwargs,
+    ) -> List[str]:
+        """
+        Post-process the model output to decode the text.
+        
+        Args:
+            generated_outputs: The token IDs generated by the model
+            skip_special_tokens: Whether to skip special tokens in the output
+            **kwargs: Additional arguments for the decoder
+            
+        Returns:
+            List of decoded strings
+        """
+        return self.tokenizer.batch_decode(
+            generated_outputs,
+            skip_special_tokens=skip_special_tokens,
+            **kwargs,
+        )
+
+    @property
+    def model_input_names(self) -> List[str]:
+        """
+        Get the input names expected by the model.
+        
+        Returns:
+            List of input names
+        """
+        tokenizer_input_names = self.tokenizer.model_input_names
+        protein_input_names = ["protein_tokenized", "batch_idx_map"]
+        
+        return list(dict.fromkeys(tokenizer_input_names + protein_input_names))

BioReason/bioreason/models/protein_llm.py

@@ -0,0 +1,375 @@
+import os
+from argparse import ArgumentParser
+import torch
+import torch.nn as nn
+from transformers import (
+    AutoTokenizer,
+    AutoModelForCausalLM,
+    AutoModelForMaskedLM,
+    AutoModel,  # 新增：用于加载BiomedBERT
+)
+
+
+from typing import Optional, List, Dict, Any, Union, Tuple
+
+from bioreason.utils.protein_utils import ProteinInput
+from bioreason.models.pl.processing_pl import ProteinLLMProcessor
+from bioreason.models.pl.chat_template_pl import CHAT_TEMPLATE
+
+class BiomedBERTQFormer(nn.Module):
+    """
+    使用预训练的BiomedBERT作为Q-Former的模型
+    """
+    
+    def __init__(self, config, biomedbert_model_name="microsoft/BiomedNLP-BiomedBERT-base-uncased-abstract-fulltext"):
+        super().__init__()
+        self.config = config
+        
+        # 直接加载预训练的BiomedBERT模型
+        self.biomedbert = AutoModel.from_pretrained(
+            biomedbert_model_name,
+            trust_remote_code=True,
+            add_cross_attention=True,  # 启用交叉注意力机制
+        )
+        
+        # 获取BiomedBERT的配置
+        self.biomedbert_config = self.biomedbert.config
+        self.biomedbert_hidden_size = self.biomedbert_config.hidden_size
+        
+        # Query tokens - 可学习的参数
+        self.query_tokens = nn.Parameter(
+            torch.zeros(1, config.qformer_config.num_query_tokens, self.biomedbert_hidden_size)
+        )
+        self.query_tokens.data.normal_(mean=0.0, std=0.02)
+        
+        # 如果蛋白质编码器的维度与BiomedBERT不匹配，需要投影层
+        if config.esm_config.esm_hidden_size != self.biomedbert_hidden_size:
+            self.protein_projection = nn.Linear(
+                config.esm_config.esm_hidden_size,
+                self.biomedbert_hidden_size
+            )
+        else:
+            self.protein_projection = None
+            
+    def forward(
+        self,
+        protein_embeds: torch.Tensor,
+        protein_attention_mask: torch.Tensor,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ):
+        """
+        Args:
+            protein_embeds: 来自ESM2的蛋白质嵌入 [batch_size, seq_len, esm_hidden_size]
+            protein_attention_mask: 蛋白质序列的注意力掩码 [batch_size, seq_len]
+        
+        Returns:
+            可用作语言模型输入的查询嵌入
+        """
+        batch_size = protein_embeds.shape[0]
+        
+        # 如果需要，将蛋白质嵌入投影到BiomedBERT维度
+        if self.protein_projection is not None:
+            protein_embeds = self.protein_projection(protein_embeds)
+        
+        # 为批次扩展查询token
+        query_tokens = self.query_tokens.expand(batch_size, -1, -1)
+        
+        # 为查询token创建注意力掩码
+        query_attention_mask = torch.ones(
+            query_tokens.shape[:2], 
+            dtype=torch.long, 
+            device=query_tokens.device
+        )
+        
+        # BiomedBERT前向传播，使用交叉注意力连接蛋白质嵌入
+        query_outputs = self.biomedbert(
+            inputs_embeds=query_tokens,  # 使用嵌入而不是input_ids
+            attention_mask=query_attention_mask,
+            encoder_hidden_states=protein_embeds,  # 蛋白质嵌入作为编码器状态
+            encoder_attention_mask=protein_attention_mask,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        
+        return query_outputs.last_hidden_state
+
+
+class ProteinLLMModel(nn.Module):
+    """
+    使用BiomedBERT作为Q-Former的蛋白质-语言模型
+    """
+
+    def __init__(
+        self,
+        text_model_name: str,
+        protein_model_name: str,
+        biomedbert_model_name: str = "microsoft/BiomedNLP-BiomedBERT-base-uncased-abstract-fulltext",
+        cache_dir: Optional[str] = None,
+        max_length_protein: int = 1024,
+        max_length_text: int = 512,
+        text_model_finetune: bool = True,
+        protein_model_finetune: bool = True,
+        biomedbert_finetune: bool = True,  # 新增：是否微调BiomedBERT
+        qformer_num_query_tokens: int = 32,
+    ):
+        """
+        初始化ProteinLLMModel，使用BiomedBERT作为Q-Former
+
+        Args:
+            text_model_name: 文本模型名称
+            protein_model_name: ESM2蛋白质模型名称
+            biomedbert_model_name: BiomedBERT模型名称
+            cache_dir: 模型缓存目录
+            max_length_protein: 蛋白质序列最大长度
+            max_length_text: 文本序列最大长度
+            text_model_finetune: 是否微调文本模型
+            protein_model_finetune: 是否微调蛋白质模型
+            biomedbert_finetune: 是否微调BiomedBERT
+            qformer_num_query_tokens: Q-Former查询token数量
+        """
+        super().__init__()
+
+        self.text_model_finetune = text_model_finetune
+        self.protein_model_finetune = protein_model_finetune
+        self.biomedbert_finetune = biomedbert_finetune
+        self.max_length_protein = max_length_protein
+        self.max_length_text = max_length_text
+        self.qformer_num_query_tokens = qformer_num_query_tokens
+
+        # 加载文本模型和分词器
+        self.text_model = AutoModelForCausalLM.from_pretrained(
+            text_model_name, cache_dir=cache_dir, trust_remote_code=True
+        )
+        self.text_tokenizer = AutoTokenizer.from_pretrained(text_model_name, trust_remote_code=True)
+        self.text_config = self.text_model.config
+        self.text_tokenizer.chat_template = CHAT_TEMPLATE
+        self.text_tokenizer.pad_token = self.text_tokenizer.eos_token
+
+        # 添加蛋白质特殊token
+        new_tokens = ["<|protein_start|>", "<|protein_pad|>", "<|protein_end|>"]
+        self.text_tokenizer.add_special_tokens({"additional_special_tokens": new_tokens})
+        self.protein_token_id = self.text_tokenizer.convert_tokens_to_ids("<|protein_pad|>")
+
+        # 加载ESM2蛋白质模型和分词器
+        self.protein_model = AutoModelForMaskedLM.from_pretrained(
+            protein_model_name, cache_dir=cache_dir, trust_remote_code=True
+        )
+        self.protein_tokenizer = AutoTokenizer.from_pretrained(protein_model_name, trust_remote_code=True)
+        self.protein_config = self.protein_model.config
+
+        # 获取模型维度
+        self.text_hidden_size = self.text_config.hidden_size
+        self.protein_hidden_size = self.protein_config.hidden_size
+
+        # 创建Q-Former配置（简化版本，因为我们直接使用BiomedBERT）
+        from types import SimpleNamespace
+        qformer_config = SimpleNamespace(
+            qformer_config=SimpleNamespace(
+                num_query_tokens=qformer_num_query_tokens,
+            ),
+            esm_config=SimpleNamespace(
+                esm_hidden_size=self.protein_hidden_size,
+            )
+        )
+
+        # 初始化BiomedBERT Q-Former
+        self.qformer = BiomedBERTQFormer(qformer_config, biomedbert_model_name)
+
+        # 获取BiomedBERT的隐藏层大小
+        biomedbert_hidden_size = self.qformer.biomedbert_hidden_size
+
+        # 创建投影层，将BiomedBERT输出映射到文本模型嵌入空间
+        self.protein_projection = nn.Linear(biomedbert_hidden_size, self.text_hidden_size)
+
+        # 创建处理器
+        self.processor = ProteinLLMProcessor(
+            tokenizer=self.text_tokenizer, 
+            protein_tokenizer=self.protein_tokenizer
+        )
+
+        # 根据微调设置冻结模型参数
+        if not self.protein_model_finetune:
+            for param in self.protein_model.parameters():
+                param.requires_grad = False
+                
+        if not self.text_model_finetune:
+            for param in self.text_model.parameters():
+                param.requires_grad = False
+                
+        if not self.biomedbert_finetune:
+            for param in self.qformer.biomedbert.parameters():
+                param.requires_grad = False
+
+    def process_protein_embeddings(
+        self,
+        protein_tokenized: Dict[str, torch.Tensor],
+        batch_idx_map: List[int],
+        batch_size: int,
+    ) -> List[torch.Tensor]:
+        """
+        通过ESM2和BiomedBERT Q-Former处理蛋白质序列获得嵌入
+
+        Args:
+            protein_tokenized: 分词后的蛋白质序列
+            batch_idx_map: 每个序列到批次项的映射
+            batch_size: 批次中的项目数量
+
+        Returns:
+            每个批次项的张量嵌入列表
+        """
+        # 通过ESM2处理所有序列获得表示
+        with torch.no_grad() if not self.protein_model_finetune else torch.enable_grad():
+            esm_outputs = self.protein_model(
+                input_ids=protein_tokenized["input_ids"],
+                attention_mask=protein_tokenized["attention_mask"],
+                output_hidden_states=True,
+            )
+            protein_hidden_states = esm_outputs.hidden_states[-1]
+
+        # 通过BiomedBERT Q-Former处理获得查询嵌入
+        result = []
+        
+        for batch_idx in range(batch_size):
+            # 找到属于此批次项的序列
+            seq_indices = [i for i, idx in enumerate(batch_idx_map) if idx == batch_idx]
+            
+            if not seq_indices:
+                # 此批次项没有序列，创建虚拟嵌入
+                dummy_embeddings = torch.zeros(
+                    (self.qformer_num_query_tokens, self.text_hidden_size),
+                    device=protein_hidden_states.device,
+                    dtype=protein_hidden_states.dtype
+                )
+                result.append(dummy_embeddings)
+                continue
+            
+            # 获取此批次项的蛋白质嵌入
+            batch_protein_embeds = protein_hidden_states[seq_indices]
+            batch_attention_mask = protein_tokenized["attention_mask"][seq_indices]
+            
+            # 如果有多个序列，沿序列维度连接
+            if len(seq_indices) > 1:
+                concat_embeds = batch_protein_embeds.view(1, -1, self.protein_hidden_size)
+                concat_mask = batch_attention_mask.view(1, -1)
+            else:
+                concat_embeds = batch_protein_embeds
+                concat_mask = batch_attention_mask
+            
+            # 通过BiomedBERT Q-Former
+            with torch.no_grad() if not self.biomedbert_finetune else torch.enable_grad():
+                query_embeddings = self.qformer(
+                    protein_embeds=concat_embeds,
+                    protein_attention_mask=concat_mask,
+                )
+            
+            # 投影到文本模型空间
+            projected_embeddings = self.protein_projection(query_embeddings.squeeze(0))
+            result.append(projected_embeddings)
+
+        return result
+
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        protein_tokenized: Optional[Dict[str, torch.Tensor]] = None,
+        batch_idx_map: Optional[List[int]] = None,
+        labels: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        """前向传播"""
+        if input_ids is None or attention_mask is None:
+            raise ValueError("input_ids and attention_mask must be provided")
+
+        batch_size = input_ids.shape[0]
+        text_inputs_embeds = self.text_model.get_input_embeddings()(input_ids)
+
+        if protein_tokenized is not None and batch_idx_map:
+            batch_protein_embeds = self.process_protein_embeddings(protein_tokenized, batch_idx_map, batch_size)
+
+            # 用BiomedBERT Q-Former输出替换蛋白质token
+            mask = input_ids == self.protein_token_id
+            n_protein_tokens = mask.sum().item()
+            expected_protein_features = batch_size * self.qformer_num_query_tokens
+
+            if n_protein_tokens != expected_protein_features:
+                protein_embeds_flat = torch.cat(batch_protein_embeds, dim=0)
+                n_protein_features = protein_embeds_flat.shape[0]
+                
+                if n_protein_features > n_protein_tokens:
+                    protein_embeds_flat = protein_embeds_flat[:n_protein_tokens]
+                elif n_protein_features < n_protein_tokens:
+                    padding = torch.zeros(
+                        (n_protein_tokens - n_protein_features, self.text_hidden_size),
+                        device=protein_embeds_flat.device,
+                        dtype=protein_embeds_flat.dtype
+                    )
+                    protein_embeds_flat = torch.cat([protein_embeds_flat, padding], dim=0)
+            else:
+                protein_embeds_flat = torch.cat(batch_protein_embeds, dim=0)
+
+            protein_embeds_flat = protein_embeds_flat.to(dtype=text_inputs_embeds.dtype)
+            text_inputs_embeds[mask] = protein_embeds_flat
+
+        outputs = self.text_model(
+            inputs_embeds=text_inputs_embeds,
+            attention_mask=attention_mask,
+            labels=labels,
+            **kwargs,
+        )
+
+        return outputs
+
+    def generate(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        protein_tokenized: Optional[Dict[str, torch.Tensor]] = None,
+        batch_idx_map: Optional[List[int]] = None,
+        **generation_kwargs,
+    ) -> Union[torch.Tensor, List[str]]:
+        """生成文本"""
+        if input_ids is None or attention_mask is None:
+            raise ValueError("input_ids and attention_mask must be provided")
+
+        batch_size = input_ids.shape[0]
+        text_inputs_embeds = self.text_model.get_input_embeddings()(input_ids)
+
+        if protein_tokenized is not None and batch_idx_map:
+            batch_protein_embeds = self.process_protein_embeddings(protein_tokenized, batch_idx_map, batch_size)
+
+            mask = input_ids == self.protein_token_id
+            n_protein_tokens = mask.sum().item()
+            expected_protein_features = batch_size * self.qformer_num_query_tokens
+
+            if n_protein_tokens != expected_protein_features:
+                protein_embeds_flat = torch.cat(batch_protein_embeds, dim=0)
+                n_protein_features = protein_embeds_flat.shape[0]
+                
+                if n_protein_features > n_protein_tokens:
+                    protein_embeds_flat = protein_embeds_flat[:n_protein_tokens]
+                elif n_protein_features < n_protein_tokens:
+                    padding = torch.zeros(
+                        (n_protein_tokens - n_protein_features, self.text_hidden_size),
+                        device=protein_embeds_flat.device,
+                        dtype=protein_embeds_flat.dtype
+                    )
+                    protein_embeds_flat = torch.cat([protein_embeds_flat, padding], dim=0)
+            else:
+                protein_embeds_flat = torch.cat(batch_protein_embeds, dim=0)
+
+            protein_embeds_flat = protein_embeds_flat.to(dtype=text_inputs_embeds.dtype)
+            text_inputs_embeds[mask] = protein_embeds_flat
+
+        with torch.no_grad():
+            outputs = self.text_model.generate(
+                inputs_embeds=text_inputs_embeds,
+                attention_mask=attention_mask,
+                use_cache=True,
+                **generation_kwargs,
+            )
+
+        return outputs

BioReason/bioreason/models/protein_utils.py

@@ -0,0 +1,233 @@
+from dataclasses import dataclass
+from typing import List, Optional, Dict, Any
+
+
+@dataclass
+class ProteinInput:
+    """
+    Data class for protein-text input pairs.
+    
+    This class encapsulates a text input along with associated protein sequences
+    for processing by the ESM2-LLM model.
+    """
+    text: str
+    protein_sequences: List[str]
+    metadata: Optional[Dict[str, Any]] = None
+    
+    def __post_init__(self):
+        """Validate inputs after initialization."""
+        if not isinstance(self.text, str):
+            raise TypeError("text must be a string")
+        
+        if not isinstance(self.protein_sequences, list):
+            raise TypeError("protein_sequences must be a list")
+        
+        for i, seq in enumerate(self.protein_sequences):
+            if not isinstance(seq, str):
+                raise TypeError(f"protein_sequences[{i}] must be a string")
+            
+            # Basic validation for protein sequence (amino acid characters)
+            valid_aa = set('ACDEFGHIKLMNPQRSTVWYX')  # X for unknown
+            if not all(aa.upper() in valid_aa for aa in seq):
+                raise ValueError(f"protein_sequences[{i}] contains invalid amino acid characters")
+    
+    def __len__(self):
+        """Return the number of protein sequences."""
+        return len(self.protein_sequences)
+    
+    def __getitem__(self, index):
+        """Get a specific protein sequence by index."""
+        return self.protein_sequences[index]
+    
+    def add_protein_sequence(self, sequence: str):
+        """Add a protein sequence to the input."""
+        if not isinstance(sequence, str):
+            raise TypeError("sequence must be a string")
+        
+        # Validate amino acid sequence
+        valid_aa = set('ACDEFGHIKLMNPQRSTVWYX')
+        if not all(aa.upper() in valid_aa for aa in sequence):
+            raise ValueError("sequence contains invalid amino acid characters")
+        
+        self.protein_sequences.append(sequence)
+    
+    def get_total_protein_length(self) -> int:
+        """Get the total length of all protein sequences."""
+        return sum(len(seq) for seq in self.protein_sequences)
+    
+    def get_protein_lengths(self) -> List[int]:
+        """Get the length of each protein sequence."""
+        return [len(seq) for seq in self.protein_sequences]
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary representation."""
+        return {
+            "text": self.text,
+            "protein_sequences": self.protein_sequences,
+            "metadata": self.metadata,
+        }
+    
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "ProteinInput":
+        """Create ProteinInput from dictionary."""
+        return cls(
+            text=data["text"],
+            protein_sequences=data["protein_sequences"],
+            metadata=data.get("metadata"),
+        )
+
+
+def validate_protein_sequence(sequence: str) -> bool:
+    """
+    Validate if a string is a valid protein sequence.
+    
+    Args:
+        sequence: String to validate
+        
+    Returns:
+        True if valid, False otherwise
+    """
+    if not isinstance(sequence, str):
+        return False
+    
+    valid_aa = set('ACDEFGHIKLMNPQRSTVWYX')  # Standard amino acids + X for unknown
+    return all(aa.upper() in valid_aa for aa in sequence)
+
+
+def clean_protein_sequence(sequence: str) -> str:
+    """
+    Clean a protein sequence by removing invalid characters and whitespace.
+    
+    Args:
+        sequence: Raw protein sequence
+        
+    Returns:
+        Cleaned protein sequence
+    """
+    if not isinstance(sequence, str):
+        raise TypeError("sequence must be a string")
+    
+    # Remove whitespace and convert to uppercase
+    cleaned = ''.join(sequence.split()).upper()
+    
+    # Keep only valid amino acid characters
+    valid_aa = set('ACDEFGHIKLMNPQRSTVWYX')
+    cleaned = ''.join(aa for aa in cleaned if aa in valid_aa)
+    
+    return cleaned
+
+
+def load_protein_sequences_from_fasta(filepath: str) -> List[Dict[str, str]]:
+    """
+    Load protein sequences from a FASTA file.
+    
+    Args:
+        filepath: Path to FASTA file
+        
+    Returns:
+        List of dictionaries with 'id' and 'sequence' keys
+    """
+    sequences = []
+    current_id = None
+    current_seq = []
+    
+    try:
+        with open(filepath, 'r') as f:
+            for line in f:
+                line = line.strip()
+                if line.startswith('>'):
+                    # Save previous sequence if exists
+                    if current_id is not None and current_seq:
+                        sequences.append({
+                            'id': current_id,
+                            'sequence': clean_protein_sequence(''.join(current_seq))
+                        })
+                    
+                    # Start new sequence
+                    current_id = line[1:]  # Remove '>'
+                    current_seq = []
+                elif line:
+                    current_seq.append(line)
+            
+            # Save last sequence
+            if current_id is not None and current_seq:
+                sequences.append({
+                    'id': current_id,
+                    'sequence': clean_protein_sequence(''.join(current_seq))
+                })
+    
+    except FileNotFoundError:
+        raise FileNotFoundError(f"FASTA file not found: {filepath}")
+    except Exception as e:
+        raise Exception(f"Error reading FASTA file: {e}")
+    
+    return sequences
+
+
+def create_protein_inputs_from_fasta(
+    fasta_filepath: str,
+    text_template: str = "Analyze this protein: {protein_id}",
+    include_sequence_in_text: bool = False,
+) -> List[ProteinInput]:
+    """
+    Create ProteinInput objects from a FASTA file.
+    
+    Args:
+        fasta_filepath: Path to FASTA file
+        text_template: Template for creating text (can use {protein_id})
+        include_sequence_in_text: Whether to include sequence in text
+        
+    Returns:
+        List of ProteinInput objects
+    """
+    sequences = load_protein_sequences_from_fasta(fasta_filepath)
+    
+    protein_inputs = []
+    for seq_data in sequences:
+        # Create text from template
+        text = text_template.format(protein_id=seq_data['id'])
+        
+        if include_sequence_in_text:
+            text += f" Sequence: {seq_data['sequence']}"
+        
+        # Create ProteinInput
+        protein_input = ProteinInput(
+            text=text,
+            protein_sequences=[seq_data['sequence']],
+            metadata={'protein_id': seq_data['id']}
+        )
+        
+        protein_inputs.append(protein_input)
+    
+    return protein_inputs
+
+
+# Example usage and utility functions
+if __name__ == "__main__":
+    # Example usage
+    protein_input = ProteinInput(
+        text="What is the function of this protein?",
+        protein_sequences=[
+            "MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLA",
+            "ACDEFGHIKLMNPQRSTVWY"
+        ],
+        metadata={"source": "example"}
+    )
+    
+    print(f"Number of protein sequences: {len(protein_input)}")
+    print(f"Total protein length: {protein_input.get_total_protein_length()}")
+    print(f"Individual lengths: {protein_input.get_protein_lengths()}")
+    
+    # Test validation
+    try:
+        invalid_input = ProteinInput(
+            text="Test",
+            protein_sequences=["INVALID123"]  # Contains numbers
+        )
+    except ValueError as e:
+        print(f"Validation caught invalid sequence: {e}")
+    
+    # Test cleaning
+    dirty_sequence = "  AC DE FG HI  123 KL  "
+    clean_sequence = clean_protein_sequence(dirty_sequence)
+    print(f"Cleaned sequence: '{clean_sequence}'")

BioReason/bioreason/trainer/__init__.py

@@ -0,0 +1,7 @@
+from .grpo_config import DNALLMGRPOConfig
+from .grpo_trainer import DNALLMGRPOTrainer
+
+__all__ = [
+    "DNALLMGRPOConfig",
+    "DNALLMGRPOTrainer",
+]

BioReason/bioreason/trainer/demo_grpo.py

@@ -0,0 +1,811 @@
+# Copyright 2025 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import textwrap
+import warnings
+from collections import defaultdict
+from typing import Any, Callable, Optional, Sized, Union
+from unittest.mock import patch
+
+import torch
+import torch.utils.data
+import transformers
+from accelerate.utils import broadcast_object_list, gather, gather_object, is_peft_model, set_seed
+from accelerate.utils.other import is_compiled_module
+from datasets import Dataset, IterableDataset
+from packaging import version
+from torch import nn
+from torch.utils.data import Sampler
+from transformers import (
+    AutoModelForCausalLM,
+    AutoModelForSequenceClassification,
+    AutoTokenizer,
+    GenerationConfig,
+    PreTrainedModel,
+    PreTrainedTokenizerBase,
+    Trainer,
+    TrainerCallback,
+    is_wandb_available,
+)
+from transformers.integrations.deepspeed import is_deepspeed_zero3_enabled
+from transformers.utils import is_peft_available
+
+from trl.data_utils import apply_chat_template, is_conversational, maybe_apply_chat_template
+from trl.import_utils import is_vllm_available
+from trl.models import create_reference_model, prepare_deepspeed, unwrap_model_for_generation
+from trl import SyncRefModelCallback
+from trl import GRPOConfig
+from trl.trainer.utils import generate_model_card, get_comet_experiment_url, pad, selective_log_softmax
+
+
+if is_peft_available():
+    from peft import PeftConfig, get_peft_model
+
+if is_vllm_available():
+    from vllm import LLM, SamplingParams
+
+if is_wandb_available():
+    import wandb
+
+# What we call a reward function is a callable that takes a list of prompts and completions and returns a list of
+# rewards. When it's a string, it's a model ID, so it's loaded as a pretrained model.
+RewardFunc = Union[str, PreTrainedModel, Callable[[list, list], list[float]]]
+
+
+class RepeatRandomSampler(Sampler):
+    """
+    Sampler that repeats the indices of a dataset N times.
+
+    Args:
+        data_source (`Sized`):
+            Dataset to sample from.
+        repeat_count (`int`):
+            Number of times to repeat each index.
+        seed (`Optional[int]`):
+            Random seed for reproducibility (only affects this sampler).
+
+    Example:
+    ```python
+    >>> sampler = RepeatRandomSampler(["a", "b", "c", "d"], repeat_count=2)
+    >>> list(sampler)
+    [2, 2, 0, 0, 3, 3, 1, 1]
+    ```
+    """
+
+    def __init__(self, data_source: Sized, repeat_count: int, seed: Optional[int] = None):
+        self.data_source = data_source
+        self.repeat_count = repeat_count
+        self.num_samples = len(data_source)
+        self.seed = seed
+        self.generator = torch.Generator()  # Create a local random generator
+        if seed is not None:
+            self.generator.manual_seed(seed)
+
+    def __iter__(self):
+        indexes = [
+            idx
+            for idx in torch.randperm(self.num_samples, generator=self.generator).tolist()
+            for _ in range(self.repeat_count)
+        ]
+        return iter(indexes)
+
+    def __len__(self):
+        return self.num_samples * self.repeat_count
+
+# made this to test out the usual pipeline of GRPOTrainer data, and add my own debug messages
+class FakeGRPOTrainer(Trainer):
+    """
+    Trainer for the Group Relative Policy Optimization (GRPO) method. This algorithm was initially proposed in the
+    paper [DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models](https://huggingface.co/papers/2402.03300).
+
+    Example:
+
+    ```python
+    from datasets import load_dataset
+    from trl import GRPOTrainer
+
+    dataset = load_dataset("trl-lib/tldr", split="train")
+
+    def reward_func(completions, **kwargs):
+        # Dummy reward function that rewards completions with more unique letters.
+        return [float(len(set(completion))) for completion in completions]
+
+    trainer = GRPOTrainer(
+        model="Qwen/Qwen2-0.5B-Instruct",
+        reward_funcs=reward_func,
+        train_dataset=dataset,
+    )
+
+    trainer.train()
+    ```
+
+    Args:
+        model (`Union[str, PreTrainedModel]`):
+            Model to be trained. Can be either:
+
+            - A string, being the *model id* of a pretrained model hosted inside a model repo on huggingface.co, or
+              a path to a *directory* containing model weights saved using
+              [`~transformers.PreTrainedModel.save_pretrained`], e.g., `'./my_model_directory/'`. The model is
+              loaded using [`~transformers.AutoModelForCausalLM.from_pretrained`] with the keywork arguments
+              in `args.model_init_kwargs`.
+            - A [`~transformers.PreTrainedModel`] object. Only causal language models are supported.
+        reward_funcs (`Union[RewardFunc, list[RewardFunc]]`):
+            Reward functions to be used for computing the rewards. To compute the rewards, we call all the reward
+            functions with the prompts and completions and sum the rewards. Can be either:
+
+            - A single reward function, such as:
+                - A string: The *model ID* of a pretrained model hosted inside a model repo on huggingface.co, or a
+                path to a *directory* containing model weights saved using
+                [`~transformers.PreTrainedModel.save_pretrained`], e.g., `'./my_model_directory/'`. The model is loaded
+                using [`~transformers.AutoModelForSequenceClassification.from_pretrained`] with `num_labels=1` and the
+                keyword arguments in `args.model_init_kwargs`.
+                - A [`~transformers.PreTrainedModel`] object: Only sequence classification models are supported.
+                - A custom reward function: The function is provided with the prompts and the generated completions,
+                  plus any additional columns in the dataset. It should return a list of rewards. For more details, see
+                  [Using a custom reward function](#using-a-custom-reward-function).
+            - A list of reward functions, where each item can independently be any of the above types. Mixing different
+            types within the list (e.g., a string model ID and a custom reward function) is allowed.
+        args ([`GRPOConfig`], *optional*, defaults to `None`):
+            Configuration for this trainer. If `None`, a default configuration is used.
+        train_dataset ([`~datasets.Dataset`] or [`~datasets.IterableDataset`]):
+            Dataset to use for training. It must include a column `"prompt"`. Any additional columns in the dataset is
+            ignored. The format of the samples can be either:
+
+            - [Standard](dataset_formats#standard): Each sample contains plain text.
+            - [Conversational](dataset_formats#conversational): Each sample contains structured messages (e.g., role
+              and content).
+        eval_dataset ([`~datasets.Dataset`], [`~datasets.IterableDataset`] or `dict[str, Union[Dataset, IterableDataset]]`):
+            Dataset to use for evaluation. It must meet the same requirements as `train_dataset`.
+        processing_class ([`~transformers.PreTrainedTokenizerBase`], *optional*, defaults to `None`):
+            Processing class used to process the data. The padding side must be set to "left". If `None`, the
+            processing class is loaded from the model's name with [`~transformers.AutoTokenizer.from_pretrained`].
+        reward_processing_classes (`Union[PreTrainedTokenizerBase, list[PreTrainedTokenizerBase]]`, *optional*, defaults to `None`):
+            Processing classes corresponding to the reward functions specified in `reward_funcs`. Can be either:
+
+            - A single processing class: Used when `reward_funcs` contains only one reward function.
+            - A list of processing classes: Must match the order and length of the reward functions in `reward_funcs`.
+            If set to `None`, or if an element of the list corresponding to a [`~transformers.PreTrainedModel`] is
+            `None`, the tokenizer for the model is automatically loaded using [`~transformers.AutoTokenizer.from_pretrained`].
+            For elements in `reward_funcs` that are custom reward functions (not [`~transformers.PreTrainedModel`]),
+            the corresponding entries in `reward_processing_classes` are ignored.
+        callbacks (list of [`~transformers.TrainerCallback`], *optional*, defaults to `None`):
+            List of callbacks to customize the training loop. Will add those to the list of default callbacks
+            detailed in [here](https://huggingface.co/docs/transformers/main_classes/callback).
+
+            If you want to remove one of the default callbacks used, use the [`~transformers.Trainer.remove_callback`]
+            method.
+        optimizers (`tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]`, *optional*, defaults to `(None, None)`):
+            A tuple containing the optimizer and the scheduler to use. Will default to an instance of [`AdamW`] on your
+            model and a scheduler given by [`get_linear_schedule_with_warmup`] controlled by `args`.
+        peft_config ([`~peft.PeftConfig`], *optional*, defaults to `None`):
+            PEFT configuration used to wrap the model. If `None`, the model is not wrapped.
+    """
+
+    _tag_names = ["trl", "grpo"]
+
+    def __init__(
+        self,
+        model: Union[str, PreTrainedModel],
+        reward_funcs: Union[RewardFunc, list[RewardFunc]],
+        args: GRPOConfig = None,
+        train_dataset: Optional[Union[Dataset, IterableDataset]] = None,
+        eval_dataset: Optional[Union[Dataset, IterableDataset, dict[str, Union[Dataset, IterableDataset]]]] = None,
+        processing_class: Optional[PreTrainedTokenizerBase] = None,
+        reward_processing_classes: Optional[Union[PreTrainedTokenizerBase, list[PreTrainedTokenizerBase]]] = None,
+        callbacks: Optional[list[TrainerCallback]] = None,
+        optimizers: tuple[Optional[torch.optim.Optimizer], Optional[torch.optim.lr_scheduler.LambdaLR]] = (None, None),
+        peft_config: Optional["PeftConfig"] = None,
+    ):
+        # Args
+        if args is None:
+            model_name = model if isinstance(model, str) else model.config._name_or_path
+            model_name = model_name.split("/")[-1]
+            args = GRPOConfig(f"{model_name}-GRPO")
+
+        # Models
+        # Trained model
+        model_init_kwargs = args.model_init_kwargs or {}
+        if isinstance(model, str):
+            model_id = model
+            torch_dtype = model_init_kwargs.get("torch_dtype")
+            if isinstance(torch_dtype, torch.dtype) or torch_dtype == "auto" or torch_dtype is None:
+                pass  # torch_dtype is already a torch.dtype or "auto" or None
+            elif isinstance(torch_dtype, str):  # it's a str, but not "auto"
+                torch_dtype = getattr(torch, torch_dtype)
+                model_init_kwargs["torch_dtype"] = torch_dtype
+            else:
+                raise ValueError(
+                    "Invalid `torch_dtype` passed to `GRPOConfig`. Expected either 'auto' or a string representing "
+                    f"a `torch.dtype` (e.g., 'float32'), but got {torch_dtype}."
+                )
+            # Disable caching if gradient checkpointing is enabled (not supported)
+            model_init_kwargs["use_cache"] = (
+                False if args.gradient_checkpointing else model_init_kwargs.get("use_cache")
+            )
+            model = AutoModelForCausalLM.from_pretrained(model, **model_init_kwargs)
+        else:
+            model_id = model.config._name_or_path
+            if args.model_init_kwargs is not None:
+                raise ValueError(
+                    "You passed `model_init_kwargs` to the `GRPOConfig`, but your model is already instantiated. "
+                    "This argument can only be used when the `model` argument is a string."
+                )
+
+        if peft_config is not None:
+            model = get_peft_model(model, peft_config)
+
+        # Reference model
+        if is_deepspeed_zero3_enabled():
+            self.ref_model = AutoModelForCausalLM.from_pretrained(model_id, **model_init_kwargs)
+        elif not is_peft_model(model):
+            # If PEFT configuration is not provided, create a reference model based on the initial model.
+            self.ref_model = create_reference_model(model)
+        else:
+            # If PEFT is used, the reference model is not needed since the adapter can be disabled
+            # to revert to the initial model.
+            self.ref_model = None
+
+        # Processing class
+        if processing_class is None:
+            processing_class = AutoTokenizer.from_pretrained(model.config._name_or_path, padding_side="left")
+
+        # Reward functions
+        if not isinstance(reward_funcs, list):
+            reward_funcs = [reward_funcs]
+        for i, reward_func in enumerate(reward_funcs):
+            if isinstance(reward_func, str):
+                reward_funcs[i] = AutoModelForSequenceClassification.from_pretrained(
+                    reward_func, num_labels=1, **model_init_kwargs
+                )
+        self.reward_funcs = reward_funcs
+
+        # Reward weights
+        if args.reward_weights is not None:
+            if len(args.reward_weights) != len(reward_funcs):
+                raise ValueError(
+                    f"Number of reward weights ({len(args.reward_weights)}) must match number of reward "
+                    f"functions ({len(reward_funcs)})"
+                )
+            self.reward_weights = torch.tensor(args.reward_weights, dtype=torch.float32)
+        else:
+            self.reward_weights = torch.ones(len(reward_funcs), dtype=torch.float32)
+
+        # Reward processing class
+        if reward_processing_classes is None:
+            reward_processing_classes = [None] * len(reward_funcs)
+        elif not isinstance(reward_processing_classes, list):
+            reward_processing_classes = [reward_processing_classes]
+        else:
+            if len(reward_processing_classes) != len(reward_funcs):
+                raise ValueError("The number of reward processing classes must match the number of reward functions.")
+
+        for i, (reward_processing_class, reward_func) in enumerate(zip(reward_processing_classes, reward_funcs)):
+            if isinstance(reward_func, PreTrainedModel):
+                if reward_processing_class is None:
+                    reward_processing_class = AutoTokenizer.from_pretrained(reward_func.config._name_or_path)
+                if reward_processing_class.pad_token_id is None:
+                    reward_processing_class.pad_token = reward_processing_class.eos_token
+                # The reward model computes the reward for the latest non-padded token in the input sequence.
+                # So it's important to set the pad token ID to the padding token ID of the processing class.
+                reward_func.config.pad_token_id = reward_processing_class.pad_token_id
+                reward_processing_classes[i] = reward_processing_class
+        self.reward_processing_classes = reward_processing_classes
+
+        # Data collator
+        def data_collator(features):  # No data collation is needed in GRPO
+            return features
+
+        # Training arguments
+        self.max_prompt_length = args.max_prompt_length
+        self.max_completion_length = args.max_completion_length  # = |o_i| in the GRPO paper
+        self.num_generations = args.num_generations  # = G in the GRPO paper
+        self.use_vllm = args.use_vllm
+
+        self.beta = args.beta
+
+        # The trainer estimates the number of FLOPs (floating-point operations) using the number of elements in the
+        # input tensor associated with the key "input_ids". However, in GRPO, the sampled data does not include the
+        # "input_ids" key. Instead, the available keys is "prompt". As a result, the trainer issues the warning:
+        # "Could not estimate the number of tokens of the input, floating-point operations will not be computed." To
+        # suppress this warning, we set the "estimate_tokens" key in the model's "warnings_issued" dictionary to True.
+        # This acts as a flag to indicate that the warning has already been issued.
+        model.warnings_issued["estimate_tokens"] = True
+
+        # Initialize the metrics
+        self._metrics = defaultdict(list)
+        self.log_completions = args.log_completions
+
+        super().__init__(
+            model=model,
+            args=args,
+            data_collator=data_collator,
+            train_dataset=train_dataset,
+            eval_dataset=eval_dataset,
+            processing_class=processing_class,
+            callbacks=callbacks,
+            optimizers=optimizers,
+        )
+
+        # Check if the per_device_train/eval_batch_size * num processes can be divided by the number of generations
+        num_processes = self.accelerator.num_processes
+        global_batch_size = args.per_device_train_batch_size * num_processes
+        possible_values = [n_gen for n_gen in range(2, global_batch_size + 1) if (global_batch_size) % n_gen == 0]
+        if self.num_generations not in possible_values:
+            raise ValueError(
+                f"The global train batch size ({num_processes} x {args.per_device_train_batch_size}) must be evenly "
+                f"divisible by the number of generations per prompt ({self.num_generations}). Given the current train "
+                f"batch size, the valid values for the number of generations are: {possible_values}."
+            )
+        if self.args.eval_strategy != "no":
+            global_batch_size = args.per_device_eval_batch_size * num_processes
+            possible_values = [n_gen for n_gen in range(2, global_batch_size + 1) if (global_batch_size) % n_gen == 0]
+            if self.num_generations not in possible_values:
+                raise ValueError(
+                    f"The global eval batch size ({num_processes} x {args.per_device_eval_batch_size}) must be evenly "
+                    f"divisible by the number of generations per prompt ({self.num_generations}). Given the current "
+                    f"eval batch size, the valid values for the number of generations are: {possible_values}."
+                )
+
+        # Ensure each process receives a unique seed to prevent duplicate completions when generating with
+        # transformers if num_generations exceeds per_device_train_batch_size. We could skip it if we use vLLM, but
+        # it's safer to set it in all cases.
+        set_seed(args.seed, device_specific=True)
+
+        if self.use_vllm:
+            if not is_vllm_available():
+                raise ImportError(
+                    "vLLM is not available and `use_vllm` is set to True. Please install vLLM with "
+                    "`pip install vllm` to use it."
+                )
+
+            if self.accelerator.is_main_process:
+                vllm_device = self.args.vllm_device
+                if vllm_device == "auto":
+                    if torch.cuda.device_count() == 1:
+                        vllm_device = "cuda:0"  # particular case when training with onyl 1 GPU: share it
+                    else:
+                        vllm_device = f"cuda:{self.accelerator.num_processes}"  # take the next GPU idx
+                # Check that the requested device is available
+                if vllm_device.split(":")[0] == "cuda" and int(vllm_device.split(":")[1]) >= torch.cuda.device_count():
+                    raise ValueError(
+                        f"The requested device for vllm ({vllm_device}) is not available. You are likely using vLLM "
+                        "without restricting the number of GPUs for training. Set the `--num_processes` argument to a "
+                        "value lower than the number of GPUs available on your machine—typically, reducing it by one "
+                        f"is sufficient. In your case: `--num_processes {torch.cuda.device_count() - 1}`."
+                    )
+                # Check that the requested device is not also used for training
+                if vllm_device in {f"cuda:{idx}" for idx in range(self.accelerator.num_processes)}:
+                    warnings.warn(
+                        f"The requested device {vllm_device} is also being used for training. For higher throughput "
+                        "and to avoid out-of-memory errors, it is recommended to use a dedicated device for vLLM. "
+                        "If this is intentional, you may ignore this warning but should adjust "
+                        "`vllm_gpu_memory_utilization` accordingly."
+                    )
+                # vLLM is not compatible with accelerate. So we need to patch it to make sure we can (1) place the vLLM
+                # model on the desired device (world_size_patch) and (2) avoid a test that is not designed for our
+                # setting (profiling_patch).
+                world_size_patch = patch("torch.distributed.get_world_size", return_value=1)
+                profiling_patch = patch(
+                    "vllm.worker.worker.Worker._assert_memory_footprint_increased_during_profiling", return_value=None
+                )
+                with world_size_patch, profiling_patch:
+                    self.llm = LLM(
+                        model=model.name_or_path,
+                        device=vllm_device,
+                        gpu_memory_utilization=self.args.vllm_gpu_memory_utilization,
+                        dtype=self.args.vllm_dtype,
+                        # Automatic Prefix Caching caches the KV cache of existing queries, so that a new query can
+                        # directly reuse the KV cache if it shares the same prefix with one of the existing queries.
+                        # This is particularly useful here because we generate completions from the same prompts.
+                        enable_prefix_caching=True,
+                        max_model_len=self.args.vllm_max_model_len,
+                    )
+                self.sampling_params = SamplingParams(
+                    temperature=args.temperature,
+                    max_tokens=self.max_completion_length,
+                )
+
+            self._last_loaded_step = 0  # tag to avoid useless loading during grad accumulation
+
+            # When using vLLM, the main process is responsible for loading the model weights. This can cause process
+            # desynchronization and seems to lead to DeepSpeed hanging during initialization. To prevent this, we
+            # synchronize all processes after vLLM has been fully initialized.
+            self.accelerator.wait_for_everyone()
+        else:
+            self.generation_config = GenerationConfig(
+                max_new_tokens=self.max_completion_length,
+                do_sample=True,
+                temperature=args.temperature,
+                pad_token_id=processing_class.pad_token_id,
+            )
+
+        # Gradient accumulation requires scaled loss. Normally, loss scaling in the parent class depends on whether the
+        # model accepts loss-related kwargs. Since we compute our own loss, this check is irrelevant. We set
+        # self.model_accepts_loss_kwargs to False to enable scaling.
+        self.model_accepts_loss_kwargs = False
+
+        # Add tags to the model
+        self.model.add_model_tags(self._tag_names)
+
+        if self.ref_model is not None:
+            if self.is_deepspeed_enabled:
+                self.ref_model = prepare_deepspeed(self.ref_model, self.accelerator)
+            else:
+                self.ref_model = self.accelerator.prepare_model(self.ref_model, evaluation_mode=True)
+
+        if args.sync_ref_model:
+            self.add_callback(SyncRefModelCallback(ref_model=self.ref_model, accelerator=self.accelerator))
+
+        for i, reward_func in enumerate(self.reward_funcs):
+            if isinstance(reward_func, PreTrainedModel):
+                self.reward_funcs[i] = self.accelerator.prepare_model(reward_func, evaluation_mode=True)
+
+    def _set_signature_columns_if_needed(self):
+        # If `self.args.remove_unused_columns` is True, non-signature columns are removed.
+        # By default, this method sets `self._signature_columns` to the model's expected inputs.
+        # In GRPOTrainer, we preprocess data, so using the model's signature columns doesn't work.
+        # Instead, we set them to the columns expected by the `training_step` method, hence the override.
+        if self._signature_columns is None:
+            self._signature_columns = ["prompt"]
+
+    def _get_train_sampler(self) -> Sampler:
+        # Returns a sampler that ensures each prompt is repeated across multiple processes. This guarantees that
+        # identical prompts are distributed to different GPUs, allowing rewards to be computed and normalized correctly
+        # within each prompt group. Using the same seed across processes ensures consistent prompt assignment,
+        # preventing discrepancies in group formation.
+        return RepeatRandomSampler(self.train_dataset, self.num_generations, seed=self.args.seed)
+
+    def _get_eval_sampler(self, eval_dataset) -> Sampler:
+        # Returns a sampler that ensures each prompt is repeated across multiple processes. This guarantees that
+        # identical prompts are distributed to different GPUs, allowing rewards to be computed and normalized correctly
+        # within each prompt group. Using the same seed across processes ensures consistent prompt assignment,
+        # preventing discrepancies in group formation.
+        return RepeatRandomSampler(eval_dataset, self.num_generations, seed=self.args.seed)
+
+    # Get the per-token log probabilities for the completions for the model and the reference model
+    def _get_per_token_logps(self, model, input_ids, attention_mask, logits_to_keep):
+        # We add 1 to `logits_to_keep` because the last logits of the sequence is later excluded
+        logits = model(input_ids=input_ids, attention_mask=attention_mask, logits_to_keep=logits_to_keep + 1).logits
+        logits = logits[:, :-1, :]  # (B, L-1, V), exclude the last logit: it corresponds to the next token pred
+
+        input_ids = input_ids[:, -logits_to_keep:]
+        # For transformers<=4.48, logits_to_keep argument isn't supported, so here we drop logits ourselves.
+        # See https://github.com/huggingface/trl/issues/2770
+        logits = logits[:, -logits_to_keep:]
+        return selective_log_softmax(logits, input_ids)  #  compute logprobs for the input tokens
+
+    def _move_model_to_vllm(self):
+        with unwrap_model_for_generation(
+            self.model, self.accelerator, gather_deepspeed3_params=self.args.ds3_gather_for_generation
+        ) as unwrapped_model:
+            if is_compiled_module(unwrapped_model):
+                unwrapped_model = unwrapped_model._orig_mod
+            if is_peft_model(unwrapped_model):
+                unwrapped_model.merge_adapter()
+                state_dict = unwrapped_model.state_dict()
+                # Remove base_model and base_layer prefixes
+                state_dict = {
+                    k.removeprefix("base_model.model.").replace(".base_layer", ""): v for k, v in state_dict.items()
+                }
+                # Remove values with adapter prefix (example: "_lora")
+                state_dict = {k: v for k, v in state_dict.items() if unwrapped_model.prefix not in k}
+                # When module to save, remove its prefix and discard the original module
+                state_dict = {
+                    k.replace("modules_to_save.default.", ""): v
+                    for k, v in state_dict.items()
+                    if "original_module" not in k
+                }
+            else:
+                state_dict = unwrapped_model.state_dict()
+            if self.accelerator.is_main_process:
+                llm_model = self.llm.llm_engine.model_executor.driver_worker.model_runner.model
+                llm_model.load_weights(state_dict.items())
+            # Unmerge the adapter to restore the model to its original state.
+            # This must be done after loading weights to ensure they correspond to the merged state.
+            if is_peft_model(unwrapped_model):
+                unwrapped_model.unmerge_adapter()
+
+    def _prepare_inputs(self, inputs: dict[str, Union[torch.Tensor, Any]]) -> dict[str, Union[torch.Tensor, Any]]:
+        device = self.accelerator.device
+        prompts = [x["prompt"] for x in inputs]
+        prompts_text = [maybe_apply_chat_template(example, self.processing_class)["prompt"] for example in inputs]
+        prompt_inputs = self.processing_class(
+            prompts_text, return_tensors="pt", padding=True, padding_side="left", add_special_tokens=False
+        )
+        prompt_inputs = super()._prepare_inputs(prompt_inputs)
+        prompt_ids, prompt_mask = prompt_inputs["input_ids"], prompt_inputs["attention_mask"]
+
+        if self.max_prompt_length is not None:
+            prompt_ids = prompt_ids[:, -self.max_prompt_length :]
+            prompt_mask = prompt_mask[:, -self.max_prompt_length :]
+
+        # Generate completions using either vLLM or regular generation
+        if self.args.use_vllm:
+            # First, have main process load weights if needed
+            if self.state.global_step != self._last_loaded_step:
+                self._move_model_to_vllm()
+                self._last_loaded_step = self.state.global_step
+
+            # Generate completions using vLLM: gather all prompts and use them in a single call in the main process
+            all_prompts_text = gather_object(prompts_text)
+            if self.accelerator.is_main_process:
+                outputs = self.llm.generate(all_prompts_text, sampling_params=self.sampling_params, use_tqdm=False)
+                completion_ids = [out.token_ids for completions in outputs for out in completions.outputs]
+            else:
+                completion_ids = [None] * len(all_prompts_text)
+            # Broadcast the completions from the main process to all processes, ensuring each process receives its
+            # corresponding slice.
+            completion_ids = broadcast_object_list(completion_ids, from_process=0)
+            process_slice = slice(
+                self.accelerator.process_index * len(prompts),
+                (self.accelerator.process_index + 1) * len(prompts),
+            )
+            completion_ids = completion_ids[process_slice]
+
+            # Pad the completions, and concatenate them with the prompts
+            completion_ids = [torch.tensor(ids, device=device) for ids in completion_ids]
+            completion_ids = pad(completion_ids, padding_value=self.processing_class.pad_token_id)
+            prompt_completion_ids = torch.cat([prompt_ids, completion_ids], dim=1)
+        else:
+            print("about to generate!!")
+            # Regular generation path
+            with unwrap_model_for_generation(self.model, self.accelerator) as unwrapped_model:
+                prompt_completion_ids = unwrapped_model.generate(
+                    prompt_ids, attention_mask=prompt_mask, generation_config=self.generation_config
+                )
+            
+            print('prompts_ids', prompt_ids, 'attention_mask', prompt_mask)
+            print('prompt_completion_ids', prompt_completion_ids)
+            print('prompt len', prompt_ids.size(1))
+
+            # Compute prompt length and extract completion ids
+            prompt_length = prompt_ids.size(1)
+            prompt_ids = prompt_completion_ids[:, :prompt_length]
+            completion_ids = prompt_completion_ids[:, prompt_length:]
+
+        # Mask everything after the first EOS token
+        is_eos = completion_ids == self.processing_class.eos_token_id
+        eos_idx = torch.full((is_eos.size(0),), is_eos.size(1), dtype=torch.long, device=device)
+        eos_idx[is_eos.any(dim=1)] = is_eos.int().argmax(dim=1)[is_eos.any(dim=1)]
+        sequence_indices = torch.arange(is_eos.size(1), device=device).expand(is_eos.size(0), -1)
+        completion_mask = (sequence_indices <= eos_idx.unsqueeze(1)).int()
+
+        # Concatenate prompt_mask with completion_mask for logit computation
+        attention_mask = torch.cat([prompt_mask, completion_mask], dim=1)  # (B*G, P+C)
+
+        logits_to_keep = completion_ids.size(1)  # we only need to compute the logits for the completion tokens
+
+        with torch.inference_mode():
+            if self.ref_model is not None:
+                ref_per_token_logps = self._get_per_token_logps(
+                    self.ref_model, prompt_completion_ids, attention_mask, logits_to_keep
+                )
+            else:
+                with self.accelerator.unwrap_model(self.model).disable_adapter():
+                    ref_per_token_logps = self._get_per_token_logps(
+                        self.model, prompt_completion_ids, attention_mask, logits_to_keep
+                    )
+
+        # Decode the generated completions
+        completions_text = self.processing_class.batch_decode(completion_ids, skip_special_tokens=True)
+        if is_conversational(inputs[0]):
+            completions = []
+            for prompt, completion in zip(prompts, completions_text):
+                bootstrap = prompt.pop()["content"] if prompt[-1]["role"] == "assistant" else ""
+                completions.append([{"role": "assistant", "content": bootstrap + completion}])
+        else:
+            completions = completions_text
+
+        rewards_per_func = torch.zeros(len(prompts), len(self.reward_funcs), device=device)
+        for i, (reward_func, reward_processing_class) in enumerate(
+            zip(self.reward_funcs, self.reward_processing_classes)
+        ):
+            if isinstance(reward_func, nn.Module):  # Module instead of PretrainedModel for compat with compiled models
+                if is_conversational(inputs[0]):
+                    messages = [{"messages": p + c} for p, c in zip(prompts, completions)]
+                    texts = [apply_chat_template(x, reward_processing_class)["text"] for x in messages]
+                else:
+                    texts = [p + c for p, c in zip(prompts, completions)]
+                reward_inputs = reward_processing_class(
+                    texts, return_tensors="pt", padding=True, padding_side="right", add_special_tokens=False
+                )
+                reward_inputs = super()._prepare_inputs(reward_inputs)
+                with torch.inference_mode():
+                    rewards_per_func[:, i] = reward_func(**reward_inputs).logits[:, 0]  # Shape (B*G,)
+            else:
+                # Repeat all input columns (but "prompt" and "completion") to match the number of generations
+                keys = [key for key in inputs[0] if key not in ["prompt", "completion"]]
+                reward_kwargs = {key: [example[key] for example in inputs] for key in keys}
+                output_reward_func = reward_func(prompts=prompts, completions=completions, **reward_kwargs)
+                rewards_per_func[:, i] = torch.tensor(output_reward_func, dtype=torch.float32, device=device)
+
+        # Gather the reward per function: this part is crucial, because the rewards are normalized per group and the
+        # completions may be distributed across processes
+        rewards_per_func = gather(rewards_per_func)
+
+        # Apply weights to each reward function's output and sum
+        rewards = (rewards_per_func * self.reward_weights.to(device).unsqueeze(0)).sum(dim=1)
+
+        # Compute grouped-wise rewards
+        mean_grouped_rewards = rewards.view(-1, self.num_generations).mean(dim=1)
+        std_grouped_rewards = rewards.view(-1, self.num_generations).std(dim=1)
+
+        # Normalize the rewards to compute the advantages
+        mean_grouped_rewards = mean_grouped_rewards.repeat_interleave(self.num_generations, dim=0)
+        std_grouped_rewards = std_grouped_rewards.repeat_interleave(self.num_generations, dim=0)
+        advantages = (rewards - mean_grouped_rewards) / (std_grouped_rewards + 1e-4)
+
+        # Slice to keep only the local part of the data
+        process_slice = slice(
+            self.accelerator.process_index * len(prompts),
+            (self.accelerator.process_index + 1) * len(prompts),
+        )
+        advantages = advantages[process_slice]
+
+        # Log the metrics
+        reward_per_func = rewards_per_func.mean(0)
+        for i, reward_func in enumerate(self.reward_funcs):
+            if isinstance(reward_func, nn.Module):  # Module instead of PretrainedModel for compat with compiled models
+                reward_func_name = reward_func.config._name_or_path.split("/")[-1]
+            else:
+                reward_func_name = reward_func.__name__
+            self._metrics[f"rewards/{reward_func_name}"].append(reward_per_func[i].item())
+
+        self._metrics["reward"].append(rewards.mean().item())
+        self._metrics["reward_std"].append(std_grouped_rewards.mean().item())
+
+        if (
+            self.log_completions
+            and self.state.global_step % self.args.logging_steps == 0
+            and "wandb" in self.args.report_to
+        ):
+            import pandas as pd
+
+            # For logging
+            table = {
+                "step": [str(self.state.global_step)] * len(rewards),
+                "prompt": gather_object(prompts_text),
+                "completion": gather_object(completions_text),
+                "reward": rewards.tolist(),
+            }
+            df = pd.DataFrame(table)
+
+            if wandb.run is not None and self.accelerator.is_main_process:
+                wandb.log({"completions": wandb.Table(dataframe=df)})
+
+        return {
+            "prompt_ids": prompt_ids,
+            "prompt_mask": prompt_mask,
+            "completion_ids": completion_ids,
+            "completion_mask": completion_mask,
+            "ref_per_token_logps": ref_per_token_logps,
+            "advantages": advantages,
+        }
+
+    def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
+        if return_outputs:
+            raise ValueError("The GRPOTrainer does not support returning outputs")
+        # Compute the per-token log probabilities for the model
+
+        prompt_ids, prompt_mask = inputs["prompt_ids"], inputs["prompt_mask"]
+        completion_ids, completion_mask = inputs["completion_ids"], inputs["completion_mask"]
+        input_ids = torch.cat([prompt_ids, completion_ids], dim=1)
+        attention_mask = torch.cat([prompt_mask, completion_mask], dim=1)
+        logits_to_keep = completion_ids.size(1)  # we only need to compute the logits for the completion tokens
+
+        per_token_logps = self._get_per_token_logps(model, input_ids, attention_mask, logits_to_keep)
+
+        # Compute the KL divergence between the model and the reference model
+        ref_per_token_logps = inputs["ref_per_token_logps"]
+        per_token_kl = torch.exp(ref_per_token_logps - per_token_logps) - (ref_per_token_logps - per_token_logps) - 1
+
+        # x - x.detach() allows for preserving gradients from x
+        advantages = inputs["advantages"]
+        per_token_loss = torch.exp(per_token_logps - per_token_logps.detach()) * advantages.unsqueeze(1)
+        per_token_loss = -(per_token_loss - self.beta * per_token_kl)
+        loss = ((per_token_loss * completion_mask).sum(dim=1) / completion_mask.sum(dim=1)).mean()
+
+        # Log the metrics
+        completion_length = self.accelerator.gather_for_metrics(completion_mask.sum(1)).float().mean().item()
+        self._metrics["completion_length"].append(completion_length)
+
+        mean_kl = ((per_token_kl * completion_mask).sum(dim=1) / completion_mask.sum(dim=1)).mean()
+        self._metrics["kl"].append(self.accelerator.gather_for_metrics(mean_kl).mean().item())
+
+        return loss
+
+    def prediction_step(self, model, inputs, prediction_loss_only, ignore_keys: Optional[list[str]] = None):
+        inputs = self._prepare_inputs(inputs)
+        print("about to loss")
+        with torch.no_grad():
+            with self.compute_loss_context_manager():
+                loss = self.compute_loss(model, inputs)
+            loss = loss.mean().detach()
+        print("loss computed")
+        return loss, None, None
+
+    def log(self, logs: dict[str, float], start_time: Optional[float] = None) -> None:
+        metrics = {key: sum(val) / len(val) for key, val in self._metrics.items()}  # average the metrics
+
+        # This method can be called both in training and evaluation. When called in evaluation, the keys in `logs`
+        # start with "eval_". We need to add the prefix "eval_" to the keys in `metrics` to match the format.
+        if next(iter(logs.keys())).startswith("eval_"):
+            metrics = {f"eval_{key}": val for key, val in metrics.items()}
+
+        logs = {**logs, **metrics}
+        if version.parse(transformers.__version__) >= version.parse("4.47.0.dev0"):
+            super().log(logs, start_time)
+        else:  # transformers<=4.46
+            super().log(logs)
+        self._metrics.clear()
+
+    def create_model_card(
+        self,
+        model_name: Optional[str] = None,
+        dataset_name: Optional[str] = None,
+        tags: Union[str, list[str], None] = None,
+    ):
+        """
+        Creates a draft of a model card using the information available to the `Trainer`.
+
+        Args:
+            model_name (`str` or `None`, *optional*, defaults to `None`):
+                Name of the model.
+            dataset_name (`str` or `None`, *optional*, defaults to `None`):
+                Name of the dataset used for training.
+            tags (`str`, `list[str]` or `None`, *optional*, defaults to `None`):
+                Tags to be associated with the model card.
+        """
+        if not self.is_world_process_zero():
+            return
+
+        if hasattr(self.model.config, "_name_or_path") and not os.path.isdir(self.model.config._name_or_path):
+            base_model = self.model.config._name_or_path
+        else:
+            base_model = None
+
+        tags = tags or []
+        if isinstance(tags, str):
+            tags = [tags]
+
+        if hasattr(self.model.config, "unsloth_version"):
+            tags.append("unsloth")
+
+        citation = textwrap.dedent(
+            """\
+            @article{zhihong2024deepseekmath,
+                title        = {{DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models}},
+                author       = {Zhihong Shao and Peiyi Wang and Qihao Zhu and Runxin Xu and Junxiao Song and Mingchuan Zhang and Y. K. Li and Y. Wu and Daya Guo},
+                year         = 2024,
+                eprint       = {arXiv:2402.03300},
+            }
+            """
+        )
+
+        model_card = generate_model_card(
+            base_model=base_model,
+            model_name=model_name,
+            hub_model_id=self.hub_model_id,
+            dataset_name=dataset_name,
+            tags=tags,
+            wandb_url=wandb.run.get_url() if is_wandb_available() and wandb.run is not None else None,
+            comet_url=get_comet_experiment_url(),
+            trainer_name="GRPO",
+            trainer_citation=citation,
+            paper_title="DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models",
+            paper_id="2402.03300",
+        )
+
+        model_card.save(os.path.join(self.args.output_dir, "README.md"))

BioReason/bioreason/trainer/grpo_config.py

@@ -0,0 +1,365 @@
+# Copyright 2025 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import dataclass, field
+from typing import Optional, Union
+
+from transformers import TrainingArguments
+
+
+@dataclass
+class DNALLMGRPOConfig(TrainingArguments):
+    r"""
+    Configuration class for the [`GRPOTrainer`].
+
+    Only the parameters specific to GRPO training are listed here. For details on other parameters, refer to the
+    [`~transformers.TrainingArguments`] documentation.
+
+    Using [`~transformers.HfArgumentParser`] we can turn this class into
+    [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the
+    command line.
+
+    Parameters:
+        > Parameters that control the model and reference model
+
+        model_init_kwargs (`dict[str, Any]` or `None`, *optional*, defaults to `None`):
+            Keyword arguments for [`~transformers.AutoModelForCausalLM.from_pretrained`], used when the `model`
+            argument of the [`GRPOTrainer`] is provided as a string.
+
+        > Parameters that control the data preprocessing
+
+        remove_unused_columns (`bool`, *optional*, defaults to `False`):
+            Whether to only keep the column `"prompt"` in the dataset. If you use a custom reward function that
+            requires any column other than `"prompts"` and `"completions"`, you should keep this to `False`.
+        max_prompt_length (`int` or `None`, *optional*, defaults to `512`):
+            Maximum length of the prompt. If the prompt is longer than this value, it will be truncated left.
+        num_generations (`int` or `None`, *optional*, defaults to `8`):
+            Number of generations per prompt to sample. The global batch size (num_processes * per_device_batch_size)
+            must be divisible by this value.
+        max_completion_length (`int` or `None`, *optional*, defaults to `256`):
+            Maximum length of the generated completion.
+        ds3_gather_for_generation (`bool`, *optional*, defaults to `True`):
+            This setting applies to DeepSpeed ZeRO-3. If enabled, the policy model weights are gathered for generation,
+            improving generation speed. However, disabling this option allows training models that exceed the VRAM
+            capacity of a single GPU, albeit at the cost of slower generation. Disabling this option is not compatible
+            with vLLM generation.
+
+        > Parameters that control generation
+
+        temperature (`float`, defaults to `0.9`):
+            Temperature for sampling. The higher the temperature, the more random the completions.
+        top_p (`float`, *optional*, defaults to `1.0`):
+            Float that controls the cumulative probability of the top tokens to consider. Must be in (0, 1]. Set to
+            `1.0` to consider all tokens.
+        top_k (`int` or `None`, *optional*, defaults to `50`):
+            Number of highest probability vocabulary tokens to keep for top-k-filtering. If `None`, top-k-filtering is
+            disabled.
+        min_p (`float` or `None`, *optional*, defaults to `None`):
+            Minimum token probability, which will be scaled by the probability of the most likely token. It must be a
+            value between `0.0` and `1.0`. Typical values are in the `0.01-0.2` range.
+        repetition_penalty (`float`, *optional*, defaults to `1.0`):
+            Float that penalizes new tokens based on whether they appear in the prompt and the generated text so far.
+            Values > `1.0` encourage the model to use new tokens, while values < `1.0` encourage the model to repeat
+            tokens.
+        cache_implementation (`str` or `None`, *optional*, defaults to `None`):
+            Implementation of the cache method for faster generation when use_vllm is set to False.
+
+        > Parameters that control generation acceleration powered by vLLM
+
+        use_vllm (`bool`, *optional*, defaults to `False`):
+            Whether to use vLLM for generating completions. If set to `True`, ensure that a GPU is kept unused for
+            training, as vLLM will require one for generation. vLLM must be installed (`pip install vllm`).
+        vllm_device (`str`, *optional*, defaults to `"auto"`):
+            Device where vLLM generation will run, e.g. `"cuda:1"`. If set to `"auto"` (default), the system will
+            automatically select the next available GPU after the last one used for training. This assumes that
+            training has not already occupied all available GPUs. If only one device is available, the device will be
+            shared between both training and vLLM.
+        vllm_gpu_memory_utilization (`float`, *optional*, defaults to `0.9`):
+            Ratio (between 0 and 1) of GPU memory to reserve for the model weights, activations, and KV cache on the
+            device dedicated to generation powered by vLLM. Higher values will increase the KV cache size and thus
+            improve the model's throughput. However, if the value is too high, it may cause out-of-memory (OOM) errors
+            during initialization.
+        vllm_dtype (`str`, *optional*, defaults to `"auto"`):
+            Data type to use for vLLM generation. If set to `"auto"`, the data type will be automatically determined
+            based on the model configuration. Find the supported values in the vLLM documentation.
+        vllm_max_model_len (`int` or `None`, *optional*, defaults to `None`):
+            If set, the `max_model_len` to use for vLLM. This could be useful when running with reduced
+            `vllm_gpu_memory_utilization`, leading to a reduced KV cache size. If not set, vLLM will use the model
+            context size, which might be much larger than the KV cache, leading to inefficiencies.
+        vllm_enable_prefix_caching (`bool`, *optional*, defaults to `True`):
+            Whether to enable prefix caching in vLLM. If set to `True` (default), ensure that the model and the hardware
+            support this feature.
+        vllm_guided_decoding_regex (`str` or `None`, *optional*, defaults to `None`):
+            Regex for vLLM guided decoding. If `None` (default), guided decoding is disabled.
+
+        > Parameters that control the training
+
+        learning_rate (`float`, *optional*, defaults to `1e-6`):
+            Initial learning rate for [`AdamW`] optimizer. The default value replaces that of
+            [`~transformers.TrainingArguments`].
+        beta (`float`, *optional*, defaults to `0.04`):
+            KL coefficient. If `0.0`, the reference model is not loaded, reducing memory usage and improving training
+            speed, but may be numerically unstable for long training runs.
+        num_iterations (`int`, *optional*, defaults to `1`):
+            Number of iterations per batch (denoted as μ in the algorithm).
+        epsilon (`float`, *optional*, defaults to `0.2`):
+            Epsilon value for clipping.
+        epsilon_high (`float` or `None`, *optional*, defaults to `None`):
+            Upper-bound epsilon value for clipping. If not specified, it defaults to the same value as the lower-bound
+            specified in argument `epsilon`. Paper [DAPO](https://huggingface.co/papers/2503.14476) recommends `0.28`.
+        reward_weights (`list[float]` or `None`, *optional*, defaults to `None`):
+            Weights for each reward function. Must match the number of reward functions. If `None`, all rewards are
+            weighted equally with weight `1.0`.
+        sync_ref_model (`bool`, *optional*, defaults to `False`):
+            Whether to synchronize the reference model with the active model every `ref_model_sync_steps` steps, using
+            the `ref_model_mixup_alpha` parameter. This synchronization originites from the
+            [TR-DPO](https://huggingface.co/papers/2404.09656) paper.
+        ref_model_mixup_alpha (`float`, *optional*, defaults to `0.6`):
+            α parameter from the [TR-DPO](https://huggingface.co/papers/2404.09656) paper, which controls the mix
+            between the current policy and the previous reference policy during updates. The reference policy is
+            updated according to the equation: `π_ref = α * π_θ + (1 - α) * π_ref_prev`. To use this parameter, you
+            must set `sync_ref_model=True`.
+        ref_model_sync_steps (`int`, *optional*, defaults to `512`):
+            τ parameter from the [TR-DPO](https://huggingface.co/papers/2404.09656) paper, which determines how
+            frequently the current policy is synchronized with the reference policy. To use this parameter, you must
+            set `sync_ref_model=True`.
+
+        > Parameters that control the logging
+
+        log_completions (`bool`, *optional*, defaults to `False`):
+            Whether to log a sample of (prompt, completion) pairs every `logging_steps` steps. If `rich` is
+            installed, it prints the sample. If `wandb` logging is enabled, it logs it to `wandb`.
+    """
+
+    # Parameters that control the model and reference model
+    model_init_kwargs: Optional[dict] = field(
+        default=None,
+        metadata={
+            "help": "Keyword arguments for `transformers.AutoModelForCausalLM.from_pretrained`, used when the `model` "
+            "argument of the `GRPOTrainer` is provided as a string."
+        },
+    )
+
+    # Parameters that control the data preprocessing
+    # The default value remove_unused_columns is overwritten from the parent class, because in GRPO we usually rely on
+    # additional columns to compute the reward
+    remove_unused_columns: Optional[bool] = field(
+        default=False,
+        metadata={
+            "help": "Whether to only keep the column 'prompt' in the dataset. If you use a custom reward function "
+            "that requires any column other than 'prompts' and 'completions', you should keep this to `False`."
+        },
+    )
+    max_prompt_length: Optional[int] = field(
+        default=512,
+        metadata={
+            "help": "Maximum length of the prompt. If the prompt is longer than this value, it will be truncated left."
+        },
+    )
+    num_generations: Optional[int] = field(
+        default=8,
+        metadata={
+            "help": "Number of generations to sample. The global batch size (num_processes * per_device_batch_size) "
+            "must be divisible by this value."
+        },
+    )
+    max_completion_length: Optional[int] = field(
+        default=800,
+        metadata={"help": "Maximum length of the generated completion."},
+    )
+    ds3_gather_for_generation: bool = field(
+        default=True,
+        metadata={
+            "help": "This setting applies to DeepSpeed ZeRO-3. If enabled, the policy model weights are gathered for "
+            "generation, improving generation speed. However, disabling this option allows training models that "
+            "exceed the VRAM capacity of a single GPU, albeit at the cost of slower generation. Disabling this option "
+            "is not compatible with vLLM generation."
+        },
+    )
+
+    # Parameters that control generation
+    temperature: float = field(
+        default=0.6,
+        metadata={"help": "Temperature for sampling. The higher the temperature, the more random the completions."},
+    )
+    top_p: float = field(
+        default=0.95,
+        metadata={
+            "help": "Float that controls the cumulative probability of the top tokens to consider. Must be in (0, 1]. "
+            "Set to 1.0 to consider all tokens."
+        },
+    )
+    top_k: Optional[int] = field(
+        default=20,
+        metadata={
+            "help": "Number of highest probability vocabulary tokens to keep for top-k-filtering. If `None`, "
+            "top-k-filtering is disabled."
+        },
+    )
+    min_p: Optional[float] = field(
+        default=None,
+        metadata={
+            "help": "Minimum token probability, which will be scaled by the probability of the most likely token. It "
+            "must be a value between 0.0 and 1.0. Typical values are in the 0.01-0.2 range."
+        },
+    )
+    repetition_penalty: float = field(
+        default=1.0,
+        metadata={
+            "help": "Float that penalizes new tokens based on whether they appear in the prompt and the generated "
+            "text so far. Values > 1.0 encourage the model to use new tokens, while values < 1.0 encourage the model "
+            "to repeat tokens."
+        },
+    )
+    cache_implementation: Optional[str] = field(
+        default=None,
+        metadata={"help": "Implementation of the cache method for faster generation when use_vllm is set to False."},
+    )
+
+    # Parameters that control generation acceleration powered by vLLM
+    use_vllm: Optional[bool] = field(
+        default=False,
+        metadata={
+            "help": "Whether to use vLLM for generating completions. If set to `True`, ensure that a GPU is kept "
+            "unused for training, as vLLM will require one for generation. vLLM must be installed "
+            "(`pip install vllm`)."
+        },
+    )
+    vllm_device: Optional[str] = field(
+        default="auto",
+        metadata={
+            "help": "Device where vLLM generation will run, e.g. 'cuda:1'. If set to 'auto' (default), the system "
+            "will automatically select the next available GPU after the last one used for training. This assumes "
+            "that training has not already occupied all available GPUs."
+        },
+    )
+    vllm_gpu_memory_utilization: float = field(
+        default=0.9,
+        metadata={
+            "help": "Ratio (between 0 and 1) of GPU memory to reserve for the model weights, activations, and KV "
+            "cache on the device dedicated to generation powered by vLLM. Higher values will increase the KV cache "
+            "size and thus improve the model's throughput. However, if the value is too high, it may cause "
+            "out-of-memory (OOM) errors during initialization."
+        },
+    )
+    vllm_dtype: Optional[str] = field(
+        default="auto",
+        metadata={
+            "help": "Data type to use for vLLM generation. If set to 'auto', the data type will be automatically "
+            "determined based on the model configuration. Find the supported values in the vLLM documentation."
+        },
+    )
+    vllm_max_model_len: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": "If set, the `max_model_len` to use for vLLM. This could be useful when running with reduced "
+            "`vllm_gpu_memory_utilization`, leading to a reduced KV cache size. If not set, vLLM will use the model "
+            "context size, which might be much larger than the KV cache, leading to inefficiencies."
+        },
+    )
+    vllm_enable_prefix_caching: Optional[bool] = field(
+        default=True,
+        metadata={
+            "help": "Whether to enable prefix caching in vLLM. If set to `True` (default), ensure that the model and "
+            "the hardware support this feature."
+        },
+    )
+    vllm_guided_decoding_regex: Optional[str] = field(
+        default=None,
+        metadata={"help": "Regex for vLLM guided decoding. If `None` (default), guided decoding is disabled."},
+    )
+
+    # Parameters that control the training
+    learning_rate: float = field(
+        default=1e-6,
+        metadata={
+            "help": "Initial learning rate for `AdamW` optimizer. The default value replaces that of "
+            "`transformers.TrainingArguments`."
+        },
+    )
+    beta: float = field(
+        default=0.04,
+        metadata={
+            "help": "KL coefficient. If `0.0`, the reference model is not loaded, reducing memory usage and improving "
+            "training speed, but may be numerically unstable for long training runs."
+        },
+    )
+    num_iterations: int = field(
+        default=1,
+        metadata={"help": "Number of iterations per batch (denoted as μ in the algorithm)."},
+    )
+    epsilon: float = field(
+        default=0.2,
+        metadata={"help": "Epsilon value for clipping."},
+    )
+    epsilon_high: Optional[float] = field(
+        default=None,
+        metadata={
+            "help": "Upper-bound epsilon value for clipping. If not specified, it defaults to the same value as the "
+            "lower-bound specified in argument `epsilon`. Paper DAPO recommends `0.28`."
+        },
+    )
+    reward_weights: Optional[list[float]] = field(
+        default=None,
+        metadata={
+            "help": "Weights for each reward function. Must match the number of reward functions. If `None`, all "
+            "rewards are weighted equally with weight `1.0`."
+        },
+    )
+    sync_ref_model: bool = field(
+        default=False,
+        metadata={
+            "help": "Whether to synchronize the reference model with the active model every `ref_model_sync_steps` "
+            "steps, using the `ref_model_mixup_alpha` parameter."
+        },
+    )
+    ref_model_mixup_alpha: float = field(
+        default=0.6,
+        metadata={
+            "help": "α parameter from the TR-DPO paper, which controls the mix between the current policy and the "
+            "previous reference policy during updates. The reference policy is updated according to the equation: "
+            "`π_ref = α * π_θ + (1 - α) * π_ref_prev`. To use this parameter, you must set `sync_ref_model=True`."
+        },
+    )
+    ref_model_sync_steps: int = field(
+        default=512,
+        metadata={
+            "help": "τ parameter from the TR-DPO paper, which determines how frequently the current policy is "
+            "synchronized with the reference policy. To use this parameter, you must set `sync_ref_model=True`."
+        },
+    )
+
+    # Parameters that control the logging
+    log_completions: bool = field(
+        default=True,
+        metadata={
+            "help": "Whether to log a sample of (prompt, completion) pairs every `logging_steps` steps. If `rich` is "
+            "installed, it prints the sample. If `wandb` logging is enabled, it logs it to `wandb`."
+        },
+    )
+
+    report_to: Union[None, str, list[str]] = field(
+        default="wandb", metadata={"help": "The list of integrations to report the results and logs to."}
+    )
+
+    logging_first_step: bool = field(default=False, metadata={"help": "Log the first global_step"})
+    logging_steps: float = field(
+        default=2,
+        metadata={
+            "help": (
+                "Log every X updates steps. Should be an integer or a float in range `[0,1)`. "
+                "If smaller than 1, will be interpreted as ratio of total training steps."
+            )
+        },
+    )

BioReason/bioreason/trainer/grpo_trainer.py

@@ -0,0 +1,905 @@
+# Copyright 2025 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import time
+import textwrap
+import pandas as pd
+from collections import defaultdict
+from typing import Any, Callable, Optional, Union, Sized
+
+import torch
+import torch.utils.data
+import transformers
+from datasets import Dataset, IterableDataset
+from packaging import version
+from transformers import (
+    AriaForConditionalGeneration,
+    AriaProcessor,
+    AutoModelForCausalLM,
+    AutoModelForSequenceClassification,
+    AutoProcessor,
+    AutoTokenizer,
+    GenerationConfig,
+    PreTrainedModel,
+    PreTrainedTokenizerBase,
+    Qwen2VLForConditionalGeneration,
+    Qwen2_5_VLForConditionalGeneration,
+    Trainer,
+    TrainerCallback,
+    is_wandb_available,
+)
+from transformers.integrations.deepspeed import is_deepspeed_zero3_enabled
+from transformers.utils import is_peft_available
+
+from trl.data_utils import apply_chat_template, is_conversational, maybe_apply_chat_template
+from trl.models import create_reference_model, prepare_deepspeed, unwrap_model_for_generation
+from trl.trainer.grpo_config import GRPOConfig
+from trl.trainer.utils import generate_model_card, get_comet_experiment_url
+# from trl import GRPOTrainer
+
+from accelerate.utils import is_peft_model, set_seed, gather_object
+import PIL.Image
+
+import copy
+from torch.utils.data import Sampler
+import warnings
+
+if is_peft_available():
+    from peft import PeftConfig, get_peft_model, prepare_model_for_kbit_training
+
+if is_wandb_available():
+    import wandb
+
+from bioreason.dna_modules.dna_module import DNABaseModule
+from bioreason.trainer import DNALLMGRPOConfig
+# What we call a reward function is a callable that takes a list of prompts and completions and returns a list of
+# rewards. When it's a string, it's a model ID, so it's loaded as a pretrained model.
+RewardFunc = Union[str, PreTrainedModel, Callable[[list, list], list[float]]]
+
+
+class RepeatRandomSampler(Sampler):
+    """
+    Sampler that repeats the indices of a dataset in a structured manner.
+
+    Args:
+        data_source (`Sized`):
+            Dataset to sample from.
+        mini_repeat_count (`int`):
+            Number of times to repeat each index per batch.
+        batch_size (`int`, *optional*, defaults to `1`):
+            Number of unique indices per batch.
+        repeat_count (`int`, *optional*, defaults to `1`):
+            Number of times to repeat the full sampling process.
+        seed (`int` or `None`, *optional*, defaults to `None`):
+            Random seed for reproducibility.
+    """
+
+    def __init__(
+        self,
+        data_source: Sized,
+        mini_repeat_count: int,
+        batch_size: int = 1,
+        repeat_count: int = 1,
+        seed: Optional[int] = None,
+    ):
+        self.data_source = data_source
+        self.mini_repeat_count = mini_repeat_count
+        self.batch_size = batch_size
+        self.repeat_count = repeat_count
+        self.num_samples = len(data_source)
+        self.seed = seed
+        self.generator = torch.Generator()
+        if seed is not None:
+            self.generator.manual_seed(seed)
+
+    def __iter__(self):
+        indexes = torch.randperm(self.num_samples, generator=self.generator).tolist()
+        indexes = [indexes[i : i + self.batch_size] for i in range(0, len(indexes), self.batch_size)]
+        indexes = [chunk for chunk in indexes if len(chunk) == self.batch_size]
+
+        for chunk in indexes:
+            for _ in range(self.repeat_count):
+                for index in chunk:
+                    for _ in range(self.mini_repeat_count):
+                        yield index
+
+    def __len__(self) -> int:
+        return self.num_samples * self.mini_repeat_count * self.repeat_count
+
+
+class DNALLMGRPOTrainer(Trainer):
+    """
+    Trainer for the Group Relative Policy Optimization (GRPO) method. This algorithm was initially proposed in the
+    paper [DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models](https://huggingface.co/papers/2402.03300).
+
+    Example:
+
+    ```python
+    from datasets import load_dataset
+    from trl import GRPOTrainer
+
+    dataset = load_dataset("trl-lib/tldr", split="train")
+
+    trainer = GRPOTrainer(
+        model="Qwen/Qwen2-0.5B-Instruct",
+        reward_funcs="weqweasdas/RM-Gemma-2B",
+        train_dataset=dataset,
+    )
+
+    trainer.train()
+    ```
+
+    Args:
+        model (`Union[str, PreTrainedModel]`):
+            Model to be trained. Can be either:
+
+            - A string, being the *model id* of a pretrained model hosted inside a model repo on huggingface.co, or
+              a path to a *directory* containing model weights saved using
+              [`~transformers.PreTrainedModel.save_pretrained`], e.g., `'./my_model_directory/'`. The model is
+              loaded using [`~transformers.AutoModelForCausalLM.from_pretrained`] with the keywork arguments
+              in `args.model_init_kwargs`.
+            - A [`~transformers.PreTrainedModel`] object. Only causal language models are supported.
+        reward_funcs (`Union[RewardFunc, list[RewardFunc]]`):
+            Reward functions to be used for computing the rewards. To compute the rewards, we call all the reward
+            functions with the prompts and completions and sum the rewards. Can be either:
+
+            - A single reward function, such as:
+                - A string: The *model ID* of a pretrained model hosted inside a model repo on huggingface.co, or a
+                path to a *directory* containing model weights saved using
+                [`~transformers.PreTrainedModel.save_pretrained`], e.g., `'./my_model_directory/'`. The model is loaded
+                using [`~transformers.AutoModelForSequenceClassification.from_pretrained`] with `num_labels=1` and the
+                keyword arguments in `args.model_init_kwargs`.
+                - A [`~transformers.PreTrainedModel`] object: Only sequence classification models are supported.
+                - A custom reward function: The function is provided with the prompts and the generated completions,
+                  plus any additional columns in the dataset. It should return a list of rewards. For more details, see
+                  [Using a custom reward function](#using-a-custom-reward-function).
+            - A list of reward functions, where each item can independently be any of the above types. Mixing different
+            types within the list (e.g., a string model ID and a custom reward function) is allowed.
+        args ([`GRPOConfig`], *optional*, defaults to `None`):
+            Configuration for this trainer. If `None`, a default configuration is used.
+        train_dataset ([`~datasets.Dataset`] or [`~datasets.IterableDataset`]):
+            Dataset to use for training. It must include a column `"prompt"`. Any additional columns in the dataset is
+            ignored. The format of the samples can be either:
+
+            - [Standard](dataset_formats#standard): Each sample contains plain text.
+            - [Conversational](dataset_formats#conversational): Each sample contains structured messages (e.g., role
+              and content).
+        eval_dataset ([`~datasets.Dataset`], [`~datasets.IterableDataset`] or `dict[str, Union[Dataset, IterableDataset]]`):
+            Dataset to use for evaluation. It must meet the same requirements as `train_dataset`.
+        processing_class ([`~transformers.PreTrainedTokenizerBase`], *optional*, defaults to `None`):
+            Processing class used to process the data. The padding side must be set to "left". If `None`, the
+            processing class is loaded from the model's name with [`~transformers.AutoTokenizer.from_pretrained`].
+        reward_processing_classes (`Union[PreTrainedTokenizerBase, list[PreTrainedTokenizerBase]]`, *optional*, defaults to `None`):
+            Processing classes corresponding to the reward functions specified in `reward_funcs`. Can be either:
+
+            - A single processing class: Used when `reward_funcs` contains only one reward function.
+            - A list of processing classes: Must match the order and length of the reward functions in `reward_funcs`.
+            If set to `None`, or if an element of the list corresponding to a [`~transformers.PreTrainedModel`] is
+            `None`, the tokenizer for the model is automatically loaded using [`~transformers.AutoTokenizer.from_pretrained`].
+            For elements in `reward_funcs` that are custom reward functions (not [`~transformers.PreTrainedModel`]),
+            the corresponding entries in `reward_processing_classes` are ignored.
+        callbacks (list of [`~transformers.TrainerCallback`], *optional*, defaults to `None`):
+            List of callbacks to customize the training loop. Will add those to the list of default callbacks
+            detailed in [here](https://huggingface.co/docs/transformers/main_classes/callback).
+
+            If you want to remove one of the default callbacks used, use the [`~transformers.Trainer.remove_callback`]
+            method.
+        optimizers (`tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]`, *optional*, defaults to `(None, None)`):
+            A tuple containing the optimizer and the scheduler to use. Will default to an instance of [`AdamW`] on your
+            model and a scheduler given by [`get_linear_schedule_with_warmup`] controlled by `args`.
+        peft_config ([`~peft.PeftConfig`], *optional*, defaults to `None`):
+            PEFT configuration used to wrap the model. If `None`, the model is not wrapped.
+    """
+
+    def __init__(
+        self,
+        model: Union[str, PreTrainedModel],
+        reward_funcs: Union[RewardFunc, list[RewardFunc]],
+        args: DNALLMGRPOConfig = None,
+        dna_module: DNABaseModule = None,
+        train_dataset: Optional[Union[Dataset, IterableDataset]] = None,
+        eval_dataset: Optional[Union[Dataset, IterableDataset, dict[str, Union[Dataset, IterableDataset]]]] = None,
+        processing_class: Optional[PreTrainedTokenizerBase] = None,
+        reward_processing_classes: Optional[Union[PreTrainedTokenizerBase, list[PreTrainedTokenizerBase]]] = None,
+        callbacks: Optional[list[TrainerCallback]] = None,
+        optimizers: tuple[Optional[torch.optim.Optimizer], Optional[torch.optim.lr_scheduler.LambdaLR]] = (None, None),
+        peft_config: Optional["PeftConfig"] = None,
+        freeze_dna_modules: Optional[bool] = False,
+        attn_implementation: str = "flash_attention_2",
+        torch_dtype: str = "bfloat16",
+        **kwargs,
+    ):
+        # Args
+        if args is None:
+            model_name = model if isinstance(model, str) else model.config._name_or_path
+            model_name = model_name.split("/")[-1]
+            args = GRPOConfig(f"{model_name}-GRPO")
+        
+        self.dna_module = dna_module
+
+        # Models
+        # Trained model
+        model_init_kwargs = args.model_init_kwargs or {}
+        # FIXME
+        # Remember to modify it in the invernvl
+        model_init_kwargs["attn_implementation"] = attn_implementation
+        if model_init_kwargs.get("torch_dtype") is None:
+            model_init_kwargs["torch_dtype"] = torch_dtype
+        
+        assert not isinstance(model, str), "model must NOT be a string in the current implementation"
+
+        torch_dtype = model_init_kwargs.get("torch_dtype")
+        if isinstance(torch_dtype, torch.dtype) or torch_dtype == "auto" or torch_dtype is None:
+            pass  # torch_dtype is already a torch.dtype or "auto" or None
+        elif isinstance(torch_dtype, str):  # it's a str, but not "auto"
+            torch_dtype = getattr(torch, torch_dtype)
+        else:
+            raise ValueError(
+                "Invalid `torch_dtype` passed to `GRPOConfig`. Expected either 'auto' or a string representing "
+                f"a `torch.dtype` (e.g., 'float32'), but got {torch_dtype}."
+            )
+        # Disable caching if gradient checkpointing is enabled (not supported)
+        model_init_kwargs["use_cache"] = (
+            False if args.gradient_checkpointing else model_init_kwargs.get("use_cache")
+        )
+
+        # LoRA
+        self.dna_modules_keywords = self.dna_module.get_dnallm_modules_keywords()
+        if peft_config is not None:
+            print("Applying LoRA...")
+            def find_all_linear_names(model, multimodal_keywords):
+                cls = torch.nn.Linear
+                lora_module_names = set()
+                for name, module in model.named_modules():
+                    print('name:', name, 'module:', module)
+                    # LoRA is not applied to the DNA modules
+                    if any(mm_keyword in name for mm_keyword in multimodal_keywords):
+                        continue
+                    if isinstance(module, cls):
+                        lora_module_names.add(name)
+                for m in lora_module_names:  # needed for 16-bit
+                    if "embed_tokens" in m:
+                        lora_module_names.remove(m)
+                return list(lora_module_names)
+            target_modules = find_all_linear_names(model, self.dna_modules_keywords)
+            peft_config.target_modules = target_modules
+            model = prepare_model_for_kbit_training(model)
+            model = get_peft_model(model, peft_config)
+
+        # Freeze DNA modules
+        if freeze_dna_modules:
+            print("Freezing DNA modules...")
+            for p in model.dna_model.parameters():
+                p.requires_grad = False
+        
+        # Make projection layer trainable
+        for p in model.dna_projection.parameters():
+            p.required_grad = True
+
+        # Compute the number of trainable parameters and print the parameter that is trainable
+        trainable_params = [p for p in model.parameters() if p.requires_grad]
+        total_params = sum(p.numel() for p in trainable_params)
+        # for n, p in model.named_parameters():
+        #     if p.requires_grad:
+        #         print(n, p.shape)
+        print(f"Total trainable parameters: {total_params}")
+
+        # Enable gradient checkpointing if requested
+        if args.gradient_checkpointing:
+            model = self._enable_gradient_checkpointing(model, args)
+
+        # Reference model
+        self.beta = args.beta
+        if self.beta == 0.0:
+            # If beta is 0.0, the reference model is not needed
+            self.ref_model = None
+        elif is_deepspeed_zero3_enabled():
+            self.ref_model = model_cls.from_pretrained(model_id, **model_init_kwargs)
+        elif is_peft_model(model):
+            # If PEFT is used, the reference model is not needed since the adapter can be disabled
+            # to revert to the initial model.
+            self.ref_model = None
+        else:
+            # If PEFT configuration is not provided, create a reference model based on the initial model.
+            self.ref_model = create_reference_model(model)
+
+        # Processing class
+        if processing_class is None:
+            processing_cls = self.dna_module.get_processing_class()
+
+            #if isinstance(model.text_model)
+            processing_class = processing_cls(tokenizer=model.text_tokenizer, dna_tokenizer=model.dna_tokenizer)
+            # print(model.tokenizer.chat_template)
+            for component, processing_keyword in self.dna_module.get_custom_processing_keywords():
+                if processing_keyword in kwargs:
+                    # If we cannot find component in processing_class, return the processing_class itself
+                    processing_component = getattr(processing_class, component, processing_class)
+                    setattr(processing_component, processing_keyword, kwargs[processing_keyword])
+            if getattr(processing_class, "tokenizer",  None) is not None:
+                pad_token_id = processing_class.tokenizer.pad_token_id
+                processing_class.pad_token_id = pad_token_id
+                processing_class.eos_token_id = processing_class.tokenizer.eos_token_id
+            else:
+                assert isinstance(processing_class, PreTrainedTokenizerBase), "processing_class must be an instance of PreTrainedTokenizerBase if it has no tokenizer attribute"
+                pad_token_id = processing_class.pad_token_id
+
+        self.dna_module.post_model_init(model, processing_class)
+        self.dna_module.post_model_init(self.ref_model, processing_class)
+
+        # Reward functions
+        if not isinstance(reward_funcs, list):
+            reward_funcs = [reward_funcs]
+        for i, reward_func in enumerate(reward_funcs):
+            if isinstance(reward_func, str):
+                reward_funcs[i] = AutoModelForSequenceClassification.from_pretrained(
+                    reward_func, num_labels=1, **model_init_kwargs
+                )
+        self.reward_funcs = reward_funcs
+
+        # Reward processing class
+        if reward_processing_classes is None:
+            reward_processing_classes = [None] * len(reward_funcs)
+        elif not isinstance(reward_processing_classes, list):
+            reward_processing_classes = [reward_processing_classes]
+        else:
+            if len(reward_processing_classes) != len(reward_funcs):
+                raise ValueError("The number of reward processing classes must match the number of reward functions.")
+
+        for i, (reward_processing_class, reward_func) in enumerate(zip(reward_processing_classes, reward_funcs)):
+            if isinstance(reward_func, PreTrainedModel):
+                if reward_processing_class is None:
+                    reward_processing_class = AutoTokenizer.from_pretrained(reward_func.config._name_or_path)
+                if reward_processing_class.pad_token_id is None:
+                    reward_processing_class.pad_token = reward_processing_class.eos_token
+                # The reward model computes the reward for the latest non-padded token in the input sequence.
+                # So it's important to set the pad token ID to the padding token ID of the processing class.
+                reward_func.config.pad_token_id = reward_processing_class.pad_token_id
+                reward_processing_classes[i] = reward_processing_class
+        self.reward_processing_classes = reward_processing_classes
+
+        # Data collator
+        def data_collator(features):  # No data collation is needed in GRPO
+            return features
+
+        # Training arguments
+        self.max_prompt_length = args.max_prompt_length
+        self.max_prompt_length = None
+        if args.max_prompt_length is not None:
+            warnings.warn("Setting max_prompt_length is currently not supported, it has been set to None")
+
+        self.max_completion_length = args.max_completion_length  # = |o_i| in the GRPO paper
+        self.num_generations = args.num_generations  # = G in the GRPO paper
+        self.generation_config = GenerationConfig(
+            max_new_tokens=self.max_completion_length,
+            do_sample=True,  
+            temperature=0.6,
+            top_p=0.95,
+            top_k=20,
+            pad_token_id=pad_token_id,
+        )
+        if hasattr(self.dna_module, "get_eos_token_id"): # For InternVL
+            self.generation_config.eos_token_id = self.dna_module.get_eos_token_id(processing_class)
+        self.beta = args.beta
+        self.epsilon_low = args.epsilon
+        self.epsilon_high = args.epsilon_high if args.epsilon_high is not None else args.epsilon
+
+        # Multi-step
+        self.num_iterations = args.num_iterations  # = 𝜇 in the GRPO paper
+        # Tracks the number of iterations (forward + backward passes), including those within a gradient accumulation cycle
+        self._step = 0
+        # Buffer the batch to reuse generated outputs across multiple updates
+        self._buffered_inputs = [None] * args.gradient_accumulation_steps
+
+        # The trainer estimates the number of FLOPs (floating-point operations) using the number of elements in the
+        # input tensor associated with the key "input_ids". However, in GRPO, the sampled data does not include the
+        # "input_ids" key. Instead, the available keys is "prompt". As a result, the trainer issues the warning:
+        # "Could not estimate the number of tokens of the input, floating-point operations will not be computed." To
+        # suppress this warning, we set the "estimate_tokens" key in the model's "warnings_issued" dictionary to True.
+        # This acts as a flag to indicate that the warning has already been issued.
+        model.warnings_issued["estimate_tokens"] = True
+
+        # Initialize the metrics
+        self._metrics = defaultdict(list)
+        self.log_completions = args.log_completions
+
+        super().__init__(
+            model=model,
+            args=args,
+            data_collator=data_collator,
+            train_dataset=train_dataset,
+            eval_dataset=eval_dataset,
+            processing_class=processing_class,
+            callbacks=callbacks,
+            optimizers=optimizers,
+        )
+
+        # Check if the per_device_train/eval_batch_size * num processes can be divided by the number of generations
+        num_processes = self.accelerator.num_processes
+        global_batch_size = args.per_device_train_batch_size * num_processes
+        possible_values = [n_gen for n_gen in range(2, global_batch_size + 1) if (global_batch_size) % n_gen == 0]
+        if self.num_generations not in possible_values:
+            raise ValueError(
+                f"The global train batch size ({num_processes} x {args.per_device_train_batch_size}) must be evenly "
+                f"divisible by the number of generations per prompt ({self.num_generations}). Given the current train "
+                f"batch size, the valid values for the number of generations are: {possible_values}."
+            )
+        if self.args.eval_strategy != "no":
+            global_batch_size = args.per_device_eval_batch_size * num_processes
+            possible_values = [n_gen for n_gen in range(2, global_batch_size + 1) if (global_batch_size) % n_gen == 0]
+            if self.num_generations not in possible_values:
+                raise ValueError(
+                    f"The global eval batch size ({num_processes} x {args.per_device_eval_batch_size}) must be evenly "
+                    f"divisible by the number of generations per prompt ({self.num_generations}). Given the current "
+                    f"eval batch size, the valid values for the number of generations are: {possible_values}."
+                )
+
+        # Ensure each process receives a unique seed to prevent duplicate completions when generating with
+        # transformers if num_generations exceeds per_device_train_batch_size. We could skip it if we use vLLM, but
+        # it's safer to set it in all cases.
+        set_seed(args.seed, device_specific=True)
+
+        # Gradient accumulation requires scaled loss. Normally, loss scaling in the parent class depends on whether the
+        # model accepts loss-related kwargs. Since we compute our own loss, this check is irrelevant. We set
+        # self.model_accepts_loss_kwargs to False to enable scaling.
+        self.model_accepts_loss_kwargs = False
+
+        if self.ref_model is not None:
+            # if self.is_deepspeed_enabled:
+            if is_deepspeed_zero3_enabled():
+                self.ref_model = prepare_deepspeed(self.ref_model, self.accelerator)
+            else:
+                self.ref_model = self.accelerator.prepare_model(self.ref_model, evaluation_mode=True)
+
+        for i, reward_func in enumerate(self.reward_funcs):
+            if isinstance(reward_func, PreTrainedModel):
+                self.reward_funcs[i] = self.accelerator.prepare_model(reward_func, evaluation_mode=True)
+
+    def _enable_gradient_checkpointing(self, model: PreTrainedModel, args: GRPOConfig) -> PreTrainedModel:
+        """Enables gradient checkpointing for the model."""
+        # Ensure use_cache is disabled
+        model.config.use_cache = False
+
+        # Enable gradient checkpointing on the base model for PEFT
+        if is_peft_model(model):
+            model.base_model.gradient_checkpointing_enable()
+        # Enable gradient checkpointing for non-PEFT models
+        else:
+            if getattr(model, "language_model", None) is not None:
+                # For InternVL; these operations are copied from the original training script of InternVL
+                model.language_model.config.use_cache = False
+                model.dna_model.gradient_checkpointing = True
+                model.dna_model.encoder.gradient_checkpointing = True
+                model.language_model._set_gradient_checkpointing()
+                # This line is necessary, otherwise the `model.gradient_checkpointing_enable()` will be executed during the training process, leading to an error since InternVL does not support this operation.
+                args.gradient_checkpointing = False
+            else:
+                model.gradient_checkpointing_enable()
+
+        gradient_checkpointing_kwargs = args.gradient_checkpointing_kwargs or {}
+        use_reentrant = (
+            "use_reentrant" not in gradient_checkpointing_kwargs or gradient_checkpointing_kwargs["use_reentrant"]
+        )
+
+        if use_reentrant:
+            model.enable_input_require_grads()
+
+        return model
+    
+    def _set_signature_columns_if_needed(self):
+        # If `self.args.remove_unused_columns` is True, non-signature columns are removed.
+        # By default, this method sets `self._signature_columns` to the model's expected inputs.
+        # In GRPOTrainer, we preprocess data, so using the model's signature columns doesn't work.
+        # Instead, we set them to the columns expected by the `training_step` method, hence the override.
+        if self._signature_columns is None:
+            self._signature_columns = ["prompt"]
+
+
+    # Get the per-token log probabilities for the completions for the model and the reference model
+    def _get_per_token_logps(self, model, input_ids, attention_mask, **custom_multimodal_inputs):
+        logits = model(input_ids=input_ids, attention_mask=attention_mask, **custom_multimodal_inputs).logits  # (B, L, V)
+        logits = logits[:, :-1, :]  # (B, L-1, V), exclude the last logit: it corresponds to the next token pred
+        input_ids = input_ids[:, 1:]  # (B, L-1), exclude the first input ID since we don't have logits for it
+        # Compute the log probabilities for the input tokens. Use a loop to reduce memory peak.
+        per_token_logps = []
+        for logits_row, input_ids_row in zip(logits, input_ids):
+            log_probs = logits_row.log_softmax(dim=-1)
+            token_log_prob = torch.gather(log_probs, dim=1, index=input_ids_row.unsqueeze(1)).squeeze(1)
+            per_token_logps.append(token_log_prob)
+        return torch.stack(per_token_logps)
+
+
+    def _prepare_inputs(self, inputs):
+        # Simple pass-through, just like original
+        return inputs
+
+    def _get_key_from_inputs(self, x, key):
+        ele = x.get(key, None)
+        assert ele is not None, f"The key {key} is not found in the input"
+        if isinstance(ele, list):
+            return [e for e in ele]
+        else:
+            return [ele]
+
+    def _generate_and_score_completions(self, inputs: dict[str, Union[torch.Tensor, Any]], model) -> dict[str, Union[torch.Tensor, Any]]:
+        device = self.accelerator.device
+        prompts = [x["prompt"] for x in inputs]
+        prompts_text = self.dna_module.prepare_prompt(self.processing_class, inputs)
+        # Handle both pre-loaded images and image paths
+        batch_dna_sequences = []
+        print("_generate_and_score_completions (GRPO):")
+        for x in inputs:
+            #print('---')
+            #print(x)
+            if 'dna_sequences' in x:
+                dnas = self._get_key_from_inputs(x, "dna_sequences")
+
+            for dna in dnas:
+                # clean if desired
+                pass
+            batch_dna_sequences.append(dnas)
+            # NOTE: typically appends dna, so dna_sequences is all the dna in one list
+            # odd. trying this instead
+                
+
+        prompt_inputs = self.dna_module.prepare_model_inputs(
+            self.processing_class,
+            model,
+            prompts_text,
+            batch_dna_sequences,
+            return_tensors="pt",
+            padding=True,
+            padding_side="left",
+            add_special_tokens=False,
+        )
+
+        prompt_inputs = super()._prepare_inputs(prompt_inputs)
+        prompt_ids, prompt_mask = prompt_inputs["input_ids"], prompt_inputs["attention_mask"]
+
+        # max_prompt_length is not supported yet
+        # if self.max_prompt_length is not None:
+        #     prompt_ids = prompt_ids[:, -self.max_prompt_length :]
+        #     prompt_inputs["input_ids"] = prompt_ids
+        #     prompt_mask = prompt_mask[:, -self.max_prompt_length :]
+        #     prompt_inputs["attention_mask"] = prompt_mask
+
+        # Generate completions
+        start = time.time()
+        with unwrap_model_for_generation(model, self.accelerator) as unwrapped_model:
+            kwargs = {k: v for k, v in prompt_inputs.items() if k not in self.dna_module.get_non_generate_params()}
+            generate_returned_result = unwrapped_model.generate(
+                **kwargs,
+                generation_config=self.generation_config
+            )
+            end = time.time()
+            print(f"Generation time: {end - start:.9f} seconds")
+            prompt_length = prompt_ids.size(1)
+            if not self.dna_module.is_embeds_input():
+                prompt_completion_ids = generate_returned_result
+                prompt_ids = prompt_completion_ids[:, :prompt_length]
+                completion_ids = prompt_completion_ids[:, prompt_length:]
+            else:
+                # In this case, the input of the LLM backbone is the embedding of the combination of the image and text prompt
+                # So the returned result of the `generate` method only contains the completion ids
+                completion_ids = generate_returned_result
+                prompt_completion_ids = torch.cat([prompt_ids, completion_ids], dim=1)
+
+        # Mask everything after the first EOS token
+        # print('completion:', completion_ids)
+        # print('generate_returned_result', generate_returned_result, generate_returned_result.shape)
+        # print('prompt_inputs["input_ids"]', prompt_inputs["input_ids"], prompt_inputs["input_ids"].shape)
+        # print('prompt_ids', prompt_ids, prompt_ids.shape)
+        # print('prompt_length', prompt_length)
+        # print('prompt_completion_ids', prompt_completion_ids, prompt_completion_ids.shape)
+        is_eos = completion_ids == self.processing_class.eos_token_id
+        eos_idx = torch.full((is_eos.size(0),), is_eos.size(1), dtype=torch.long, device=device)
+        eos_idx[is_eos.any(dim=1)] = is_eos.int().argmax(dim=1)[is_eos.any(dim=1)]
+        sequence_indices = torch.arange(is_eos.size(1), device=device).expand(is_eos.size(0), -1)
+        completion_mask = (sequence_indices <= eos_idx.unsqueeze(1)).int()
+
+        # Concatenate prompt_mask with completion_mask for logit computation
+        attention_mask = torch.cat([prompt_mask, completion_mask], dim=1)  # (B, P+C)
+
+        # Get the multimodal inputs
+        multimodal_keywords = self.dna_module.get_custom_multimodal_keywords()
+        multimodal_inputs = {k: prompt_inputs[k] if k in prompt_inputs else None for k in multimodal_keywords}
+        with torch.no_grad():
+            # When using num_iterations == 1, old_per_token_logps == per_token_logps, so we can skip its
+            # computation here, and use per_token_logps.detach() instead.
+            if self.num_iterations > 1:
+                old_per_token_logps = self._get_per_token_logps(
+                    model, prompt_completion_ids, attention_mask, **multimodal_inputs
+                )
+                old_per_token_logps = old_per_token_logps[:, prompt_length - 1:]
+            else:
+                old_per_token_logps = None
+
+            if self.beta == 0.0:
+                ref_per_token_logps = None
+            elif self.ref_model is not None:
+                ref_per_token_logps = self._get_per_token_logps(
+                    self.ref_model, prompt_completion_ids, attention_mask, **multimodal_inputs
+                )
+            else:
+                with self.accelerator.unwrap_model(model).disable_adapter():
+                    ref_per_token_logps = self._get_per_token_logps(
+                        model, prompt_completion_ids, attention_mask, **multimodal_inputs
+                    )
+        if ref_per_token_logps is not None:
+            ref_per_token_logps = ref_per_token_logps[:, prompt_length - 1:]
+
+        # Decode the generated completions
+        completions_text = self.processing_class.batch_decode(completion_ids, skip_special_tokens=True)
+        if is_conversational(inputs[0]):
+            completions = [[{"role": "assistant", "content": completion}] for completion in completions_text]
+        else:
+            completions = completions_text
+        # Compute the rewards
+        # No need to duplicate prompts as we're not generating multiple completions per prompt
+        print("Reward calculation...")
+        rewards_per_func = torch.zeros(len(prompts), len(self.reward_funcs), device=device)
+        for i, (reward_func, reward_processing_class) in enumerate(
+            zip(self.reward_funcs, self.reward_processing_classes)
+        ):
+            if isinstance(reward_func, PreTrainedModel):
+                if is_conversational(inputs[0]):
+                    messages = [{"messages": p + c} for p, c in zip(prompts, completions)]
+                    texts = [apply_chat_template(x, reward_processing_class)["text"] for x in messages]
+                else:
+                    texts = [p + c for p, c in zip(prompts, completions)]
+                reward_inputs = reward_processing_class(
+                    texts, return_tensors="pt", padding=True, padding_side="right", add_special_tokens=False
+                )
+                reward_inputs = super()._prepare_inputs(reward_inputs)
+                with torch.inference_mode():
+                    rewards_per_func[:, i] = reward_func(**reward_inputs).logits[:, 0]  # Shape (B*G,)
+            else:
+                # Repeat all input columns (but "prompt" and "completion") to match the number of generations
+                reward_kwargs = {key: [] for key in inputs[0].keys() if key not in ["prompt", "completion"]}
+                for key in reward_kwargs:
+                    for example in inputs:
+                        # No need to duplicate prompts as we're not generating multiple completions per prompt
+                        # reward_kwargs[key].extend([example[key]] * self.num_generations)
+                        reward_kwargs[key].extend([example[key]])
+                output_reward_func = reward_func(prompts=prompts, completions=completions, **reward_kwargs)
+                rewards_per_func[:, i] = torch.tensor(output_reward_func, dtype=torch.float32, device=device)
+
+        # Gather rewards across processes
+        rewards_per_func = self.accelerator.gather(rewards_per_func)
+        
+        # Sum the rewards from all reward functions
+        rewards = rewards_per_func.sum(dim=1)
+        
+        # Compute grouped-wise rewards
+        # Each group consists of num_generations completions for the same prompt
+        mean_grouped_rewards = rewards.view(-1, self.num_generations).mean(dim=1)
+        std_grouped_rewards = rewards.view(-1, self.num_generations).std(dim=1)
+        
+        # Normalize the rewards to compute the advantages
+        mean_grouped_rewards = mean_grouped_rewards.repeat_interleave(self.num_generations, dim=0)
+        std_grouped_rewards = std_grouped_rewards.repeat_interleave(self.num_generations, dim=0)
+        advantages = (rewards - mean_grouped_rewards) / (std_grouped_rewards + 1e-4)
+        
+        # Get only the local slice of advantages
+        process_slice = slice(
+            self.accelerator.process_index * len(prompts),
+            (self.accelerator.process_index + 1) * len(prompts),
+        )
+        advantages = advantages[process_slice]
+
+        # Log the metrics
+        print("Logging metrics...")
+        completion_length = self.accelerator.gather_for_metrics(completion_mask.sum(1)).float().mean().item()
+        self._metrics["completion_length"].append(completion_length)
+
+        reward_per_func = self.accelerator.gather_for_metrics(rewards_per_func).mean(0)
+        for i, reward_func in enumerate(self.reward_funcs):
+            if isinstance(reward_func, PreTrainedModel):
+                reward_func_name = reward_func.config._name_or_path.split("/")[-1]
+            else:
+                reward_func_name = reward_func.__name__
+            self._metrics[f"rewards/{reward_func_name}"].append(reward_per_func[i].item())
+
+        self._metrics["reward"].append(self.accelerator.gather_for_metrics(rewards).mean().item())
+
+        self._metrics["reward_std"].append(self.accelerator.gather_for_metrics(std_grouped_rewards).mean().item())
+
+        print(self.log_completions, self.state.global_step, self.args.logging_steps, self.args.report_to)
+        if (
+            self.log_completions
+            and self.state.global_step % self.args.logging_steps == 0
+            and "wandb" in self.args.report_to
+        ):
+            timestamp = time.time()
+
+            # Get the length of one of the other arrays
+            num_items = len(gather_object(prompts_text))
+
+            table = {
+                "step": [f"{self.state.global_step}_{timestamp}"] * num_items,  # Repeat to match length
+                "prompt": gather_object(prompts_text),
+                "completion": gather_object(completions_text),
+                "reward": rewards.tolist(),
+            }
+            df = pd.DataFrame(table)
+
+            if wandb.run is not None and self.accelerator.is_main_process:
+                wandb.log({f"completions_{self.state.global_step}_{timestamp}": wandb.Table(dataframe=df)})
+
+        return {
+            "prompt_ids": prompt_ids,
+            "prompt_mask": prompt_mask,
+            "completion_ids": completion_ids,
+            "completion_mask": completion_mask,
+            "old_per_token_logps": old_per_token_logps,
+            "ref_per_token_logps": ref_per_token_logps,
+            "advantages": advantages,
+            "multimodal_inputs": multimodal_inputs
+        }
+
+    def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
+        if return_outputs:
+            raise ValueError("The GRPOTrainer does not support returning outputs")
+    
+        # Check if we need to generate new completions or use buffered ones
+        print("index 1")
+        if self.state.global_step % self.num_iterations == 0:
+            inputs = self._generate_and_score_completions(inputs, model)
+            self._buffered_inputs[self._step % self.args.gradient_accumulation_steps] = inputs
+        else:
+            inputs = self._buffered_inputs[self._step % self.args.gradient_accumulation_steps]
+        self._step += 1
+        
+        print("index 2")
+        # Get the prepared inputs
+        prompt_ids, prompt_mask = inputs["prompt_ids"], inputs["prompt_mask"]
+        completion_ids, completion_mask = inputs["completion_ids"], inputs["completion_mask"]
+        multimodal_inputs = inputs["multimodal_inputs"]
+        
+        # Concatenate for full sequence
+        input_ids = torch.cat([prompt_ids, completion_ids], dim=1)
+        attention_mask = torch.cat([prompt_mask, completion_mask], dim=1)
+        print("index 3")
+        # Get the current policy's log probabilities
+        
+        print("index 4")
+        per_token_logps = self._get_per_token_logps(model, input_ids, attention_mask, **multimodal_inputs)
+        # Get rid of the prompt (-1 because of the shift done in get_per_token_logps)
+        per_token_logps = per_token_logps[:, prompt_ids.size(1) - 1:]
+
+        # Get the advantages from inputs
+        advantages = inputs["advantages"]
+        print("index 5")
+        # When using num_iterations == 1, old_per_token_logps == per_token_logps, so we can skip its computation
+        # and use per_token_logps.detach() instead
+        old_per_token_logps = inputs["old_per_token_logps"] if self.num_iterations > 1 else per_token_logps.detach()
+
+        # Compute the policy ratio and clipped version
+        coef_1 = torch.exp(per_token_logps - old_per_token_logps)
+        coef_2 = torch.clamp(coef_1, 1 - self.epsilon_low, 1 + self.epsilon_high)
+        per_token_loss1 = coef_1 * advantages.unsqueeze(1)
+        per_token_loss2 = coef_2 * advantages.unsqueeze(1)
+        per_token_loss = -torch.min(per_token_loss1, per_token_loss2)
+        print("index 6")
+        # Add KL penalty if beta > 0
+        if self.beta > 0:
+            ref_per_token_logps = inputs["ref_per_token_logps"]
+            per_token_kl = torch.exp(ref_per_token_logps - per_token_logps) - (ref_per_token_logps - per_token_logps) - 1
+            per_token_loss = per_token_loss + self.beta * per_token_kl
+
+            # Log KL divergence
+            mean_kl = ((per_token_kl * completion_mask).sum(dim=1) / completion_mask.sum(dim=1)).mean()
+            self._metrics["kl"].append(self.accelerator.gather_for_metrics(mean_kl).mean().item())
+
+        # Compute final loss
+        print("Computing final loss...")
+        loss = ((per_token_loss * completion_mask).sum(dim=1) / completion_mask.sum(dim=1)).mean()
+
+        # Log clip ratio
+        is_clipped = (per_token_loss1 < per_token_loss2).float()
+        clip_ratio = (is_clipped * completion_mask).sum() / completion_mask.sum()
+        self._metrics["clip_ratio"].append(self.accelerator.gather_for_metrics(clip_ratio).mean().item())
+
+        return loss
+
+    def log(self, logs: dict[str, float], start_time: Optional[float] = None) -> None:
+        metrics = {key: sum(val) / len(val) for key, val in self._metrics.items()}  # average the metrics
+        logs = {**logs, **metrics}
+        if version.parse(transformers.__version__) >= version.parse("4.47.0.dev0"):
+            super().log(logs, start_time)
+        else:  # transformers<=4.46
+            super().log(logs)
+        self._metrics.clear()
+
+    def create_model_card(
+        self,
+        model_name: Optional[str] = None,
+        dataset_name: Optional[str] = None,
+        tags: Union[str, list[str], None] = None,
+    ):
+        """
+        Creates a draft of a model card using the information available to the `Trainer`.
+
+        Args:
+            model_name (`str` or `None`, *optional*, defaults to `None`):
+                Name of the model.
+            dataset_name (`str` or `None`, *optional*, defaults to `None`):
+                Name of the dataset used for training.
+            tags (`str`, `list[str]` or `None`, *optional*, defaults to `None`):
+                Tags to be associated with the model card.
+        """
+        if not self.is_world_process_zero():
+            return
+
+        if hasattr(self.model.config, "_name_or_path") and not os.path.isdir(self.model.config._name_or_path):
+            base_model = self.model.config._name_or_path
+        else:
+            base_model = None
+
+        tags = tags or []
+        if isinstance(tags, str):
+            tags = [tags]
+
+        if hasattr(self.model.config, "unsloth_version"):
+            tags.append("unsloth")
+
+        citation = textwrap.dedent(
+            """\
+            @article{zhihong2024deepseekmath,
+                title        = {{DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models}},
+                author       = {Zhihong Shao and Peiyi Wang and Qihao Zhu and Runxin Xu and Junxiao Song and Mingchuan Zhang and Y. K. Li and Y. Wu and Daya Guo},
+                year         = 2024,
+                eprint       = {arXiv:2402.03300},
+            """
+        )
+
+        model_card = generate_model_card(
+            base_model=base_model,
+            model_name=model_name,
+            hub_model_id=self.hub_model_id,
+            dataset_name=dataset_name,
+            tags=tags,
+            wandb_url=wandb.run.get_url() if is_wandb_available() and wandb.run is not None else None,
+            comet_url=get_comet_experiment_url(),
+            trainer_name="GRPO",
+            trainer_citation=citation,
+            paper_title="DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models",
+            paper_id="2402.03300",
+        )
+
+        model_card.save(os.path.join(self.args.output_dir, "README.md"))
+
+    def _get_train_sampler(self) -> Sampler:
+        """Returns a sampler that ensures proper data sampling for GRPO training."""
+        effective_batch_size = (
+            self.args.per_device_train_batch_size
+            * self.accelerator.num_processes
+            * self.args.gradient_accumulation_steps
+        )
+        
+        return RepeatRandomSampler(
+            data_source=self.train_dataset,
+            mini_repeat_count=self.num_generations,
+            batch_size=effective_batch_size // self.num_generations,
+            repeat_count=self.num_iterations,
+            seed=self.args.seed,
+        )
+
+    def _get_eval_sampler(self, eval_dataset) -> Sampler:
+        """Returns a sampler for evaluation."""
+        return RepeatRandomSampler(
+            data_source=eval_dataset,
+            mini_repeat_count=self.num_generations,
+            seed=self.args.seed,
+        )

BioReason/bioreason/utils/dna_utils.py

@@ -0,0 +1,12 @@
+from typing import TYPE_CHECKING, Callable, Optional, Union
+
+import numpy as np
+
+from transformers.utils import is_torch_available
+
+if is_torch_available():
+    import torch
+
+DNAInput = Union[
+    str, list[int], np.ndarray, "torch.Tensor", list[str], list[list[int]], list[np.ndarray], list["torch.Tensor"]
+]  # noqa

BioReason/data/Clinvar_Coding.ipynb

@@ -0,0 +1,2481 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "83c9cd1f",
+   "metadata": {},
+   "source": [
+    "## Setup and Data Preparation\n",
+    "\n",
+    "Initial setup steps to prepare the working environment and extract ClinVar data."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "81a36253-9050-4d58-96cd-8238aae51e0e",
+   "metadata": {},
+   "source": [
+    "# ClinVar Coding Variants Data Processing\n",
+    "\n",
+    "This notebook processes ClinVar coding variants data by extracting additional information including gene names, gene IDs, and associated diseases from ClinVar XML records.\n",
+    "\n",
+    "## Overview\n",
+    "\n",
+    "The workflow includes:\n",
+    "1. **Data Extraction**: Filter ClinVar entries from VEP-annotated pathogenic coding variants\n",
+    "2. **XML Processing**: Parse ClinVar XML records to extract gene and disease information\n",
+    "3. **Gene Annotation**: Map gene IDs to gene names using NCBI Entrez utilities\n",
+    "4. **Data Integration**: Combine all information into a comprehensive dataset\n",
+    "\n",
+    "## Requirements\n",
+    "\n",
+    "- Python 3.7+\n",
+    "- pandas library\n",
+    "- xml.etree.ElementTree (built-in)\n",
+    "- NCBI Entrez Direct tools (for gene name mapping)\n",
+    "- Input data: VEP-annotated pathogenic coding variants CSV file\n",
+    "\n",
+    "## Data Structure\n",
+    "\n",
+    "The processing creates a dataset with the following key columns:\n",
+    "- Variant information (chromosome, position, alleles)\n",
+    "- ClinVar ID and significance\n",
+    "- Gene symbols and IDs\n",
+    "- Associated disease/phenotype information"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "cb351234-50a3-4061-81ce-bdce5343e790",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Create working directory for ClinVar data processing\n",
+    "import os\n",
+    "os.makedirs('clinvar', exist_ok=True)\n",
+    "print(\"✅ Created 'clinvar' directory\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "443ccab8-50a1-45ae-950c-8425eb318e93",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "\n",
+    "# Navigate to clinvar directory\n",
+    "os.chdir('clinvar')\n",
+    "print(f\"📁 Current working directory: {os.getcwd()}\")\n",
+    "\n",
+    "with open('vep_pathogenic_coding.csv') as infile, open('clinvar_coding_raw.csv', 'w') as outfile:\n",
+    "    for line in infile:\n",
+    "        if 'ClinVar' in line:\n",
+    "            outfile.write(line)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e1f92675-b85c-4baa-8680-9c3776e04ac9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Extract ClinVar entries from VEP-annotated pathogenic coding variants\n",
+    "# Note: Update the input file path to match your data location\n",
+    "input_file = \"../data/vep_pathogenic_coding.csv\"  # Adjust path as needed\n",
+    "output_file = \"clinvar_coding_raw.csv\"\n",
+    "\n",
+    "# Use shell command to filter ClinVar entries\n",
+    "import subprocess\n",
+    "try:\n",
+    "    result = subprocess.run(\n",
+    "        [\"grep\", \"ClinVar\", input_file],\n",
+    "        capture_output=True,\n",
+    "        text=True,\n",
+    "        check=True\n",
+    "    )\n",
+    "    \n",
+    "    with open(output_file, 'w') as f:\n",
+    "        f.write(result.stdout)\n",
+    "    \n",
+    "    print(f\"✅ Extracted ClinVar entries to {output_file}\")\n",
+    "    print(f\"📊 Found {len(result.stdout.strip().split('\\n'))} ClinVar entries\")\n",
+    "    \n",
+    "except subprocess.CalledProcessError:\n",
+    "    print(f\"❌ Error: Could not find ClinVar entries in {input_file}\")\n",
+    "    print(\"Please ensure the input file exists and contains ClinVar annotations\")\n",
+    "except FileNotFoundError:\n",
+    "    print(f\"❌ Error: Input file {input_file} not found\")\n",
+    "    print(\"Please update the input_file path to point to your VEP-annotated data\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7e560308-135b-4189-9146-ff50845839a4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Extract ClinVar IDs from the filtered data (assuming ID is in column 8)\n",
+    "# Note: Adjust column number if your data structure is different\n",
+    "import pandas as pd\n",
+    "\n",
+    "try:\n",
+    "    # Read the raw ClinVar data to determine structure\n",
+    "    df_temp = pd.read_csv(\"clinvar_coding_raw.csv\")\n",
+    "    print(f\"📋 Data shape: {df_temp.shape}\")\n",
+    "    print(f\"📋 Columns: {list(df_temp.columns)}\")\n",
+    "    \n",
+    "    # Extract ClinVar IDs (adjust column index as needed)\n",
+    "    # Column 8 corresponds to index 7 in Python (0-based)\n",
+    "    if df_temp.shape[1] >= 8:\n",
+    "        clinvar_ids = df_temp.iloc[:, 7]  # 8th column (0-based index 7)\n",
+    "        \n",
+    "        # Save IDs to file\n",
+    "        with open(\"Clinvar_ID.txt\", 'w') as f:\n",
+    "            for id_val in clinvar_ids:\n",
+    "                if pd.notna(id_val):\n",
+    "                    f.write(f\"{id_val}\\n\")\n",
+    "        \n",
+    "        print(f\"✅ Extracted {len(clinvar_ids.dropna())} ClinVar IDs to Clinvar_ID.txt\")\n",
+    "    else:\n",
+    "        print(f\"❌ Error: Expected at least 8 columns, found {df_temp.shape[1]}\")\n",
+    "        \n",
+    "except FileNotFoundError:\n",
+    "    print(\"❌ Error: clinvar_coding_raw.csv not found\")\n",
+    "    print(\"Please run the previous cell first to extract ClinVar data\")\n",
+    "except Exception as e:\n",
+    "    print(f\"❌ Error processing ClinVar data: {e}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "53b0dfd8-8d49-4c3f-adb4-4c6bfbffcfa9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "chmod +x Clinvar_esearch.sh\n",
+    "\n",
+    "## XML Data Retrieval\n",
+    "\n",
+    "**Note**: This step requires creating a shell script (`Clinvar_esearch.sh`) to fetch XML data from NCBI.\n",
+    "\n",
+    "The script should:\n",
+    "1. Read ClinVar IDs from `Clinvar_ID.txt`\n",
+    "2. Use NCBI Entrez Direct tools to fetch XML records\n",
+    "3. Save XML files in a `data/` subdirectory\n",
+    "\n",
+    "Example script content:\n",
+    "```bash\n",
+    "#!/bin/bash\n",
+    "mkdir -p data\n",
+    "while read -r id; do\n",
+    "    esearch -db clinvar -query \"$id\" | efetch -format xml > \"data/${id}.xml\"\n",
+    "    echo \"Downloaded XML for ClinVar ID: $id\"\n",
+    "done < Clinvar_ID.txt\n",
+    "```\n",
+    "\n",
+    "**Prerequisites**: Install NCBI Entrez Direct tools:\n",
+    "- macOS: `brew install brewsci/bio/edirect`\n",
+    "- Linux: Follow NCBI EDirect installation guide"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0755ad6d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Parsing XML for Gene and Disease\n",
+    "\n",
+    "# Make the ClinVar search script executable and run it\n",
+    "# Note: This assumes you have created the Clinvar_esearch.sh script\n",
+    "\n",
+    "import os\n",
+    "import subprocess\n",
+    "\n",
+    "script_path = \"Clinvar_esearch.sh\"\n",
+    "\n",
+    "if os.path.exists(script_path):\n",
+    "    # Make script executable\n",
+    "    os.chmod(script_path, 0o755)\n",
+    "    print(f\"✅ Made {script_path} executable\")\n",
+    "    \n",
+    "    # Optionally run the script (uncomment if you want to execute automatically)\n",
+    "    # print(\"🚀 Running ClinVar XML download script...\")\n",
+    "    # result = subprocess.run([f\"./{script_path}\"], capture_output=True, text=True)\n",
+    "    # if result.returncode == 0:\n",
+    "    #     print(\"✅ XML download completed successfully\")\n",
+    "    # else:\n",
+    "    #     print(f\"❌ Script execution failed: {result.stderr}\")\n",
+    "else:\n",
+    "    print(f\"⚠️ Warning: {script_path} not found\")\n",
+    "    print(\"Please create this script manually to download ClinVar XML data\")\n",
+    "    print(\"See the documentation in the previous cell for script template\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d21a188b-a0dc-4af2-9b71-5a44d8cd4673",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Import required libraries\n",
+    "import pandas as pd\n",
+    "import xml.etree.ElementTree as ET\n",
+    "import json\n",
+    "import os\n",
+    "from pathlib import Path\n",
+    "\n",
+    "print(\"📚 Libraries imported successfully\")\n",
+    "print(f\"📁 Current directory: {os.getcwd()}\")\n",
+    "print(f\"📊 Pandas version: {pd.__version__}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1365615b-ee81-4df0-9fca-df001e9f01d4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Load the raw ClinVar data\n",
+    "try:\n",
+    "    clinvar_raw = pd.read_csv(\"clinvar_coding_raw.csv\")\n",
+    "    print(f\"✅ Loaded ClinVar data: {clinvar_raw.shape[0]} rows, {clinvar_raw.shape[1]} columns\")\n",
+    "    print(f\"📋 Columns: {list(clinvar_raw.columns)[:10]}\")  # Show first 10 columns\n",
+    "    \n",
+    "except FileNotFoundError:\n",
+    "    print(\"❌ Error: clinvar_coding_raw.csv not found\")\n",
+    "    print(\"Please run the data extraction steps first\")\n",
+    "    clinvar_raw = None\n",
+    "except Exception as e:\n",
+    "    print(f\"❌ Error loading data: {e}\")\n",
+    "    clinvar_raw = None"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7144ddf2-abf7-4680-b578-d4bd4b7195ea",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Remove unnecessary columns to streamline the dataset\n",
+    "# Note: Adjust column names based on your actual data structure\n",
+    "\n",
+    "if clinvar_raw is not None:\n",
+    "    columns_to_remove = [\n",
+    "        \"GENOMIC_MUTATION_ID\", \"N_SAMPLES\", \"TOTAL_SAMPLES\", \"FREQ\", \n",
+    "        \"OMIM\", \"PMID\", \"AC\", \"AN\", \"AF\", \"MAF\", \"MAC\"\n",
+    "    ]\n",
+    "    \n",
+    "    # Only remove columns that actually exist in the dataset\n",
+    "    existing_columns = [col for col in columns_to_remove if col in clinvar_raw.columns]\n",
+    "    missing_columns = [col for col in columns_to_remove if col not in clinvar_raw.columns]\n",
+    "    \n",
+    "    if existing_columns:\n",
+    "        clinvar_raw = clinvar_raw.drop(columns=existing_columns)\n",
+    "        print(f\"✅ Removed {len(existing_columns)} columns: {existing_columns}\")\n",
+    "    \n",
+    "    if missing_columns:\n",
+    "        print(f\"ℹ️ Columns not found (skipped): {missing_columns}\")\n",
+    "    \n",
+    "    print(f\"📊 Remaining columns: {clinvar_raw.shape[1]}\")\n",
+    "else:\n",
+    "    print(\"⚠️ Skipping column removal - data not loaded\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "fbffd3cd-7df3-43e2-8d73-01f54e8d1da6",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "<div>\n",
+       "<style scoped>\n",
+       "    .dataframe tbody tr th:only-of-type {\n",
+       "        vertical-align: middle;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe tbody tr th {\n",
+       "        vertical-align: top;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe thead th {\n",
+       "        text-align: right;\n",
+       "    }\n",
+       "</style>\n",
+       "<table border=\"1\" class=\"dataframe\">\n",
+       "  <thead>\n",
+       "    <tr style=\"text-align: right;\">\n",
+       "      <th></th>\n",
+       "      <th>CHROM</th>\n",
+       "      <th>POS</th>\n",
+       "      <th>REF</th>\n",
+       "      <th>ALT</th>\n",
+       "      <th>LABEL</th>\n",
+       "      <th>SOURCE</th>\n",
+       "      <th>CONSEQUENCE</th>\n",
+       "      <th>ID</th>\n",
+       "      <th>REVIEW_STATUS</th>\n",
+       "      <th>GENE</th>\n",
+       "      <th>split</th>\n",
+       "      <th>INT_LABEL</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>chr1</td>\n",
+       "      <td>976215</td>\n",
+       "      <td>A</td>\n",
+       "      <td>G</td>\n",
+       "      <td>Pathogenic</td>\n",
+       "      <td>ClinVar</td>\n",
+       "      <td>missense_variant</td>\n",
+       "      <td>1320032</td>\n",
+       "      <td>no_assertion_criteria_provided</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>train</td>\n",
+       "      <td>1</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>chr1</td>\n",
+       "      <td>1050449</td>\n",
+       "      <td>G</td>\n",
+       "      <td>A</td>\n",
+       "      <td>Pathogenic</td>\n",
+       "      <td>ClinVar</td>\n",
+       "      <td>missense_variant</td>\n",
+       "      <td>1284257</td>\n",
+       "      <td>no_assertion_criteria_provided</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>train</td>\n",
+       "      <td>1</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>2</th>\n",
+       "      <td>chr1</td>\n",
+       "      <td>1050575</td>\n",
+       "      <td>G</td>\n",
+       "      <td>C</td>\n",
+       "      <td>Pathogenic</td>\n",
+       "      <td>ClinVar</td>\n",
+       "      <td>missense_variant</td>\n",
+       "      <td>18241</td>\n",
+       "      <td>no_assertion_criteria_provided</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>train</td>\n",
+       "      <td>1</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>3</th>\n",
+       "      <td>chr1</td>\n",
+       "      <td>1213738</td>\n",
+       "      <td>G</td>\n",
+       "      <td>A</td>\n",
+       "      <td>Pathogenic</td>\n",
+       "      <td>ClinVar</td>\n",
+       "      <td>missense_variant</td>\n",
+       "      <td>96692</td>\n",
+       "      <td>no_assertion_criteria_provided</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>train</td>\n",
+       "      <td>1</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>4</th>\n",
+       "      <td>chr1</td>\n",
+       "      <td>1232279</td>\n",
+       "      <td>A</td>\n",
+       "      <td>G</td>\n",
+       "      <td>Pathogenic</td>\n",
+       "      <td>ClinVar</td>\n",
+       "      <td>initiatior_codon_variant,missense_variant</td>\n",
+       "      <td>60484</td>\n",
+       "      <td>criteria_provided,_multiple_submitters,_no_con...</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>train</td>\n",
+       "      <td>1</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>...</th>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>22249</th>\n",
+       "      <td>chrY</td>\n",
+       "      <td>2787412</td>\n",
+       "      <td>C</td>\n",
+       "      <td>T</td>\n",
+       "      <td>Pathogenic</td>\n",
+       "      <td>ClinVar</td>\n",
+       "      <td>missense_variant</td>\n",
+       "      <td>9747</td>\n",
+       "      <td>no_assertion_criteria_provided</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>train</td>\n",
+       "      <td>1</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>22250</th>\n",
+       "      <td>chrY</td>\n",
+       "      <td>2787426</td>\n",
+       "      <td>C</td>\n",
+       "      <td>G</td>\n",
+       "      <td>Pathogenic</td>\n",
+       "      <td>ClinVar</td>\n",
+       "      <td>missense_variant</td>\n",
+       "      <td>9739</td>\n",
+       "      <td>criteria_provided,_single_submitter</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>train</td>\n",
+       "      <td>1</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>22251</th>\n",
+       "      <td>chrY</td>\n",
+       "      <td>2787515</td>\n",
+       "      <td>C</td>\n",
+       "      <td>A</td>\n",
+       "      <td>Pathogenic</td>\n",
+       "      <td>ClinVar</td>\n",
+       "      <td>missense_variant</td>\n",
+       "      <td>492908</td>\n",
+       "      <td>no_assertion_criteria_provided</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>train</td>\n",
+       "      <td>1</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>22252</th>\n",
+       "      <td>chrY</td>\n",
+       "      <td>2787551</td>\n",
+       "      <td>C</td>\n",
+       "      <td>T</td>\n",
+       "      <td>Pathogenic</td>\n",
+       "      <td>ClinVar</td>\n",
+       "      <td>missense_variant</td>\n",
+       "      <td>9754</td>\n",
+       "      <td>no_assertion_criteria_provided</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>train</td>\n",
+       "      <td>1</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>22253</th>\n",
+       "      <td>chrY</td>\n",
+       "      <td>7063898</td>\n",
+       "      <td>A</td>\n",
+       "      <td>T</td>\n",
+       "      <td>Pathogenic</td>\n",
+       "      <td>ClinVar</td>\n",
+       "      <td>missense_variant</td>\n",
+       "      <td>625467</td>\n",
+       "      <td>no_assertion_criteria_provided</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>train</td>\n",
+       "      <td>1</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "<p>22254 rows × 12 columns</p>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "      CHROM      POS REF ALT       LABEL   SOURCE  \\\n",
+       "0      chr1   976215   A   G  Pathogenic  ClinVar   \n",
+       "1      chr1  1050449   G   A  Pathogenic  ClinVar   \n",
+       "2      chr1  1050575   G   C  Pathogenic  ClinVar   \n",
+       "3      chr1  1213738   G   A  Pathogenic  ClinVar   \n",
+       "4      chr1  1232279   A   G  Pathogenic  ClinVar   \n",
+       "...     ...      ...  ..  ..         ...      ...   \n",
+       "22249  chrY  2787412   C   T  Pathogenic  ClinVar   \n",
+       "22250  chrY  2787426   C   G  Pathogenic  ClinVar   \n",
+       "22251  chrY  2787515   C   A  Pathogenic  ClinVar   \n",
+       "22252  chrY  2787551   C   T  Pathogenic  ClinVar   \n",
+       "22253  chrY  7063898   A   T  Pathogenic  ClinVar   \n",
+       "\n",
+       "                                     CONSEQUENCE       ID  \\\n",
+       "0                               missense_variant  1320032   \n",
+       "1                               missense_variant  1284257   \n",
+       "2                               missense_variant    18241   \n",
+       "3                               missense_variant    96692   \n",
+       "4      initiatior_codon_variant,missense_variant    60484   \n",
+       "...                                          ...      ...   \n",
+       "22249                           missense_variant     9747   \n",
+       "22250                           missense_variant     9739   \n",
+       "22251                           missense_variant   492908   \n",
+       "22252                           missense_variant     9754   \n",
+       "22253                           missense_variant   625467   \n",
+       "\n",
+       "                                           REVIEW_STATUS  GENE  split  \\\n",
+       "0                         no_assertion_criteria_provided   NaN  train   \n",
+       "1                         no_assertion_criteria_provided   NaN  train   \n",
+       "2                         no_assertion_criteria_provided   NaN  train   \n",
+       "3                         no_assertion_criteria_provided   NaN  train   \n",
+       "4      criteria_provided,_multiple_submitters,_no_con...   NaN  train   \n",
+       "...                                                  ...   ...    ...   \n",
+       "22249                     no_assertion_criteria_provided   NaN  train   \n",
+       "22250                criteria_provided,_single_submitter   NaN  train   \n",
+       "22251                     no_assertion_criteria_provided   NaN  train   \n",
+       "22252                     no_assertion_criteria_provided   NaN  train   \n",
+       "22253                     no_assertion_criteria_provided   NaN  train   \n",
+       "\n",
+       "       INT_LABEL  \n",
+       "0              1  \n",
+       "1              1  \n",
+       "2              1  \n",
+       "3              1  \n",
+       "4              1  \n",
+       "...          ...  \n",
+       "22249          1  \n",
+       "22250          1  \n",
+       "22251          1  \n",
+       "22252          1  \n",
+       "22253          1  \n",
+       "\n",
+       "[22254 rows x 12 columns]"
+      ]
+     },
+     "execution_count": 17,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "clinvar_raw\n",
+    "\n",
+    "# Preview the cleaned dataset\n",
+    "if clinvar_raw is not None:\n",
+    "    print(f\"📊 Dataset shape: {clinvar_raw.shape}\")\n",
+    "    print(f\"📋 Column names: {list(clinvar_raw.columns)}\")\n",
+    "    print(\"\\n🔍 First few rows:\")\n",
+    "    display(clinvar_raw.head())\n",
+    "    \n",
+    "    # Check for any null values\n",
+    "    null_counts = clinvar_raw.isnull().sum()\n",
+    "    if null_counts.sum() > 0:\n",
+    "        print(\"\\n⚠️ Null values found:\")\n",
+    "        print(null_counts[null_counts > 0])\n",
+    "else:\n",
+    "    print(\"❌ No data to display\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e380634b-0c22-4d1e-8520-6fc5728e7de5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Add new columns for gene information\n",
+    "if clinvar_raw is not None:\n",
+    "    clinvar_raw['GENE_ID'] = \"\"\n",
+    "    clinvar_raw['GENE'] = \"\"\n",
+    "    print(\"✅ Added GENE_ID and GENE columns\")\n",
+    "    print(f\"📊 Updated dataset shape: {clinvar_raw.shape}\")\n",
+    "else:\n",
+    "    print(\"⚠️ Cannot add columns - data not loaded\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "92b159f5-694d-4ee4-9616-1ebf00f71904",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "<div>\n",
+       "<style scoped>\n",
+       "    .dataframe tbody tr th:only-of-type {\n",
+       "        vertical-align: middle;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe tbody tr th {\n",
+       "        vertical-align: top;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe thead th {\n",
+       "        text-align: right;\n",
+       "    }\n",
+       "</style>\n",
+       "<table border=\"1\" class=\"dataframe\">\n",
+       "  <thead>\n",
+       "    <tr style=\"text-align: right;\">\n",
+       "      <th></th>\n",
+       "      <th>CHROM</th>\n",
+       "      <th>POS</th>\n",
+       "      <th>REF</th>\n",
+       "      <th>ALT</th>\n",
+       "      <th>LABEL</th>\n",
+       "      <th>SOURCE</th>\n",
+       "      <th>CONSEQUENCE</th>\n",
+       "      <th>ID</th>\n",
+       "      <th>REVIEW_STATUS</th>\n",
+       "      <th>GENE</th>\n",
+       "      <th>split</th>\n",
+       "      <th>INT_LABEL</th>\n",
+       "      <th>GENE_ID</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>chr1</td>\n",
+       "      <td>976215</td>\n",
+       "      <td>A</td>\n",
+       "      <td>G</td>\n",
+       "      <td>Pathogenic</td>\n",
+       "      <td>ClinVar</td>\n",
+       "      <td>missense_variant</td>\n",
+       "      <td>1320032</td>\n",
+       "      <td>no_assertion_criteria_provided</td>\n",
+       "      <td></td>\n",
+       "      <td>train</td>\n",
+       "      <td>1</td>\n",
+       "      <td></td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>chr1</td>\n",
+       "      <td>1050449</td>\n",
+       "      <td>G</td>\n",
+       "      <td>A</td>\n",
+       "      <td>Pathogenic</td>\n",
+       "      <td>ClinVar</td>\n",
+       "      <td>missense_variant</td>\n",
+       "      <td>1284257</td>\n",
+       "      <td>no_assertion_criteria_provided</td>\n",
+       "      <td></td>\n",
+       "      <td>train</td>\n",
+       "      <td>1</td>\n",
+       "      <td></td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>2</th>\n",
+       "      <td>chr1</td>\n",
+       "      <td>1050575</td>\n",
+       "      <td>G</td>\n",
+       "      <td>C</td>\n",
+       "      <td>Pathogenic</td>\n",
+       "      <td>ClinVar</td>\n",
+       "      <td>missense_variant</td>\n",
+       "      <td>18241</td>\n",
+       "      <td>no_assertion_criteria_provided</td>\n",
+       "      <td></td>\n",
+       "      <td>train</td>\n",
+       "      <td>1</td>\n",
+       "      <td></td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>3</th>\n",
+       "      <td>chr1</td>\n",
+       "      <td>1213738</td>\n",
+       "      <td>G</td>\n",
+       "      <td>A</td>\n",
+       "      <td>Pathogenic</td>\n",
+       "      <td>ClinVar</td>\n",
+       "      <td>missense_variant</td>\n",
+       "      <td>96692</td>\n",
+       "      <td>no_assertion_criteria_provided</td>\n",
+       "      <td></td>\n",
+       "      <td>train</td>\n",
+       "      <td>1</td>\n",
+       "      <td></td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>4</th>\n",
+       "      <td>chr1</td>\n",
+       "      <td>1232279</td>\n",
+       "      <td>A</td>\n",
+       "      <td>G</td>\n",
+       "      <td>Pathogenic</td>\n",
+       "      <td>ClinVar</td>\n",
+       "      <td>initiatior_codon_variant,missense_variant</td>\n",
+       "      <td>60484</td>\n",
+       "      <td>criteria_provided,_multiple_submitters,_no_con...</td>\n",
+       "      <td></td>\n",
+       "      <td>train</td>\n",
+       "      <td>1</td>\n",
+       "      <td></td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>...</th>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>22249</th>\n",
+       "      <td>chrY</td>\n",
+       "      <td>2787412</td>\n",
+       "      <td>C</td>\n",
+       "      <td>T</td>\n",
+       "      <td>Pathogenic</td>\n",
+       "      <td>ClinVar</td>\n",
+       "      <td>missense_variant</td>\n",
+       "      <td>9747</td>\n",
+       "      <td>no_assertion_criteria_provided</td>\n",
+       "      <td></td>\n",
+       "      <td>train</td>\n",
+       "      <td>1</td>\n",
+       "      <td></td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>22250</th>\n",
+       "      <td>chrY</td>\n",
+       "      <td>2787426</td>\n",
+       "      <td>C</td>\n",
+       "      <td>G</td>\n",
+       "      <td>Pathogenic</td>\n",
+       "      <td>ClinVar</td>\n",
+       "      <td>missense_variant</td>\n",
+       "      <td>9739</td>\n",
+       "      <td>criteria_provided,_single_submitter</td>\n",
+       "      <td></td>\n",
+       "      <td>train</td>\n",
+       "      <td>1</td>\n",
+       "      <td></td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>22251</th>\n",
+       "      <td>chrY</td>\n",
+       "      <td>2787515</td>\n",
+       "      <td>C</td>\n",
+       "      <td>A</td>\n",
+       "      <td>Pathogenic</td>\n",
+       "      <td>ClinVar</td>\n",
+       "      <td>missense_variant</td>\n",
+       "      <td>492908</td>\n",
+       "      <td>no_assertion_criteria_provided</td>\n",
+       "      <td></td>\n",
+       "      <td>train</td>\n",
+       "      <td>1</td>\n",
+       "      <td></td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>22252</th>\n",
+       "      <td>chrY</td>\n",
+       "      <td>2787551</td>\n",
+       "      <td>C</td>\n",
+       "      <td>T</td>\n",
+       "      <td>Pathogenic</td>\n",
+       "      <td>ClinVar</td>\n",
+       "      <td>missense_variant</td>\n",
+       "      <td>9754</td>\n",
+       "      <td>no_assertion_criteria_provided</td>\n",
+       "      <td></td>\n",
+       "      <td>train</td>\n",
+       "      <td>1</td>\n",
+       "      <td></td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>22253</th>\n",
+       "      <td>chrY</td>\n",
+       "      <td>7063898</td>\n",
+       "      <td>A</td>\n",
+       "      <td>T</td>\n",
+       "      <td>Pathogenic</td>\n",
+       "      <td>ClinVar</td>\n",
+       "      <td>missense_variant</td>\n",
+       "      <td>625467</td>\n",
+       "      <td>no_assertion_criteria_provided</td>\n",
+       "      <td></td>\n",
+       "      <td>train</td>\n",
+       "      <td>1</td>\n",
+       "      <td></td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "<p>22254 rows × 13 columns</p>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "      CHROM      POS REF ALT       LABEL   SOURCE  \\\n",
+       "0      chr1   976215   A   G  Pathogenic  ClinVar   \n",
+       "1      chr1  1050449   G   A  Pathogenic  ClinVar   \n",
+       "2      chr1  1050575   G   C  Pathogenic  ClinVar   \n",
+       "3      chr1  1213738   G   A  Pathogenic  ClinVar   \n",
+       "4      chr1  1232279   A   G  Pathogenic  ClinVar   \n",
+       "...     ...      ...  ..  ..         ...      ...   \n",
+       "22249  chrY  2787412   C   T  Pathogenic  ClinVar   \n",
+       "22250  chrY  2787426   C   G  Pathogenic  ClinVar   \n",
+       "22251  chrY  2787515   C   A  Pathogenic  ClinVar   \n",
+       "22252  chrY  2787551   C   T  Pathogenic  ClinVar   \n",
+       "22253  chrY  7063898   A   T  Pathogenic  ClinVar   \n",
+       "\n",
+       "                                     CONSEQUENCE       ID  \\\n",
+       "0                               missense_variant  1320032   \n",
+       "1                               missense_variant  1284257   \n",
+       "2                               missense_variant    18241   \n",
+       "3                               missense_variant    96692   \n",
+       "4      initiatior_codon_variant,missense_variant    60484   \n",
+       "...                                          ...      ...   \n",
+       "22249                           missense_variant     9747   \n",
+       "22250                           missense_variant     9739   \n",
+       "22251                           missense_variant   492908   \n",
+       "22252                           missense_variant     9754   \n",
+       "22253                           missense_variant   625467   \n",
+       "\n",
+       "                                           REVIEW_STATUS GENE  split  \\\n",
+       "0                         no_assertion_criteria_provided       train   \n",
+       "1                         no_assertion_criteria_provided       train   \n",
+       "2                         no_assertion_criteria_provided       train   \n",
+       "3                         no_assertion_criteria_provided       train   \n",
+       "4      criteria_provided,_multiple_submitters,_no_con...       train   \n",
+       "...                                                  ...  ...    ...   \n",
+       "22249                     no_assertion_criteria_provided       train   \n",
+       "22250                criteria_provided,_single_submitter       train   \n",
+       "22251                     no_assertion_criteria_provided       train   \n",
+       "22252                     no_assertion_criteria_provided       train   \n",
+       "22253                     no_assertion_criteria_provided       train   \n",
+       "\n",
+       "       INT_LABEL GENE_ID  \n",
+       "0              1          \n",
+       "1              1          \n",
+       "2              1          \n",
+       "3              1          \n",
+       "4              1          \n",
+       "...          ...     ...  \n",
+       "22249          1          \n",
+       "22250          1          \n",
+       "22251          1          \n",
+       "22252          1          \n",
+       "22253          1          \n",
+       "\n",
+       "[22254 rows x 13 columns]"
+      ]
+     },
+     "execution_count": 34,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "clinvar_raw\n",
+    "\n",
+    "# Display updated dataset with new columns\n",
+    "if clinvar_raw is not None:\n",
+    "    print(f\"📊 Dataset with new columns: {clinvar_raw.shape}\")\n",
+    "    print(f\"📋 All columns: {list(clinvar_raw.columns)}\")\n",
+    "    display(clinvar_raw.head())\n",
+    "else:\n",
+    "    print(\"❌ No data to display\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f36db716-392a-46a8-a404-d78165a4623c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import pandas as pd\n",
+    "import xml.etree.ElementTree as ET\n",
+    "import os\n",
+    "\n",
+    "# Parse ClinVar XML files to extract gene information\n",
+    "# This processes each ClinVar ID and extracts gene symbols and IDs from XML records\n",
+    "\n",
+    "if clinvar_raw is not None:\n",
+    "    # Load list of ClinVar IDs\n",
+    "    try:\n",
+    "        with open(\"Clinvar_ID.txt\", \"r\") as f:\n",
+    "            clinvar_ids = [line.strip() for line in f if line.strip()]\n",
+    "        \n",
+    "        print(f\"📋 Processing {len(clinvar_ids)} ClinVar IDs\")\n",
+    "        \n",
+    "        processed_count = 0\n",
+    "        error_count = 0\n",
+    "        \n",
+    "        # Process each ClinVar ID\n",
+    "        for i, clinvar_id in enumerate(clinvar_ids):\n",
+    "            if i % 100 == 0:  # Progress indicator\n",
+    "                print(f\"📊 Processing ID {i+1}/{len(clinvar_ids)}...\")\n",
+    "            \n",
+    "            try:\n",
+    "                id_int = int(clinvar_id)\n",
+    "                xml_path = f'data/{clinvar_id}.xml'\n",
+    "                \n",
+    "                # Check if XML file exists\n",
+    "                if not os.path.exists(xml_path):\n",
+    "                    print(f\"⚠️ XML file not found: {xml_path}\")\n",
+    "                    continue\n",
+    "                \n",
+    "                # Parse XML file\n",
+    "                with open(xml_path, 'r', encoding='utf-8') as file:\n",
+    "                    tree = ET.parse(file)\n",
+    "                    root = tree.getroot()\n",
+    "                \n",
+    "                # Check for error in XML\n",
+    "                error_element = root.find(\".//error\")\n",
+    "                if error_element is not None:\n",
+    "                    # Remove entries with errors\n",
+    "                    clinvar_raw = clinvar_raw[clinvar_raw[\"ID\"] != id_int]\n",
+    "                    error_count += 1\n",
+    "                    continue\n",
+    "                \n",
+    "                # Extract gene information\n",
+    "                gene_names = []\n",
+    "                gene_ids = []\n",
+    "                \n",
+    "                for gene in root.findall(\".//genes/gene\"):\n",
+    "                    symbol = gene.findtext(\"symbol\")\n",
+    "                    gene_id_data = gene.findtext(\"GeneID\")\n",
+    "                    \n",
+    "                    if symbol:\n",
+    "                        gene_names.append(symbol)\n",
+    "                    if gene_id_data:\n",
+    "                        gene_ids.append(gene_id_data)\n",
+    "                \n",
+    "                # Join multiple entries with commas\n",
+    "                gene_name_str = \", \".join(gene_names) if gene_names else \"\"\n",
+    "                gene_id_str = \", \".join(gene_ids) if gene_ids else \"\"\n",
+    "                \n",
+    "                # Update DataFrame\n",
+    "                mask = clinvar_raw[\"ID\"] == id_int\n",
+    "                if mask.any():\n",
+    "                    clinvar_raw.loc[mask, \"GENE\"] = gene_name_str\n",
+    "                    clinvar_raw.loc[mask, \"GENE_ID\"] = gene_id_str\n",
+    "                    processed_count += 1\n",
+    "                \n",
+    "            except ET.ParseError as e:\n",
+    "                print(f\"⚠️ XML parsing error for {clinvar_id}: {e}\")\n",
+    "                error_count += 1\n",
+    "            except ValueError as e:\n",
+    "                print(f\"⚠️ Invalid ClinVar ID {clinvar_id}: {e}\")\n",
+    "                error_count += 1\n",
+    "            except Exception as e:\n",
+    "                print(f\"⚠️ Unexpected error processing {clinvar_id}: {e}\")\n",
+    "                error_count += 1\n",
+    "        \n",
+    "        print(f\"\\n✅ Processing complete:\")\n",
+    "        print(f\"   📊 Successfully processed: {processed_count}\")\n",
+    "        print(f\"   ❌ Errors encountered: {error_count}\")\n",
+    "        print(f\"   📋 Final dataset shape: {clinvar_raw.shape}\")\n",
+    "        \n",
+    "    except FileNotFoundError:\n",
+    "        print(\"❌ Error: Clinvar_ID.txt not found\")\n",
+    "        print(\"Please run the ID extraction step first\")\n",
+    "    except Exception as e:\n",
+    "        print(f\"❌ Error during XML processing: {e}\")\n",
+    "else:\n",
+    "    print(\"⚠️ Cannot process XML files - ClinVar data not loaded\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ae0c9d8b-1b12-40a4-82ec-c3452e9dda90",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "<div>\n",
+       "<style scoped>\n",
+       "    .dataframe tbody tr th:only-of-type {\n",
+       "        vertical-align: middle;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe tbody tr th {\n",
+       "        vertical-align: top;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe thead th {\n",
+       "        text-align: right;\n",
+       "    }\n",
+       "</style>\n",
+       "<table border=\"1\" class=\"dataframe\">\n",
+       "  <thead>\n",
+       "    <tr style=\"text-align: right;\">\n",
+       "      <th></th>\n",
+       "      <th>CHROM</th>\n",
+       "      <th>POS</th>\n",
+       "      <th>REF</th>\n",
+       "      <th>ALT</th>\n",
+       "      <th>LABEL</th>\n",
+       "      <th>SOURCE</th>\n",
+       "      <th>CONSEQUENCE</th>\n",
+       "      <th>ID</th>\n",
+       "      <th>REVIEW_STATUS</th>\n",
+       "      <th>GENE</th>\n",
+       "      <th>split</th>\n",
+       "      <th>INT_LABEL</th>\n",
+       "      <th>GENE_ID</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>chr1</td>\n",
+       "      <td>976215</td>\n",
+       "      <td>A</td>\n",
+       "      <td>G</td>\n",
+       "      <td>Pathogenic</td>\n",
+       "      <td>ClinVar</td>\n",
+       "      <td>missense_variant</td>\n",
+       "      <td>1320032</td>\n",
+       "      <td>no_assertion_criteria_provided</td>\n",
+       "      <td>PERM1</td>\n",
+       "      <td>train</td>\n",
+       "      <td>1</td>\n",
+       "      <td>84808</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>chr1</td>\n",
+       "      <td>1050449</td>\n",
+       "      <td>G</td>\n",
+       "      <td>A</td>\n",
+       "      <td>Pathogenic</td>\n",
+       "      <td>ClinVar</td>\n",
+       "      <td>missense_variant</td>\n",
+       "      <td>1284257</td>\n",
+       "      <td>no_assertion_criteria_provided</td>\n",
+       "      <td>AGRN</td>\n",
+       "      <td>train</td>\n",
+       "      <td>1</td>\n",
+       "      <td>375790</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>2</th>\n",
+       "      <td>chr1</td>\n",
+       "      <td>1050575</td>\n",
+       "      <td>G</td>\n",
+       "      <td>C</td>\n",
+       "      <td>Pathogenic</td>\n",
+       "      <td>ClinVar</td>\n",
+       "      <td>missense_variant</td>\n",
+       "      <td>18241</td>\n",
+       "      <td>no_assertion_criteria_provided</td>\n",
+       "      <td>AGRN</td>\n",
+       "      <td>train</td>\n",
+       "      <td>1</td>\n",
+       "      <td>375790</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>3</th>\n",
+       "      <td>chr1</td>\n",
+       "      <td>1213738</td>\n",
+       "      <td>G</td>\n",
+       "      <td>A</td>\n",
+       "      <td>Pathogenic</td>\n",
+       "      <td>ClinVar</td>\n",
+       "      <td>missense_variant</td>\n",
+       "      <td>96692</td>\n",
+       "      <td>no_assertion_criteria_provided</td>\n",
+       "      <td>TNFRSF4</td>\n",
+       "      <td>train</td>\n",
+       "      <td>1</td>\n",
+       "      <td>7293</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>4</th>\n",
+       "      <td>chr1</td>\n",
+       "      <td>1232279</td>\n",
+       "      <td>A</td>\n",
+       "      <td>G</td>\n",
+       "      <td>Pathogenic</td>\n",
+       "      <td>ClinVar</td>\n",
+       "      <td>initiatior_codon_variant,missense_variant</td>\n",
+       "      <td>60484</td>\n",
+       "      <td>criteria_provided,_multiple_submitters,_no_con...</td>\n",
+       "      <td>B3GALT6</td>\n",
+       "      <td>train</td>\n",
+       "      <td>1</td>\n",
+       "      <td>126792</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>...</th>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>22249</th>\n",
+       "      <td>chrY</td>\n",
+       "      <td>2787412</td>\n",
+       "      <td>C</td>\n",
+       "      <td>T</td>\n",
+       "      <td>Pathogenic</td>\n",
+       "      <td>ClinVar</td>\n",
+       "      <td>missense_variant</td>\n",
+       "      <td>9747</td>\n",
+       "      <td>no_assertion_criteria_provided</td>\n",
+       "      <td>SRY</td>\n",
+       "      <td>train</td>\n",
+       "      <td>1</td>\n",
+       "      <td>6736</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>22250</th>\n",
+       "      <td>chrY</td>\n",
+       "      <td>2787426</td>\n",
+       "      <td>C</td>\n",
+       "      <td>G</td>\n",
+       "      <td>Pathogenic</td>\n",
+       "      <td>ClinVar</td>\n",
+       "      <td>missense_variant</td>\n",
+       "      <td>9739</td>\n",
+       "      <td>criteria_provided,_single_submitter</td>\n",
+       "      <td>SRY</td>\n",
+       "      <td>train</td>\n",
+       "      <td>1</td>\n",
+       "      <td>6736</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>22251</th>\n",
+       "      <td>chrY</td>\n",
+       "      <td>2787515</td>\n",
+       "      <td>C</td>\n",
+       "      <td>A</td>\n",
+       "      <td>Pathogenic</td>\n",
+       "      <td>ClinVar</td>\n",
+       "      <td>missense_variant</td>\n",
+       "      <td>492908</td>\n",
+       "      <td>no_assertion_criteria_provided</td>\n",
+       "      <td>SRY</td>\n",
+       "      <td>train</td>\n",
+       "      <td>1</td>\n",
+       "      <td>6736</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>22252</th>\n",
+       "      <td>chrY</td>\n",
+       "      <td>2787551</td>\n",
+       "      <td>C</td>\n",
+       "      <td>T</td>\n",
+       "      <td>Pathogenic</td>\n",
+       "      <td>ClinVar</td>\n",
+       "      <td>missense_variant</td>\n",
+       "      <td>9754</td>\n",
+       "      <td>no_assertion_criteria_provided</td>\n",
+       "      <td>SRY</td>\n",
+       "      <td>train</td>\n",
+       "      <td>1</td>\n",
+       "      <td>6736</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>22253</th>\n",
+       "      <td>chrY</td>\n",
+       "      <td>7063898</td>\n",
+       "      <td>A</td>\n",
+       "      <td>T</td>\n",
+       "      <td>Pathogenic</td>\n",
+       "      <td>ClinVar</td>\n",
+       "      <td>missense_variant</td>\n",
+       "      <td>625467</td>\n",
+       "      <td>no_assertion_criteria_provided</td>\n",
+       "      <td>LOC126057105, TBL1Y</td>\n",
+       "      <td>train</td>\n",
+       "      <td>1</td>\n",
+       "      <td>126057105, 90665</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "<p>22150 rows × 13 columns</p>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "      CHROM      POS REF ALT       LABEL   SOURCE  \\\n",
+       "0      chr1   976215   A   G  Pathogenic  ClinVar   \n",
+       "1      chr1  1050449   G   A  Pathogenic  ClinVar   \n",
+       "2      chr1  1050575   G   C  Pathogenic  ClinVar   \n",
+       "3      chr1  1213738   G   A  Pathogenic  ClinVar   \n",
+       "4      chr1  1232279   A   G  Pathogenic  ClinVar   \n",
+       "...     ...      ...  ..  ..         ...      ...   \n",
+       "22249  chrY  2787412   C   T  Pathogenic  ClinVar   \n",
+       "22250  chrY  2787426   C   G  Pathogenic  ClinVar   \n",
+       "22251  chrY  2787515   C   A  Pathogenic  ClinVar   \n",
+       "22252  chrY  2787551   C   T  Pathogenic  ClinVar   \n",
+       "22253  chrY  7063898   A   T  Pathogenic  ClinVar   \n",
+       "\n",
+       "                                     CONSEQUENCE       ID  \\\n",
+       "0                               missense_variant  1320032   \n",
+       "1                               missense_variant  1284257   \n",
+       "2                               missense_variant    18241   \n",
+       "3                               missense_variant    96692   \n",
+       "4      initiatior_codon_variant,missense_variant    60484   \n",
+       "...                                          ...      ...   \n",
+       "22249                           missense_variant     9747   \n",
+       "22250                           missense_variant     9739   \n",
+       "22251                           missense_variant   492908   \n",
+       "22252                           missense_variant     9754   \n",
+       "22253                           missense_variant   625467   \n",
+       "\n",
+       "                                           REVIEW_STATUS                 GENE  \\\n",
+       "0                         no_assertion_criteria_provided                PERM1   \n",
+       "1                         no_assertion_criteria_provided                 AGRN   \n",
+       "2                         no_assertion_criteria_provided                 AGRN   \n",
+       "3                         no_assertion_criteria_provided              TNFRSF4   \n",
+       "4      criteria_provided,_multiple_submitters,_no_con...              B3GALT6   \n",
+       "...                                                  ...                  ...   \n",
+       "22249                     no_assertion_criteria_provided                  SRY   \n",
+       "22250                criteria_provided,_single_submitter                  SRY   \n",
+       "22251                     no_assertion_criteria_provided                  SRY   \n",
+       "22252                     no_assertion_criteria_provided                  SRY   \n",
+       "22253                     no_assertion_criteria_provided  LOC126057105, TBL1Y   \n",
+       "\n",
+       "       split  INT_LABEL           GENE_ID  \n",
+       "0      train          1             84808  \n",
+       "1      train          1            375790  \n",
+       "2      train          1            375790  \n",
+       "3      train          1              7293  \n",
+       "4      train          1            126792  \n",
+       "...      ...        ...               ...  \n",
+       "22249  train          1              6736  \n",
+       "22250  train          1              6736  \n",
+       "22251  train          1              6736  \n",
+       "22252  train          1              6736  \n",
+       "22253  train          1  126057105, 90665  \n",
+       "\n",
+       "[22150 rows x 13 columns]"
+      ]
+     },
+     "execution_count": 39,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "clinvar_raw\n",
+    "\n",
+    "# Display the dataset with extracted gene information\n",
+    "if clinvar_raw is not None:\n",
+    "    print(f\"📊 Dataset after gene extraction: {clinvar_raw.shape}\")\n",
+    "    \n",
+    "    # Show statistics\n",
+    "    gene_filled = (clinvar_raw['GENE'] != '').sum()\n",
+    "    gene_id_filled = (clinvar_raw['GENE_ID'] != '').sum()\n",
+    "    \n",
+    "    print(f\"📋 Entries with gene names: {gene_filled} ({gene_filled/len(clinvar_raw)*100:.1f}%)\")\n",
+    "    print(f\"📋 Entries with gene IDs: {gene_id_filled} ({gene_id_filled/len(clinvar_raw)*100:.1f}%)\")\n",
+    "    \n",
+    "    # Show sample data\n",
+    "    display(clinvar_raw.head(10))\n",
+    "else:\n",
+    "    print(\"❌ No data to display\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b76910bd-aa86-4943-a0f2-dcf9756ad81d",
+   "metadata": {},
+   "source": [
+    "## Disease/Phenotype Information Extraction\n",
+    "\n",
+    "This section extracts disease and phenotype information from the ClinVar XML records. Each variant may be associated with multiple diseases, so the data is expanded to create one row per variant-disease combination.\n",
+    "\n",
+    "### Putting in the Disease Name"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "54ccd972-5804-4d63-9012-5531034d2b60",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Extract disease/phenotype information from ClinVar XML files\n",
+    "# This creates multiple rows for variants associated with multiple diseases\n",
+    "\n",
+    "if clinvar_raw is not None:\n",
+    "    try:\n",
+    "        # Load ClinVar IDs\n",
+    "        with open(\"Clinvar_ID.txt\", \"r\") as f:\n",
+    "            clinvar_ids = [line.strip() for line in f if line.strip()]\n",
+    "        \n",
+    "        print(f\"📋 Processing {len(clinvar_ids)} ClinVar IDs for disease extraction\")\n",
+    "        \n",
+    "        # Ensure ID column is integer type\n",
+    "        clinvar_raw[\"ID\"] = clinvar_raw[\"ID\"].astype(int)\n",
+    "        \n",
+    "        # Create new DataFrame to store expanded data\n",
+    "        clinvar_data = pd.DataFrame(columns=clinvar_raw.columns.tolist() + [\"Disease\"])\n",
+    "        \n",
+    "        processed_count = 0\n",
+    "        disease_count = 0\n",
+    "        \n",
+    "        # Process each ClinVar ID\n",
+    "        for i, clinvar_id in enumerate(clinvar_ids):\n",
+    "            if i % 100 == 0:  # Progress indicator\n",
+    "                print(f\"📊 Processing disease info {i+1}/{len(clinvar_ids)}...\")\n",
+    "            \n",
+    "            try:\n",
+    "                id_int = int(clinvar_id)\n",
+    "                xml_path = f\"data/{clinvar_id}.xml\"\n",
+    "                \n",
+    "                if not os.path.exists(xml_path):\n",
+    "                    continue\n",
+    "                \n",
+    "                # Parse XML\n",
+    "                tree = ET.parse(xml_path)\n",
+    "                root = tree.getroot()\n",
+    "                \n",
+    "                # Extract all trait names (diseases/phenotypes)\n",
+    "                trait_names = []\n",
+    "                for trait in root.findall(\".//trait\"):\n",
+    "                    trait_name = trait.findtext(\"trait_name\")\n",
+    "                    if trait_name:\n",
+    "                        trait_names.append(trait_name)\n",
+    "                \n",
+    "                # Filter out 'not provided' if other traits exist\n",
+    "                filtered_traits = [t for t in trait_names if t.lower() != \"not provided\"]\n",
+    "                if not filtered_traits and \"not provided\" in [t.lower() for t in trait_names]:\n",
+    "                    filtered_traits = [\"not provided\"]\n",
+    "                \n",
+    "                # If no traits found, use empty string\n",
+    "                if not filtered_traits:\n",
+    "                    filtered_traits = [\"\"]\n",
+    "                \n",
+    "                # Create one row for each disease/trait\n",
+    "                base_row = clinvar_raw[clinvar_raw[\"ID\"] == id_int]\n",
+    "                if not base_row.empty:\n",
+    "                    for disease_name in filtered_traits:\n",
+    "                        new_row = base_row.copy()\n",
+    "                        new_row[\"Disease\"] = disease_name\n",
+    "                        clinvar_data = pd.concat([clinvar_data, new_row], ignore_index=True)\n",
+    "                        disease_count += 1\n",
+    "                    processed_count += 1\n",
+    "                \n",
+    "            except ET.ParseError as e:\n",
+    "                print(f\"⚠️ XML parsing error for {clinvar_id}: {e}\")\n",
+    "            except Exception as e:\n",
+    "                print(f\"⚠️ Error processing {clinvar_id}: {e}\")\n",
+    "        \n",
+    "        print(f\"\\n✅ Disease extraction complete:\")\n",
+    "        print(f\"   📊 Variants processed: {processed_count}\")\n",
+    "        print(f\"   🔬 Total variant-disease pairs: {disease_count}\")\n",
+    "        print(f\"   📋 Final dataset shape: {clinvar_data.shape}\")\n",
+    "        \n",
+    "        # Save intermediate results\n",
+    "        clinvar_data.to_csv(\"clinvar_with_disease.csv\", sep='\\t', index=False)\n",
+    "        print(\"💾 Saved results to clinvar_with_disease.csv\")\n",
+    "        \n",
+    "    except FileNotFoundError:\n",
+    "        print(\"❌ Error: Required files not found\")\n",
+    "        print(\"Please ensure Clinvar_ID.txt exists and XML files are downloaded\")\n",
+    "        clinvar_data = None\n",
+    "    except Exception as e:\n",
+    "        print(f\"❌ Error during disease extraction: {e}\")\n",
+    "        clinvar_data = None\n",
+    "else:\n",
+    "    print(\"⚠️ Cannot extract diseases - ClinVar data not loaded\")\n",
+    "    clinvar_data = None"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "277445cd-72b9-44a4-a257-49cd3202e501",
+   "metadata": {
+    "scrolled": true
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "<div>\n",
+       "<style scoped>\n",
+       "    .dataframe tbody tr th:only-of-type {\n",
+       "        vertical-align: middle;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe tbody tr th {\n",
+       "        vertical-align: top;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe thead th {\n",
+       "        text-align: right;\n",
+       "    }\n",
+       "</style>\n",
+       "<table border=\"1\" class=\"dataframe\">\n",
+       "  <thead>\n",
+       "    <tr style=\"text-align: right;\">\n",
+       "      <th></th>\n",
+       "      <th>CHROM</th>\n",
+       "      <th>POS</th>\n",
+       "      <th>REF</th>\n",
+       "      <th>ALT</th>\n",
+       "      <th>LABEL</th>\n",
+       "      <th>SOURCE</th>\n",
+       "      <th>CONSEQUENCE</th>\n",
+       "      <th>ID</th>\n",
+       "      <th>REVIEW_STATUS</th>\n",
+       "      <th>GENE</th>\n",
+       "      <th>split</th>\n",
+       "      <th>INT_LABEL</th>\n",
+       "      <th>GENE_ID</th>\n",
+       "      <th>Disease</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>chr1</td>\n",
+       "      <td>976215</td>\n",
+       "      <td>A</td>\n",
+       "      <td>G</td>\n",
+       "      <td>Pathogenic</td>\n",
+       "      <td>ClinVar</td>\n",
+       "      <td>missense_variant</td>\n",
+       "      <td>1320032</td>\n",
+       "      <td>no_assertion_criteria_provided</td>\n",
+       "      <td>PERM1</td>\n",
+       "      <td>train</td>\n",
+       "      <td>1</td>\n",
+       "      <td>84808</td>\n",
+       "      <td>Renal tubular epithelial cell apoptosis</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>chr1</td>\n",
+       "      <td>976215</td>\n",
+       "      <td>A</td>\n",
+       "      <td>G</td>\n",
+       "      <td>Pathogenic</td>\n",
+       "      <td>ClinVar</td>\n",
+       "      <td>missense_variant</td>\n",
+       "      <td>1320032</td>\n",
+       "      <td>no_assertion_criteria_provided</td>\n",
+       "      <td>PERM1</td>\n",
+       "      <td>train</td>\n",
+       "      <td>1</td>\n",
+       "      <td>84808</td>\n",
+       "      <td>Neutrophil inclusion bodies</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>2</th>\n",
+       "      <td>chr1</td>\n",
+       "      <td>1050449</td>\n",
+       "      <td>G</td>\n",
+       "      <td>A</td>\n",
+       "      <td>Pathogenic</td>\n",
+       "      <td>ClinVar</td>\n",
+       "      <td>missense_variant</td>\n",
+       "      <td>1284257</td>\n",
+       "      <td>no_assertion_criteria_provided</td>\n",
+       "      <td>AGRN</td>\n",
+       "      <td>train</td>\n",
+       "      <td>1</td>\n",
+       "      <td>375790</td>\n",
+       "      <td>Congenital myasthenic syndrome 8</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>3</th>\n",
+       "      <td>chr1</td>\n",
+       "      <td>1050575</td>\n",
+       "      <td>G</td>\n",
+       "      <td>C</td>\n",
+       "      <td>Pathogenic</td>\n",
+       "      <td>ClinVar</td>\n",
+       "      <td>missense_variant</td>\n",
+       "      <td>18241</td>\n",
+       "      <td>no_assertion_criteria_provided</td>\n",
+       "      <td>AGRN</td>\n",
+       "      <td>train</td>\n",
+       "      <td>1</td>\n",
+       "      <td>375790</td>\n",
+       "      <td>Congenital myasthenic syndrome 8</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>4</th>\n",
+       "      <td>chr1</td>\n",
+       "      <td>1213738</td>\n",
+       "      <td>G</td>\n",
+       "      <td>A</td>\n",
+       "      <td>Pathogenic</td>\n",
+       "      <td>ClinVar</td>\n",
+       "      <td>missense_variant</td>\n",
+       "      <td>96692</td>\n",
+       "      <td>no_assertion_criteria_provided</td>\n",
+       "      <td>TNFRSF4</td>\n",
+       "      <td>train</td>\n",
+       "      <td>1</td>\n",
+       "      <td>7293</td>\n",
+       "      <td>Combined immunodeficiency due to OX40 deficiency</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>...</th>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>32680</th>\n",
+       "      <td>chrY</td>\n",
+       "      <td>2787412</td>\n",
+       "      <td>C</td>\n",
+       "      <td>T</td>\n",
+       "      <td>Pathogenic</td>\n",
+       "      <td>ClinVar</td>\n",
+       "      <td>missense_variant</td>\n",
+       "      <td>9747</td>\n",
+       "      <td>no_assertion_criteria_provided</td>\n",
+       "      <td>SRY</td>\n",
+       "      <td>train</td>\n",
+       "      <td>1</td>\n",
+       "      <td>6736</td>\n",
+       "      <td>46,XY sex reversal 1</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>32681</th>\n",
+       "      <td>chrY</td>\n",
+       "      <td>2787426</td>\n",
+       "      <td>C</td>\n",
+       "      <td>G</td>\n",
+       "      <td>Pathogenic</td>\n",
+       "      <td>ClinVar</td>\n",
+       "      <td>missense_variant</td>\n",
+       "      <td>9739</td>\n",
+       "      <td>criteria_provided,_single_submitter</td>\n",
+       "      <td>SRY</td>\n",
+       "      <td>train</td>\n",
+       "      <td>1</td>\n",
+       "      <td>6736</td>\n",
+       "      <td>not provided</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>32682</th>\n",
+       "      <td>chrY</td>\n",
+       "      <td>2787515</td>\n",
+       "      <td>C</td>\n",
+       "      <td>A</td>\n",
+       "      <td>Pathogenic</td>\n",
+       "      <td>ClinVar</td>\n",
+       "      <td>missense_variant</td>\n",
+       "      <td>492908</td>\n",
+       "      <td>no_assertion_criteria_provided</td>\n",
+       "      <td>SRY</td>\n",
+       "      <td>train</td>\n",
+       "      <td>1</td>\n",
+       "      <td>6736</td>\n",
+       "      <td>46,XY sex reversal 1</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>32683</th>\n",
+       "      <td>chrY</td>\n",
+       "      <td>2787551</td>\n",
+       "      <td>C</td>\n",
+       "      <td>T</td>\n",
+       "      <td>Pathogenic</td>\n",
+       "      <td>ClinVar</td>\n",
+       "      <td>missense_variant</td>\n",
+       "      <td>9754</td>\n",
+       "      <td>no_assertion_criteria_provided</td>\n",
+       "      <td>SRY</td>\n",
+       "      <td>train</td>\n",
+       "      <td>1</td>\n",
+       "      <td>6736</td>\n",
+       "      <td>46,XY sex reversal 1</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>32684</th>\n",
+       "      <td>chrY</td>\n",
+       "      <td>7063898</td>\n",
+       "      <td>A</td>\n",
+       "      <td>T</td>\n",
+       "      <td>Pathogenic</td>\n",
+       "      <td>ClinVar</td>\n",
+       "      <td>missense_variant</td>\n",
+       "      <td>625467</td>\n",
+       "      <td>no_assertion_criteria_provided</td>\n",
+       "      <td>LOC126057105, TBL1Y</td>\n",
+       "      <td>train</td>\n",
+       "      <td>1</td>\n",
+       "      <td>126057105, 90665</td>\n",
+       "      <td>Deafness, Y-linked 2</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "<p>32685 rows × 14 columns</p>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "      CHROM      POS REF ALT       LABEL   SOURCE       CONSEQUENCE       ID  \\\n",
+       "0      chr1   976215   A   G  Pathogenic  ClinVar  missense_variant  1320032   \n",
+       "1      chr1   976215   A   G  Pathogenic  ClinVar  missense_variant  1320032   \n",
+       "2      chr1  1050449   G   A  Pathogenic  ClinVar  missense_variant  1284257   \n",
+       "3      chr1  1050575   G   C  Pathogenic  ClinVar  missense_variant    18241   \n",
+       "4      chr1  1213738   G   A  Pathogenic  ClinVar  missense_variant    96692   \n",
+       "...     ...      ...  ..  ..         ...      ...               ...      ...   \n",
+       "32680  chrY  2787412   C   T  Pathogenic  ClinVar  missense_variant     9747   \n",
+       "32681  chrY  2787426   C   G  Pathogenic  ClinVar  missense_variant     9739   \n",
+       "32682  chrY  2787515   C   A  Pathogenic  ClinVar  missense_variant   492908   \n",
+       "32683  chrY  2787551   C   T  Pathogenic  ClinVar  missense_variant     9754   \n",
+       "32684  chrY  7063898   A   T  Pathogenic  ClinVar  missense_variant   625467   \n",
+       "\n",
+       "                             REVIEW_STATUS                 GENE  split  \\\n",
+       "0           no_assertion_criteria_provided                PERM1  train   \n",
+       "1           no_assertion_criteria_provided                PERM1  train   \n",
+       "2           no_assertion_criteria_provided                 AGRN  train   \n",
+       "3           no_assertion_criteria_provided                 AGRN  train   \n",
+       "4           no_assertion_criteria_provided              TNFRSF4  train   \n",
+       "...                                    ...                  ...    ...   \n",
+       "32680       no_assertion_criteria_provided                  SRY  train   \n",
+       "32681  criteria_provided,_single_submitter                  SRY  train   \n",
+       "32682       no_assertion_criteria_provided                  SRY  train   \n",
+       "32683       no_assertion_criteria_provided                  SRY  train   \n",
+       "32684       no_assertion_criteria_provided  LOC126057105, TBL1Y  train   \n",
+       "\n",
+       "      INT_LABEL           GENE_ID  \\\n",
+       "0             1             84808   \n",
+       "1             1             84808   \n",
+       "2             1            375790   \n",
+       "3             1            375790   \n",
+       "4             1              7293   \n",
+       "...         ...               ...   \n",
+       "32680         1              6736   \n",
+       "32681         1              6736   \n",
+       "32682         1              6736   \n",
+       "32683         1              6736   \n",
+       "32684         1  126057105, 90665   \n",
+       "\n",
+       "                                                Disease  \n",
+       "0               Renal tubular epithelial cell apoptosis  \n",
+       "1                           Neutrophil inclusion bodies  \n",
+       "2                      Congenital myasthenic syndrome 8  \n",
+       "3                      Congenital myasthenic syndrome 8  \n",
+       "4      Combined immunodeficiency due to OX40 deficiency  \n",
+       "...                                                 ...  \n",
+       "32680                              46,XY sex reversal 1  \n",
+       "32681                                      not provided  \n",
+       "32682                              46,XY sex reversal 1  \n",
+       "32683                              46,XY sex reversal 1  \n",
+       "32684                              Deafness, Y-linked 2  \n",
+       "\n",
+       "[32685 rows x 14 columns]"
+      ]
+     },
+     "execution_count": 51,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "clinvar_data\n",
+    "\n",
+    "# Display the dataset with disease information\n",
+    "if 'clinvar_data' in locals() and clinvar_data is not None:\n",
+    "    print(f\"📊 Dataset with diseases: {clinvar_data.shape}\")\n",
+    "    \n",
+    "    # Show disease statistics\n",
+    "    disease_counts = clinvar_data['Disease'].value_counts()\n",
+    "    print(f\"\\n🔬 Disease distribution (top 10):\")\n",
+    "    print(disease_counts.head(10))\n",
+    "    \n",
+    "    # Show sample data\n",
+    "    print(\"\\n🔍 Sample data:\")\n",
+    "    display(clinvar_data.head())\n",
+    "else:\n",
+    "    print(\"❌ No disease data to display\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c6b1c6dc-33ed-4f57-a385-29816f4c9984",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "np.int64(2749)"
+      ]
+     },
+     "execution_count": 53,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# Count entries with 'not provided' disease information\n",
+    "if 'clinvar_data' in locals() and clinvar_data is not None:\n",
+    "    not_provided_count = (clinvar_data[\"Disease\"] == \"not provided\").sum()\n",
+    "    total_count = len(clinvar_data)\n",
+    "    \n",
+    "    print(f\"📊 Entries with 'not provided' disease: {not_provided_count}\")\n",
+    "    print(f\"📊 Total entries: {total_count}\")\n",
+    "    print(f\"📊 Percentage: {not_provided_count/total_count*100:.1f}%\")\n",
+    "else:\n",
+    "    print(\"❌ Cannot calculate statistics - data not available\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "8a7513ee-96b2-4c7d-8678-0195eb826aa5",
+   "metadata": {},
+   "source": [
+    "## Gene ID to Gene Name Mapping\n",
+    "\n",
+    "This section converts gene IDs to human-readable gene names using NCBI Entrez utilities.\n",
+    "\n",
+    "**Prerequisites**: NCBI Entrez Direct tools must be installed:\n",
+    "- macOS: `brew install brewsci/bio/edirect`\n",
+    "- Linux: Follow NCBI EDirect installation guide\n",
+    "\n",
+    "The process:\n",
+    "1. Extract unique gene IDs from the dataset\n",
+    "2. Use `esummary` to fetch gene descriptions from NCBI\n",
+    "3. Create a mapping dictionary\n",
+    "4. Apply the mapping to add gene names to the dataset"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ee0d3632-d11e-4429-bb50-5eb9ba55d424",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#!/usr/bin/env python3\n",
+    "\n",
+    "import os\n",
+    "import pandas as pd\n",
+    "\n",
+    "# Extract unique gene IDs and create mapping file\n",
+    "# This prepares the gene ID list for NCBI lookup\n",
+    "\n",
+    "if 'clinvar_data' in locals() and clinvar_data is not None:\n",
+    "    # Extract all unique gene IDs\n",
+    "    all_gene_ids = set()\n",
+    "    \n",
+    "    for gene_id_str in clinvar_data['GENE_ID'].dropna():\n",
+    "        if gene_id_str.strip():  # Skip empty strings\n",
+    "            # Split comma-separated IDs\n",
+    "            ids = [gid.strip() for gid in gene_id_str.split(',') if gid.strip()]\n",
+    "            all_gene_ids.update(ids)\n",
+    "    \n",
+    "    # Save unique gene IDs to file\n",
+    "    with open(\"gene_id.txt\", 'w') as f:\n",
+    "        for gene_id in sorted(all_gene_ids):\n",
+    "            f.write(f\"{gene_id}\\n\")\n",
+    "    \n",
+    "    print(f\"✅ Extracted {len(all_gene_ids)} unique gene IDs to gene_id.txt\")\n",
+    "    \n",
+    "    # Create the shell script for NCBI lookup\n",
+    "    script_content = '''#!/bin/bash\n",
+    "\n",
+    "input_file=\"gene_id.txt\"\n",
+    "output_file=\"gene_id_to_name.json\"\n",
+    "\n",
+    "# Check if input file exists\n",
+    "if [ ! -f \"$input_file\" ]; then\n",
+    "    echo \"❌ Error: $input_file not found\"\n",
+    "    exit 1\n",
+    "fi\n",
+    "\n",
+    "# Check if EDirect tools are available\n",
+    "if ! command -v esummary &> /dev/null; then\n",
+    "    echo \"❌ Error: NCBI EDirect tools not found\"\n",
+    "    echo \"Please install: brew install brewsci/bio/edirect (macOS)\"\n",
+    "    exit 1\n",
+    "fi\n",
+    "\n",
+    "echo \"🚀 Starting gene ID to name mapping...\"\n",
+    "\n",
+    "# Start JSON object\n",
+    "echo \"{\" > \"$output_file\"\n",
+    "\n",
+    "first_entry=true\n",
+    "total_lines=$(wc -l < \"$input_file\")\n",
+    "current_line=0\n",
+    "\n",
+    "while IFS= read -r gene_id; do\n",
+    "    # Skip empty lines\n",
+    "    [[ -z \"$gene_id\" ]] && continue\n",
+    "    \n",
+    "    current_line=$((current_line + 1))\n",
+    "    \n",
+    "    # Progress indicator\n",
+    "    if (( current_line % 50 == 0 )); then\n",
+    "        echo \"📊 Processing $current_line/$total_lines gene IDs...\"\n",
+    "    fi\n",
+    "    \n",
+    "    # Fetch gene description using Entrez Direct\n",
+    "    description=$(esummary -db gene -id \"$gene_id\" 2>/dev/null | xtract -pattern DocumentSummary -element Description)\n",
+    "    \n",
+    "    # Handle empty description\n",
+    "    if [ -z \"$description\" ]; then\n",
+    "        description=\"Unknown\"\n",
+    "    fi\n",
+    "    \n",
+    "    # JSON escape quotes and other special characters\n",
+    "    description=$(printf '%s' \"$description\" | sed 's/\"/\\\\\"/g')\n",
+    "    \n",
+    "    # Add comma if not the first entry\n",
+    "    if [ \"$first_entry\" = true ]; then\n",
+    "        first_entry=false\n",
+    "    else\n",
+    "        echo \",\" >> \"$output_file\"\n",
+    "    fi\n",
+    "    \n",
+    "    # Append key-value pair\n",
+    "    echo \"  \\\"$gene_id\\\": \\\"$description\\\"\" >> \"$output_file\"\n",
+    "    \n",
+    "done < \"$input_file\"\n",
+    "\n",
+    "# Close JSON object\n",
+    "echo \"\" >> \"$output_file\"\n",
+    "echo \"}\" >> \"$output_file\"\n",
+    "\n",
+    "echo \"✅ Gene ID to name mapping completed\"\n",
+    "echo \"💾 Results saved to $output_file\"\n",
+    "'''\n",
+    "    \n",
+    "    # Write the script\n",
+    "    with open(\"gene_mapping.sh\", 'w') as f:\n",
+    "        f.write(script_content)\n",
+    "    \n",
+    "    # Make executable\n",
+    "    os.chmod(\"gene_mapping.sh\", 0o755)\n",
+    "    \n",
+    "    print(\"✅ Created gene_mapping.sh script\")\n",
+    "    print(\"\\n🚀 To run the gene mapping:\")\n",
+    "    print(\"   ./gene_mapping.sh\")\n",
+    "    print(\"\\n⚠️ Note: This requires NCBI EDirect tools to be installed\")\n",
+    "    \n",
+    "else:\n",
+    "    print(\"⚠️ Cannot create gene mapping - data not available\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "1957ef57-1af8-46a1-8d1b-147f6b423619",
+   "metadata": {},
+   "source": [
+    "## Apply Gene Name Mapping\n",
+    "\n",
+    "Load the gene ID to name mapping and apply it to the dataset to add human-readable gene names.\n",
+    "\n",
+    "Read json and add it to the clinvar_data df"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b39be718-c0ae-4aae-b1d8-d0c872947ec2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import json\n",
+    "\n",
+    "# Load gene ID to name mapping and apply to dataset\n",
+    "\n",
+    "if 'clinvar_data' in locals() and clinvar_data is not None:\n",
+    "    try:\n",
+    "        # Load gene ID → name mapping\n",
+    "        with open(\"gene_id_to_name.json\", \"r\") as f:\n",
+    "            gene_id_dict = json.load(f)\n",
+    "        \n",
+    "        print(f\"✅ Loaded mapping for {len(gene_id_dict)} gene IDs\")\n",
+    "        \n",
+    "        # Function to convert gene IDs to gene names\n",
+    "        def get_gene_names(gene_id_str):\n",
+    "            if pd.isna(gene_id_str) or not gene_id_str.strip():\n",
+    "                return \"\"\n",
+    "            \n",
+    "            gene_ids = [gid.strip() for gid in gene_id_str.split(\",\") if gid.strip()]\n",
+    "            gene_names = []\n",
+    "            \n",
+    "            for gid in gene_ids:\n",
+    "                gene_name = gene_id_dict.get(gid, f\"Unknown_ID_{gid}\")\n",
+    "                gene_names.append(gene_name)\n",
+    "            \n",
+    "            return \" | \".join(gene_names)\n",
+    "        \n",
+    "        # Apply mapping to create gene names column\n",
+    "        print(\"📊 Applying gene name mapping...\")\n",
+    "        clinvar_data[\"GENE_Name\"] = clinvar_data[\"GENE_ID\"].apply(get_gene_names)\n",
+    "        \n",
+    "        # Statistics\n",
+    "        mapped_count = (clinvar_data[\"GENE_Name\"] != \"\").sum()\n",
+    "        print(f\"✅ Gene names mapped for {mapped_count} entries ({mapped_count/len(clinvar_data)*100:.1f}%)\")\n",
+    "        \n",
+    "        # Show sample mappings\n",
+    "        sample_data = clinvar_data[clinvar_data[\"GENE_Name\"] != \"\"][[\"GENE_ID\", \"GENE_Name\"]].head()\n",
+    "        if not sample_data.empty:\n",
+    "            print(\"\\n🔍 Sample gene ID to name mappings:\")\n",
+    "            for _, row in sample_data.iterrows():\n",
+    "                print(f\"   {row['GENE_ID']} → {row['GENE_Name'][:100]}{'...' if len(row['GENE_Name']) > 100 else ''}\")\n",
+    "        \n",
+    "    except FileNotFoundError:\n",
+    "        print(\"❌ Error: gene_id_to_name.json not found\")\n",
+    "        print(\"Please run the gene mapping script first: ./gene_mapping.sh\")\n",
+    "        # Create empty column as fallback\n",
+    "        clinvar_data[\"GENE_Name\"] = \"\"\n",
+    "    except json.JSONDecodeError as e:\n",
+    "        print(f\"❌ Error parsing JSON mapping file: {e}\")\n",
+    "        clinvar_data[\"GENE_Name\"] = \"\"\n",
+    "    except Exception as e:\n",
+    "        print(f\"❌ Error applying gene mapping: {e}\")\n",
+    "        clinvar_data[\"GENE_Name\"] = \"\"\n",
+    "else:\n",
+    "    print(\"⚠️ Cannot apply gene mapping - data not available\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4b7a44c2-7823-47c1-b268-22a1815ffd09",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "<div>\n",
+       "<style scoped>\n",
+       "    .dataframe tbody tr th:only-of-type {\n",
+       "        vertical-align: middle;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe tbody tr th {\n",
+       "        vertical-align: top;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe thead th {\n",
+       "        text-align: right;\n",
+       "    }\n",
+       "</style>\n",
+       "<table border=\"1\" class=\"dataframe\">\n",
+       "  <thead>\n",
+       "    <tr style=\"text-align: right;\">\n",
+       "      <th></th>\n",
+       "      <th>CHROM</th>\n",
+       "      <th>POS</th>\n",
+       "      <th>REF</th>\n",
+       "      <th>ALT</th>\n",
+       "      <th>LABEL</th>\n",
+       "      <th>SOURCE</th>\n",
+       "      <th>CONSEQUENCE</th>\n",
+       "      <th>ID</th>\n",
+       "      <th>REVIEW_STATUS</th>\n",
+       "      <th>GENE</th>\n",
+       "      <th>split</th>\n",
+       "      <th>INT_LABEL</th>\n",
+       "      <th>GENE_ID</th>\n",
+       "      <th>Disease</th>\n",
+       "      <th>GENE_Name</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>chr1</td>\n",
+       "      <td>976215</td>\n",
+       "      <td>A</td>\n",
+       "      <td>G</td>\n",
+       "      <td>Pathogenic</td>\n",
+       "      <td>ClinVar</td>\n",
+       "      <td>missense_variant</td>\n",
+       "      <td>1320032</td>\n",
+       "      <td>no_assertion_criteria_provided</td>\n",
+       "      <td>PERM1</td>\n",
+       "      <td>train</td>\n",
+       "      <td>1</td>\n",
+       "      <td>84808</td>\n",
+       "      <td>Renal tubular epithelial cell apoptosis</td>\n",
+       "      <td>PPARGC1 and ESRR induced regulator, muscle 1</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>chr1</td>\n",
+       "      <td>976215</td>\n",
+       "      <td>A</td>\n",
+       "      <td>G</td>\n",
+       "      <td>Pathogenic</td>\n",
+       "      <td>ClinVar</td>\n",
+       "      <td>missense_variant</td>\n",
+       "      <td>1320032</td>\n",
+       "      <td>no_assertion_criteria_provided</td>\n",
+       "      <td>PERM1</td>\n",
+       "      <td>train</td>\n",
+       "      <td>1</td>\n",
+       "      <td>84808</td>\n",
+       "      <td>Neutrophil inclusion bodies</td>\n",
+       "      <td>PPARGC1 and ESRR induced regulator, muscle 1</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>2</th>\n",
+       "      <td>chr1</td>\n",
+       "      <td>1050449</td>\n",
+       "      <td>G</td>\n",
+       "      <td>A</td>\n",
+       "      <td>Pathogenic</td>\n",
+       "      <td>ClinVar</td>\n",
+       "      <td>missense_variant</td>\n",
+       "      <td>1284257</td>\n",
+       "      <td>no_assertion_criteria_provided</td>\n",
+       "      <td>AGRN</td>\n",
+       "      <td>train</td>\n",
+       "      <td>1</td>\n",
+       "      <td>375790</td>\n",
+       "      <td>Congenital myasthenic syndrome 8</td>\n",
+       "      <td>agrin</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>3</th>\n",
+       "      <td>chr1</td>\n",
+       "      <td>1050575</td>\n",
+       "      <td>G</td>\n",
+       "      <td>C</td>\n",
+       "      <td>Pathogenic</td>\n",
+       "      <td>ClinVar</td>\n",
+       "      <td>missense_variant</td>\n",
+       "      <td>18241</td>\n",
+       "      <td>no_assertion_criteria_provided</td>\n",
+       "      <td>AGRN</td>\n",
+       "      <td>train</td>\n",
+       "      <td>1</td>\n",
+       "      <td>375790</td>\n",
+       "      <td>Congenital myasthenic syndrome 8</td>\n",
+       "      <td>agrin</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>4</th>\n",
+       "      <td>chr1</td>\n",
+       "      <td>1213738</td>\n",
+       "      <td>G</td>\n",
+       "      <td>A</td>\n",
+       "      <td>Pathogenic</td>\n",
+       "      <td>ClinVar</td>\n",
+       "      <td>missense_variant</td>\n",
+       "      <td>96692</td>\n",
+       "      <td>no_assertion_criteria_provided</td>\n",
+       "      <td>TNFRSF4</td>\n",
+       "      <td>train</td>\n",
+       "      <td>1</td>\n",
+       "      <td>7293</td>\n",
+       "      <td>Combined immunodeficiency due to OX40 deficiency</td>\n",
+       "      <td>TNF receptor superfamily member 4</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>...</th>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>32680</th>\n",
+       "      <td>chrY</td>\n",
+       "      <td>2787412</td>\n",
+       "      <td>C</td>\n",
+       "      <td>T</td>\n",
+       "      <td>Pathogenic</td>\n",
+       "      <td>ClinVar</td>\n",
+       "      <td>missense_variant</td>\n",
+       "      <td>9747</td>\n",
+       "      <td>no_assertion_criteria_provided</td>\n",
+       "      <td>SRY</td>\n",
+       "      <td>train</td>\n",
+       "      <td>1</td>\n",
+       "      <td>6736</td>\n",
+       "      <td>46,XY sex reversal 1</td>\n",
+       "      <td>sex determining region Y</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>32681</th>\n",
+       "      <td>chrY</td>\n",
+       "      <td>2787426</td>\n",
+       "      <td>C</td>\n",
+       "      <td>G</td>\n",
+       "      <td>Pathogenic</td>\n",
+       "      <td>ClinVar</td>\n",
+       "      <td>missense_variant</td>\n",
+       "      <td>9739</td>\n",
+       "      <td>criteria_provided,_single_submitter</td>\n",
+       "      <td>SRY</td>\n",
+       "      <td>train</td>\n",
+       "      <td>1</td>\n",
+       "      <td>6736</td>\n",
+       "      <td>not provided</td>\n",
+       "      <td>sex determining region Y</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>32682</th>\n",
+       "      <td>chrY</td>\n",
+       "      <td>2787515</td>\n",
+       "      <td>C</td>\n",
+       "      <td>A</td>\n",
+       "      <td>Pathogenic</td>\n",
+       "      <td>ClinVar</td>\n",
+       "      <td>missense_variant</td>\n",
+       "      <td>492908</td>\n",
+       "      <td>no_assertion_criteria_provided</td>\n",
+       "      <td>SRY</td>\n",
+       "      <td>train</td>\n",
+       "      <td>1</td>\n",
+       "      <td>6736</td>\n",
+       "      <td>46,XY sex reversal 1</td>\n",
+       "      <td>sex determining region Y</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>32683</th>\n",
+       "      <td>chrY</td>\n",
+       "      <td>2787551</td>\n",
+       "      <td>C</td>\n",
+       "      <td>T</td>\n",
+       "      <td>Pathogenic</td>\n",
+       "      <td>ClinVar</td>\n",
+       "      <td>missense_variant</td>\n",
+       "      <td>9754</td>\n",
+       "      <td>no_assertion_criteria_provided</td>\n",
+       "      <td>SRY</td>\n",
+       "      <td>train</td>\n",
+       "      <td>1</td>\n",
+       "      <td>6736</td>\n",
+       "      <td>46,XY sex reversal 1</td>\n",
+       "      <td>sex determining region Y</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>32684</th>\n",
+       "      <td>chrY</td>\n",
+       "      <td>7063898</td>\n",
+       "      <td>A</td>\n",
+       "      <td>T</td>\n",
+       "      <td>Pathogenic</td>\n",
+       "      <td>ClinVar</td>\n",
+       "      <td>missense_variant</td>\n",
+       "      <td>625467</td>\n",
+       "      <td>no_assertion_criteria_provided</td>\n",
+       "      <td>LOC126057105, TBL1Y</td>\n",
+       "      <td>train</td>\n",
+       "      <td>1</td>\n",
+       "      <td>126057105, 90665</td>\n",
+       "      <td>Deafness, Y-linked 2</td>\n",
+       "      <td>P300/CBP strongly-dependent group 1 enhancer G...</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "<p>32685 rows × 15 columns</p>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "      CHROM      POS REF ALT       LABEL   SOURCE       CONSEQUENCE       ID  \\\n",
+       "0      chr1   976215   A   G  Pathogenic  ClinVar  missense_variant  1320032   \n",
+       "1      chr1   976215   A   G  Pathogenic  ClinVar  missense_variant  1320032   \n",
+       "2      chr1  1050449   G   A  Pathogenic  ClinVar  missense_variant  1284257   \n",
+       "3      chr1  1050575   G   C  Pathogenic  ClinVar  missense_variant    18241   \n",
+       "4      chr1  1213738   G   A  Pathogenic  ClinVar  missense_variant    96692   \n",
+       "...     ...      ...  ..  ..         ...      ...               ...      ...   \n",
+       "32680  chrY  2787412   C   T  Pathogenic  ClinVar  missense_variant     9747   \n",
+       "32681  chrY  2787426   C   G  Pathogenic  ClinVar  missense_variant     9739   \n",
+       "32682  chrY  2787515   C   A  Pathogenic  ClinVar  missense_variant   492908   \n",
+       "32683  chrY  2787551   C   T  Pathogenic  ClinVar  missense_variant     9754   \n",
+       "32684  chrY  7063898   A   T  Pathogenic  ClinVar  missense_variant   625467   \n",
+       "\n",
+       "                             REVIEW_STATUS                 GENE  split  \\\n",
+       "0           no_assertion_criteria_provided                PERM1  train   \n",
+       "1           no_assertion_criteria_provided                PERM1  train   \n",
+       "2           no_assertion_criteria_provided                 AGRN  train   \n",
+       "3           no_assertion_criteria_provided                 AGRN  train   \n",
+       "4           no_assertion_criteria_provided              TNFRSF4  train   \n",
+       "...                                    ...                  ...    ...   \n",
+       "32680       no_assertion_criteria_provided                  SRY  train   \n",
+       "32681  criteria_provided,_single_submitter                  SRY  train   \n",
+       "32682       no_assertion_criteria_provided                  SRY  train   \n",
+       "32683       no_assertion_criteria_provided                  SRY  train   \n",
+       "32684       no_assertion_criteria_provided  LOC126057105, TBL1Y  train   \n",
+       "\n",
+       "      INT_LABEL           GENE_ID  \\\n",
+       "0             1             84808   \n",
+       "1             1             84808   \n",
+       "2             1            375790   \n",
+       "3             1            375790   \n",
+       "4             1              7293   \n",
+       "...         ...               ...   \n",
+       "32680         1              6736   \n",
+       "32681         1              6736   \n",
+       "32682         1              6736   \n",
+       "32683         1              6736   \n",
+       "32684         1  126057105, 90665   \n",
+       "\n",
+       "                                                Disease  \\\n",
+       "0               Renal tubular epithelial cell apoptosis   \n",
+       "1                           Neutrophil inclusion bodies   \n",
+       "2                      Congenital myasthenic syndrome 8   \n",
+       "3                      Congenital myasthenic syndrome 8   \n",
+       "4      Combined immunodeficiency due to OX40 deficiency   \n",
+       "...                                                 ...   \n",
+       "32680                              46,XY sex reversal 1   \n",
+       "32681                                      not provided   \n",
+       "32682                              46,XY sex reversal 1   \n",
+       "32683                              46,XY sex reversal 1   \n",
+       "32684                              Deafness, Y-linked 2   \n",
+       "\n",
+       "                                               GENE_Name  \n",
+       "0           PPARGC1 and ESRR induced regulator, muscle 1  \n",
+       "1           PPARGC1 and ESRR induced regulator, muscle 1  \n",
+       "2                                                  agrin  \n",
+       "3                                                  agrin  \n",
+       "4                      TNF receptor superfamily member 4  \n",
+       "...                                                  ...  \n",
+       "32680                           sex determining region Y  \n",
+       "32681                           sex determining region Y  \n",
+       "32682                           sex determining region Y  \n",
+       "32683                           sex determining region Y  \n",
+       "32684  P300/CBP strongly-dependent group 1 enhancer G...  \n",
+       "\n",
+       "[32685 rows x 15 columns]"
+      ]
+     },
+     "execution_count": 59,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# Display final dataset with all extracted information\n",
+    "if 'clinvar_data' in locals() and clinvar_data is not None:\n",
+    "    print(f\"📊 Final dataset shape: {clinvar_data.shape}\")\n",
+    "    print(f\"📋 Columns: {list(clinvar_data.columns)}\")\n",
+    "    \n",
+    "    # Data completeness statistics\n",
+    "    print(\"\\n📈 Data Completeness:\")\n",
+    "    for col in ['GENE', 'GENE_ID', 'GENE_Name', 'Disease']:\n",
+    "        if col in clinvar_data.columns:\n",
+    "            filled_count = (clinvar_data[col] != '').sum()\n",
+    "            print(f\"   {col}: {filled_count}/{len(clinvar_data)} ({filled_count/len(clinvar_data)*100:.1f}%)\")\n",
+    "    \n",
+    "    # Sample data\n",
+    "    print(\"\\n🔍 Sample data:\")\n",
+    "    display(clinvar_data.head())\n",
+    "    \n",
+    "    # Memory usage\n",
+    "    memory_mb = clinvar_data.memory_usage(deep=True).sum() / 1024 / 1024\n",
+    "    print(f\"\\n💾 Dataset memory usage: {memory_mb:.1f} MB\")\n",
+    "else:\n",
+    "    print(\"❌ No final data to display\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c545ae83-5cd1-4e29-87fd-69389bdb153f",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "'P300/CBP strongly-dependent group 1 enhancer GRCh37_chrY:6931456-6932655| transducin beta like 1 Y-linked'"
+      ]
+     },
+     "execution_count": 60,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# Show example of gene name mapping\n",
+    "if 'clinvar_data' in locals() and clinvar_data is not None and len(clinvar_data) > 32684:\n",
+    "    example_gene_name = clinvar_data.iloc[32684]['GENE_Name']\n",
+    "    example_gene_id = clinvar_data.iloc[32684]['GENE_ID']\n",
+    "    \n",
+    "    print(f\"🔍 Example gene mapping for row 32684:\")\n",
+    "    print(f\"   Gene ID: {example_gene_id}\")\n",
+    "    print(f\"   Gene Name: {example_gene_name}\")\n",
+    "else:\n",
+    "    # Show any available example\n",
+    "    if 'clinvar_data' in locals() and clinvar_data is not None and not clinvar_data.empty:\n",
+    "        # Find first row with gene name data\n",
+    "        example_row = clinvar_data[clinvar_data['GENE_Name'] != ''].iloc[0] if (clinvar_data['GENE_Name'] != '').any() else clinvar_data.iloc[0]\n",
+    "        \n",
+    "        print(f\"🔍 Example gene mapping:\")\n",
+    "        print(f\"   Gene ID: {example_row.get('GENE_ID', 'N/A')}\")\n",
+    "        print(f\"   Gene Name: {example_row.get('GENE_Name', 'N/A')}\")\n",
+    "    else:\n",
+    "        print(\"❌ No data available for example\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a214c29d-a4f1-4af6-a914-e6b4a14a1c49",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "\n",
+    "# Save the final processed dataset\n",
+    "if 'clinvar_data' in locals() and clinvar_data is not None:\n",
+    "    output_file = \"clinvar_with_disease.csv\"\n",
+    "    \n",
+    "    try:\n",
+    "        clinvar_data.to_csv(output_file, index=False)\n",
+    "        \n",
+    "        print(f\"✅ Final dataset saved to {output_file}\")\n",
+    "        print(f\"📊 Saved {len(clinvar_data)} records with {len(clinvar_data.columns)} columns\")\n",
+    "        \n",
+    "        # File size\n",
+    "        file_size = os.path.getsize(output_file) / 1024 / 1024\n",
+    "        print(f\"💾 File size: {file_size:.1f} MB\")\n",
+    "        \n",
+    "        # Summary of what was accomplished\n",
+    "        print(\"\\n🎯 Processing Summary:\")\n",
+    "        print(f\"   ✓ Extracted ClinVar coding variants\")\n",
+    "        print(f\"   ✓ Parsed XML records for gene information\")\n",
+    "        print(f\"   ✓ Mapped diseases/phenotypes\")\n",
+    "        print(f\"   ✓ Added human-readable gene names\")\n",
+    "        print(f\"   ✓ Created comprehensive dataset\")\n",
+    "        \n",
+    "    except Exception as e:\n",
+    "        print(f\"❌ Error saving dataset: {e}\")\n",
+    "else:\n",
+    "    print(\"⚠️ No data available to save\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b6c4c1f4-4b87-4624-8f8a-c568e40b2e63",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import shutil\n",
+    "\n",
+    "# Optional: Clean up temporary XML data directory\n",
+    "# Uncomment the following lines if you want to remove the XML files to save space\n",
+    "\n",
+    "if os.path.exists(\"data\") and os.path.isdir(\"data\"):\n",
+    "    # Count files before cleanup\n",
+    "    xml_files = [f for f in os.listdir(\"data\") if f.endswith('.xml')]\n",
+    "    \n",
+    "    print(f\"🗂️ Found {len(xml_files)} XML files in data directory\")\n",
+    "    \n",
+    "    # Uncomment to actually remove the directory\n",
+    "    # shutil.rmtree(\"data\")\n",
+    "    # print(\"🗑️ Removed temporary XML data directory\")\n",
+    "    \n",
+    "    print(\"ℹ️ XML files preserved. Uncomment the cleanup code to remove them.\")\n",
+    "else:\n",
+    "    print(\"ℹ️ No XML data directory found to clean up\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c08beea6-6ff7-4900-a8b8-8a719db36189",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "## Processing Complete ✅\n",
+    "\n",
+    "The ClinVar coding variants have been successfully processed with the following enhancements:\n",
+    "\n",
+    "### Generated Files:\n",
+    "- `clinvar_coding_raw.csv` - Raw ClinVar entries extracted from VEP data\n",
+    "- `Clinvar_ID.txt` - List of ClinVar IDs for processing\n",
+    "- `gene_id.txt` - Unique gene IDs for name mapping\n",
+    "- `gene_id_to_name.json` - Gene ID to name mapping dictionary\n",
+    "- `clinvar_with_disease.csv` - **Final comprehensive dataset**\n",
+    "\n",
+    "### Dataset Features:\n",
+    "- **Variant Information**: Genomic coordinates, alleles, and annotations\n",
+    "- **Gene Data**: Symbols, IDs, and human-readable names\n",
+    "- **Disease/Phenotype**: Associated conditions and clinical significance\n",
+    "- **Expanded Format**: One row per variant-disease combination\n",
+    "\n",
+    "### Next Steps:\n",
+    "1. **Quality Control**: Review the data for completeness and accuracy\n",
+    "2. **Analysis**: Use the dataset for downstream genetic analysis\n",
+    "3. **Integration**: Combine with other datasets as needed\n",
+    "4. **Documentation**: Update metadata and create data dictionary\n",
+    "\n",
+    "### File Cleanup:\n",
+    "- XML files in `data/` directory can be removed to save space\n",
+    "- Intermediate files can be archived or removed as needed"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.9"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

BioReason/pyproject.toml

@@ -0,0 +1,57 @@
+[build-system]
+requires = ["setuptools>=42", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "bioreason"
+version = "0.1.0"
+description = "Bio-related Reasoning with Language Models"
+readme = "README.md"
+requires-python = ">=3.11"
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.11",
+    "License :: OSI Approved :: MIT License",
+    "Operating System :: OS Independent",
+]
+dependencies = [
+    "torch",
+    "torchvision",
+    "transformers",
+    "accelerate",
+    "qwen-vl-utils",
+    "jupyter",
+    "datasets",
+    "peft",
+    "pytorch_lightning",
+    "wandb",
+    "trl[vllm]",
+    "bitsandbytes",
+    "deepspeed",
+]
+
+[project.optional-dependencies]
+dev = [
+    "pytest",
+    "black",
+    "isort",
+    "mypy",
+]
+
+[tool.setuptools]
+packages = ["bioreason"]
+
+[tool.black]
+line-length = 88
+target-version = ["py311"]
+
+[tool.isort]
+profile = "black"
+line_length = 88
+
+[tool.mypy]
+python_version = "3.11"
+warn_return_any = true
+warn_unused_configs = true
+disallow_untyped_defs = true
+disallow_incomplete_defs = true 

BioReason/reason.py

@@ -0,0 +1,636 @@
+import os
+import re
+
+import pathlib
+from argparse import ArgumentParser
+from typing import List, Dict, Optional
+from dataclasses import dataclass, field
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+from torch.optim import AdamW
+from torch.utils.data import DataLoader, Dataset
+from transformers import get_cosine_schedule_with_warmup, AutoTokenizer
+
+from transformers import (
+    AutoTokenizer,
+    AutoModelForCausalLM,
+    AutoModelForMaskedLM,
+    AutoProcessor,
+)
+
+from datasets import load_dataset, DatasetDict
+
+from peft import get_peft_model, LoraConfig, prepare_model_for_kbit_training
+from transformers import BitsAndBytesConfig
+
+import pytorch_lightning as pl
+from pytorch_lightning.callbacks import ModelCheckpoint, LearningRateMonitor
+from pytorch_lightning.loggers import WandbLogger
+
+from trl import GRPOConfig, GRPOTrainer, ModelConfig, ScriptArguments, TrlParser, get_peft_config
+#from unsloth import FastLanguageModel, is_bfloat16_supported
+
+from bioreason.models.dna_llm import DNALLMModel
+from bioreason.dna_modules import NucleotideDNAModule
+from bioreason.models.dl.processing_dl import DLProcessor
+from bioreason.trainer import DNALLMGRPOTrainer, DNALLMGRPOConfig
+from bioreason.models.evo2_tokenizer import Evo2Tokenizer, register_evo2_tokenizer
+register_evo2_tokenizer()
+
+# Custom TrainerCallback to override the saving mechanism
+from transformers import TrainerCallback, TrainerState, TrainerControl
+from transformers.trainer_utils import PREFIX_CHECKPOINT_DIR
+
+class SaveWithPyTorchCallback(TrainerCallback):
+    """Custom callback to save models with PyTorch's native save mechanism instead of safetensors"""
+    def on_save(self, args, state, control, **kwargs):
+        # Get the checkpoint folder
+        checkpoint_folder = os.path.join(
+            args.output_dir, f"{PREFIX_CHECKPOINT_DIR}-{state.global_step}"
+        )
+        os.makedirs(checkpoint_folder, exist_ok=True)
+        
+        # Save with PyTorch instead of safetensors
+        checkpoint_path = os.path.join(checkpoint_folder, "pytorch_model.bin")
+        model = kwargs.get("model")
+        
+        # Get model unwrapped from accelerator etc.
+        unwrapped_model = model.module if hasattr(model, "module") else model
+        
+        # Save using PyTorch directly
+        torch.save(unwrapped_model.state_dict(), checkpoint_path)
+        
+        # DNALLMModel doesn't have a direct config attribute, so we need to save
+        # the configs of its sub-models
+        if hasattr(unwrapped_model, "text_model"):
+            if hasattr(unwrapped_model.text_model, "config"):
+                unwrapped_model.text_model.config.save_pretrained(checkpoint_folder)
+            # Handle PEFT models which might have base_model
+            elif hasattr(unwrapped_model.text_model, "base_model") and hasattr(unwrapped_model.text_model.base_model, "config"):
+                unwrapped_model.text_model.base_model.config.save_pretrained(checkpoint_folder)
+        
+        # Print info about what's being saved
+        print(f"Saved model checkpoint to {checkpoint_folder}")
+        lora_params = [k for k in unwrapped_model.state_dict().keys() if "lora" in k]
+        print(f"Checkpoint contains {len(lora_params)} LoRA parameters")
+        
+        # Signal that we've saved
+        control.should_save = False
+        return control
+
+def _get_target_modules(model: DNALLMModel):
+    # Apply LoRA to all linear layers in the text model
+    target_modules = []
+
+    # Get all unique linear layer names
+    seen_names = set()
+    for name, module in model.text.named_modules():
+        if isinstance(module, torch.nn.Linear):
+            names = name.split(".")
+            target_name = names[-1]  # Use the last part of the name
+
+            # Skip output head but include all other linear layers
+            if target_name != "lm_head" and target_name not in seen_names:
+                target_modules.append(target_name)
+                seen_names.add(target_name)
+
+    # Add attention-specific layers
+    attention_patterns = [
+        "q_proj",
+        "k_proj",
+        "v_proj",
+        "out_proj",
+        "query",
+        "key",
+        "value",
+    ]
+    for pattern in attention_patterns:
+        if pattern not in seen_names:
+            target_modules.append(pattern)
+
+    # Return all unique layer names to apply LoRA to all layers
+    return list(target_modules)
+
+
+def extract_xml_answer(text: str) -> str:
+    # answer = text.split("<answer>")[-1]
+    # answer = answer.split("</answer>")[0]
+    answer = text.split("</think>")[-1]
+    return answer.strip()
+
+def extract_hash_answer(text: str) -> str | None:
+    if "####" not in text:
+        return None
+    return text.split("####")[1].strip()
+
+def get_kegg_questions() -> Dataset:
+    data = load_dataset('wanglab/kegg', 'default') # type: ignore
+    example_dna_sequences = ["ATCTACATGCAT", "CAGCAGCTACAG", "CATCACATCGACATCGAC"]
+    num_dna_sequences = 2 # TODO: Change to 2!
+
+    data = data.map(lambda x: { # type: ignore
+        'prompt': [
+     
+            {
+                'role': 'user',
+                'content': [
+                    *({'type': 'dna', 'text': None} for _ in range(num_dna_sequences)),
+                    {'type': 'text', 'text': x['question']},
+                ],
+            },
+        ],
+        'dna_sequences': [x['reference_sequence'], x['variant_sequence']],
+        'answer': x['answer'],
+    })  # type: ignore
+
+    return data
+
+# uncomment middle messages for 1-shot prompting
+def get_gsm8k_questions(question_prompt: str) -> Dataset:
+    data = load_dataset('openai/gsm8k', 'main') # type: ignore
+
+    example_dna_sequences = ["ATCTACATGCAT", "CAGCAGCTACAG", "CATCACATCGACATCGAC"]
+    data = data.map(lambda x: { # type: ignore
+        'prompt': [ 
+
+            {
+                'role': 'user',
+                'content': [
+                    *({'type': 'dna', 'text': None} for _ in range(len(example_dna_sequences))),
+                    {'type': 'text', 'text': 'Give me a short introduction to large language model.'}
+                ]
+            },
+        ],
+        'dna_sequences': [dna for dna in example_dna_sequences],
+        'answer': extract_hash_answer(x['answer']),
+    }) # type: ignore
+    
+    return data # type: ignore
+
+def get_gsm8k_questions_old(question_prompt: str) -> Dataset:
+    data = load_dataset('openai/gsm8k', 'main') # type: ignore
+
+    example_dna_sequences = ["ATCTACATGCAT", "CAGCAGCTACAG", "CATCACATCGACATCGAC"]
+    data = data.map(lambda x: { # type: ignore
+        'prompt': [
+            {
+                'role': 'user',
+                'content': [
+                    *({'type': 'dna', 'text': None} for _ in range(len(example_dna_sequences))),
+                    {'type': 'text', 'text': question_prompt.format(Question=x['question'])}
+                ]
+            },
+        ],
+        'dna_sequences': [dna for dna in example_dna_sequences],
+        'answer': extract_hash_answer(x['answer']),
+    }) # type: ignore
+    
+    return data # type: ignore
+
+# Reward functions
+def correctness_reward_func(prompts, completions, answer, **kwargs) -> list[float]:
+    responses = [completion[0]['content'] for completion in completions]
+    q = prompts[0][-1]['content']
+    extracted_responses = [extract_xml_answer(r) for r in responses]
+    # extracted_responses = [r.lower().replace("answer:", "").strip() for r in extracted_responses]
+    print('-'*20, f"Question:\n{q}", f"\nAnswer:\n{answer[0]}", f"\nResponse:\n{responses[0]}", f"\nExtracted:\n{extracted_responses[0]}")
+    return [2.0 if a.lower() in r.lower() else 0.0 for r, a in zip(extracted_responses, answer[0])]
+
+def less_than_4_reward_func(completions, **kwargs) -> list[float]:
+    responses = [completion[0]['content'] for completion in completions]
+    extracted_responses = [extract_xml_answer(r) for r in responses]
+    return [0.5 if len(r.split(' ')) <= 4 else 0.0 for r in extracted_responses]
+
+def strict_format_reward_func(completions, **kwargs) -> list[float]:
+    """Reward function that checks if the completion has a specific format."""
+    pattern = r"^<think>\n.*?\n</think>\n.*?\n$"
+    responses = [completion[0]["content"] for completion in completions]
+    matches = [re.match(pattern, r) for r in responses]
+    return [0.5 if match else 0.0 for match in matches]
+
+def soft_format_reward_func(completions, **kwargs) -> list[float]:
+    """Reward function that checks if the completion has a specific format."""
+    pattern = r"<think>.*?</think>\s*.*?"
+    responses = [completion[0]["content"] for completion in completions]
+    matches = [re.match(pattern, r) for r in responses]
+    return [0.5 if match else 0.0 for match in matches]
+
+def count_xml(text) -> float:
+    count = 0.0
+    if text.count("<think>\n") == 1:
+        count += 0.125
+    if text.count("\n</think>\n") == 1:
+        count += 0.125
+    return count
+
+def xmlcount_reward_func(completions, **kwargs) -> list[float]:
+    contents = [completion[0]["content"] for completion in completions]
+    return [count_xml(c) for c in contents]
+
+# Format into conversation
+def make_conversation(example):
+    return {
+        "prompt": [
+            {"role": "system", "content": SYSTEM_PROMPT},
+            {"role": "user", "content": example["problem"]},
+        ],
+    }
+
+def make_conversation_image(example):
+    return {
+        "prompt": [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "image"},
+                ],
+            },
+        ],
+    }
+
+@dataclass
+# class GRPOModelConfig(ModelConfig):
+    
+#     # "HuggingFaceTB/SmolLM-135M-Instruct"
+#     # "Qwen/Qwen2.5-0.5B-Instruct"
+#     model_name_or_path: str = field(default="Qwen/Qwen3-0.6B", metadata={"help": "Model checkpoint for weights initialization."})
+#     dna_model_name_or_path: str = field(default="InstaDeepAI/nucleotide-transformer-v2-100m-multi-species", metadata={"help": "Model checkpoint for weights initialization."})
+#     cache_dir: str = field(default=None, metadata={"help": "Path to model cache directory."})
+#     max_length_text: int = field(default=800, metadata={"help": "Maximum length of text sequences."})
+#     max_length_dna: int = field(default=800, metadata={"help": "Maximum length of DNA sequences, in groups of 6 nucleotides."})
+#     sft_checkpoint: str = field(default=None, metadata={"help": "Path to the checkpoint for SFT."})
+#     lora_r: int = field(default=32, metadata={"help": "LoRA R value."})
+#     lora_alpha: int = field(default=64, metadata={"help": "LoRA alpha."})
+#     lora_dropout: float = field(default=0.05, metadata={"help": "LoRA dropout."})
+#     lora_modules_to_save: Optional[list[str]] = field(
+#         default="embed_tokens",
+#         metadata={"help": "Model layers to unfreeze & train."},
+#     )
+#     freeze_dna_modules: bool = False
+
+class GRPOModelConfig(ModelConfig):
+
+    model_name_or_path: str = field(default="Qwen/Qwen3-0.6B", metadata={"help": "Model checkpoint for LLM weights initialization."})
+    protein_model_name_or_path: str = field(default="esm2_t33_650M_UR50D", metadata={"help": "Model checkpoint for ESM-2 protein weights initialization."})
+    cache_dir: str = field(default=None, metadata={"help": "Path to model cache directory."})
+    max_length_text: int = field(default=800, metadata={"help": "Maximum length of text sequences."})
+    max_length_protein: int = field(default=800, metadata={"help": "Maximum length of protein sequences (number of amino acids)."})
+    sft_checkpoint: str = field(default=None, metadata={"help": "Path to the checkpoint for SFT."})
+    lora_r: int = field(default=32, metadata={"help": "LoRA R value."})
+    lora_alpha: int = field(default=64, metadata={"help": "LoRA alpha."})
+    lora_dropout: float = field(default=0.05, metadata={"help": "LoRA dropout."})
+    lora_modules_to_save: Optional[list[str]] = field(
+        default_factory=lambda: ["embed_tokens", "lm_head"],
+        metadata={"help": "Model layers to unfreeze & train with LoRA."},
+    )
+    
+    # Updated: Renamed `freeze_dna_modules` to `freeze_protein_model`
+    freeze_protein_model: bool = field(default=True, metadata={"help": "Whether to freeze the ESM-2 protein model during training."})
+    num_query_tokens: int = field(default=32, metadata={"help": "The number of query tokens used by the Q-Former to summarize protein features. These tokens will be injected into the LLM input."})
+    # New: Parameters for the projector layer
+    projector_hidden_size: int = field(default=1280, metadata={"help": "Hidden size of the projector layer. It should match the ESM-2's output hidden size."})
+    projector_output_size: int = field(default=1024, metadata={"help": "Output size of the projector layer. It should match the LLM's hidden size."})
+    
+    # New: Parameter to control projector training
+    freeze_projector: bool = field(default=False, metadata={"help": "Whether to freeze the projector layer during training."})
+
+@dataclass
+class GRPOScriptArguments(ScriptArguments):
+    """
+    Script arguments for the GRPO training script.
+    """
+    dataset_name: str = field(default="wanglab/kegg", metadata={"help": "Dataset name with default."})
+    data_file_paths: str = field(
+        default=None,
+        metadata={"help": "Paths to data files, separated by ':'"},
+    )
+    arrow_cache_dir: str = field(
+        default=None,
+        metadata={"help": "Path to arrow cache directory"},
+    )
+    val_split_ratio: float = field(
+        default=0.0,
+        metadata={"help": "Ratio of validation split, default 0.0"},
+    )
+    reward_funcs: list[str] = field(
+        #default_factory=lambda: ["accuracy", "format"],
+        default_factory=lambda: ["xmlcount", "soft_format", "strict_format", "less_than_4", "correctness"],
+        #metadata={"help": "List of reward functions. Possible values: 'accuracy', 'format'"},
+        metadata={"help": "List of reward functions. Possible values: 'accuracy', 'xmlcount', 'soft_format', 'strict_format', 'less_than_4', 'correctness'"},
+    )
+    # max_pixels: Optional[int] = field(
+    #     default=12845056,
+    #     metadata={"help": "Maximum number of pixels for the image (for QwenVL)"},
+    # )
+    # min_pixels: Optional[int] = field(
+    #     default=3136,
+    #     metadata={"help": "Minimum number of pixels for the image (for QwenVL)"},
+    # )
+    # task_type: Optional[str] = field(
+    #     default=None,
+    #     metadata={"help": "Choose task type: 'default', 'gui', ..."},
+    # )
+
+
+
+reward_funcs_registry = {
+    # "accuracy": accuracy_reward,
+    # "format": format_reward,
+    "xmlcount": xmlcount_reward_func,
+    "soft_format": soft_format_reward_func,
+    "strict_format": strict_format_reward_func,
+    "less_than_4": less_than_4_reward_func,
+    "correctness": correctness_reward_func,
+}
+
+def get_vlm_module(model_name_or_path):
+    if any(mini_name in model_name_or_path.lower() for mini_name in ["qwen", "smol"]):
+        return NucleotideDNAModule
+    else:
+        raise ValueError(f"Unsupported model: {model_name_or_path}")
+    
+def _get_target_modules(model):
+    # Apply LoRA to all linear layers in the text model
+    target_modules = []
+
+    # Get all unique linear layer names
+    seen_names = set()
+    for name, module in model.text_model.named_modules():
+        if isinstance(module, torch.nn.Linear):
+            names = name.split(".")
+            target_name = names[-1]  # Use the last part of the name
+
+            # Skip output head but include all other linear layers
+            if target_name != "lm_head" and target_name not in seen_names:
+                target_modules.append(target_name)
+                seen_names.add(target_name)
+
+    # Add attention-specific layers
+    attention_patterns = [
+        "q_proj",
+        "k_proj",
+        "v_proj",
+        "out_proj",
+        "query",
+        "key",
+        "value",
+    ]
+    for pattern in attention_patterns:
+        if pattern not in seen_names:
+            target_modules.append(pattern)
+
+    # Return all unique layer names to apply LoRA to all layers
+    return list(target_modules)
+
+
+def _prep_for_training(model, training_args, dna_model_finetune: bool = False) -> LoraConfig:
+    """
+    Load and configure the DNALLMModel.
+    """
+
+    # Freeze DNA encoder parameters
+    if dna_model_finetune:
+        pass
+    else:
+        for param in model.dna_model.parameters():
+            param.requires_grad = False
+
+    target_modules = _get_target_modules(model)
+
+    lora_config = LoraConfig(
+        r=training_args.lora_r,
+        lora_alpha=training_args.lora_alpha,
+        lora_dropout=training_args.lora_dropout,
+        target_modules=target_modules,
+        init_lora_weights="gaussian",
+        bias="none",
+        task_type="CAUSAL_LM",
+    )
+
+    # Prepare text model for training
+    model.text_model = prepare_model_for_kbit_training(model.text_model)
+    model.text_model = get_peft_model(model.text_model, lora_config)
+
+    # Make projection layer trainable
+    for param in model.dna_projection.parameters():
+        param.requires_grad = True
+
+    return lora_config
+
+def main(script_args, training_args, model_args):
+
+    print(training_args.output_dir)
+    #pl.seed_everything(args.seed)
+    # os.environ["CUDA_LAUNCH_BLOCKING"] = "1"
+    torch.cuda.empty_cache()
+    torch.set_float32_matmul_precision("medium")
+
+    # Initialize model
+    # Load tokenizer for target text
+    # tokenizer = AutoTokenizer.from_pretrained(model_args.model_name_or_path)
+    # tokenizer.pad_token = tokenizer.eos_token
+
+    # Load model
+    model = DNALLMModel(
+        text_model_name=model_args.model_name_or_path,
+        dna_model_name=model_args.dna_model_name_or_path,
+        cache_dir=model_args.cache_dir,
+        max_length_text=model_args.max_length_text,
+        max_length_dna=model_args.max_length_dna,
+        text_model_finetune=True,
+        dna_model_finetune=not model_args.freeze_dna_modules,
+        debug=False,
+    )
+
+    # load checkpoint
+    if model_args.sft_checkpoint is not None:
+        print(f"Loading SFT checkpoint from {model_args.sft_checkpoint}")
+        
+        # Determine if it's a directory (PEFT format) or file (PyTorch state dict)
+        is_directory = os.path.isdir(model_args.sft_checkpoint)
+        
+        if is_directory:
+            # It's a PEFT checkpoint directory - load properly with PEFT
+            from peft import PeftModel
+            
+            # First initialize the text model with PEFT
+            print("Loading as PEFT checkpoint directory")
+            model.text_model = PeftModel.from_pretrained(
+                model.text_model,
+                model_args.sft_checkpoint,
+                is_trainable=True
+            )
+            
+            # Verify loaded adapters
+            print("Loaded LoRA adapters:", model.text_model.active_adapter)
+            
+            # Optional: Merge weights into base model
+            print("Merging SFT LoRA weights into base model...")
+            model.text_model = model.text_model.merge_and_unload()
+            print("Successfully merged SFT knowledge into base model")
+            
+        else:
+            # It's a PyTorch state dict file
+            print("Loading as PyTorch state dict file")
+            checkpoint = torch.load(model_args.sft_checkpoint)
+            
+            # replace model.text_model with text_model for all in state dict
+            def new_key(k):
+                if k.startswith("=model."): return k[6:]
+                elif k.startswith("_forward_module."): return k[len("_forward_module."):]
+                else: return k
+            
+            if "state_dict" in checkpoint:
+                magic = {new_key(k): v for k, v in checkpoint["state_dict"].items()}
+            elif "module" in checkpoint:
+                magic = {new_key(k): v for k, v in checkpoint["module"].items()}
+            elif isinstance(checkpoint, dict) and all(isinstance(k, str) for k in checkpoint.keys()):
+                # Direct state dict - the checkpoint itself is the state dict
+                print("Detected direct state dict format")
+                magic = {new_key(k): v for k, v in checkpoint.items()}
+            else:
+                raise ValueError(f"Unsupported checkpoint format: {model_args.sft_checkpoint}")
+            
+            # Handle prefix mapping for different model architectures
+            lora_prefix = False
+            for key in magic.keys():
+                if "lora" in key:
+                    lora_prefix = True
+                    break
+            
+            if lora_prefix:
+                print("Detected LoRA weights in state dict")
+                # First prepare model for LoRA training
+                _prep_for_training(model, model_args, dna_model_finetune=model_args.freeze_dna_modules)
+                
+                # Print some diagnostic info about the keys
+                model_keys = set(model.state_dict().keys())
+                checkpoint_keys = set(magic.keys())
+                print(f"Model has {len(model_keys)} keys")
+                print(f"Checkpoint has {len(checkpoint_keys)} keys")
+                
+                # Try to map LoRA keys more intelligently
+                new_magic = {}
+                for k, v in magic.items():
+                    # Try different prefix mappings based on common patterns
+                    if "base_model.model" in k and k not in model_keys:
+                        new_k = k.replace("text_model.base_model.model", "text_model")
+                        if new_k in model_keys:
+                            new_magic[new_k] = v
+                            continue
+                    
+                    # Try removing common prefixes
+                    if k.startswith("text_model.") and k not in model_keys:
+                        new_k = "text_model.base_model.model." + k[len("text_model."):]
+                        if new_k in model_keys:
+                            new_magic[new_k] = v
+                            continue
+                    
+                    # Keep original key if no mapping found
+                    new_magic[k] = v
+                
+                # Include missing target modules in diagnostic info
+                magic = new_magic
+                print(f"After key mapping: {len(magic)} keys")
+                
+                # Then load weights, allowing missing/extra keys
+                result = model.load_state_dict(magic, strict=False)
+                
+                if len(result.unexpected_keys) > 0:
+                    print(f"Sample unexpected keys: {result.unexpected_keys[:5]}")
+                if len(result.missing_keys) > 0:
+                    print(f"Sample missing keys: {result.missing_keys[:5]}")
+                    
+                print(f"Loaded checkpoint with {len(result.missing_keys)} missing keys and {len(result.unexpected_keys)} unexpected keys")
+            else:
+                print("Standard weights detected - remapping keys")
+                # Map keys to model structure
+                magic = {k.replace("text_model", "text_model.base_model.model"): v for k, v in magic.items()}
+                magic = {k.replace("dna_model", "dna_model"): v for k, v in magic.items()}
+                
+                # Fix the shared memory tensors issue by making a copy of weights
+                for key in list(magic.keys()):
+                    if 'lm_head.weight' in key:
+                        magic[key] = magic[key].clone()
+                
+                # Load weights before setting up LoRA
+                result = model.load_state_dict(magic, strict=False)
+                print(f"Loaded checkpoint with {len(result.missing_keys)} missing keys and {len(result.unexpected_keys)} unexpected keys")
+                
+                # Now prepare for LoRA training
+                _prep_for_training(model, model_args, dna_model_finetune=model_args.freeze_dna_modules)
+    else:
+        # No checkpoint, just prepare for training
+        _prep_for_training(model, model_args, dna_model_finetune=model_args.freeze_dna_modules)
+
+    # Get reward functions
+    reward_funcs = [reward_funcs_registry[func] for func in script_args.reward_funcs]
+    # reward_funcs = [
+    #         xmlcount_reward_func,
+    #         soft_format_reward_func,
+    #         strict_format_reward_func,
+    #         int_reward_func,
+    #         correctness_reward_func,
+    #     ]
+    print("reward_funcs:", reward_funcs)
+
+    vlm_module_cls = get_vlm_module(model_args.model_name_or_path)
+    print("using vlm module:", vlm_module_cls.__name__)
+    question_prompt = vlm_module_cls.get_question_template()
+
+
+    dataset = get_kegg_questions()
+    
+    #dataset = get_gsm8k_questions(question_prompt)
+
+    print(dataset)
+
+    #print('ITEM ONE OF THE DATASET', dataset['train'][0])
+
+    # Custom callback to handle saving with PyTorch's native mechanism
+    custom_save_callback = SaveWithPyTorchCallback()
+
+    # Initialize the GRPO trainer with custom callback
+    trainer = DNALLMGRPOTrainer(
+        model=model,
+        reward_funcs=reward_funcs,
+        args=training_args,
+        dna_module=vlm_module_cls(),
+        train_dataset=dataset['train'],
+        eval_dataset=dataset['val'] if training_args.eval_strategy != "no" else None,
+        peft_config=get_peft_config(model_args),
+        attn_implementation=model_args.attn_implementation,
+        torch_dtype=model_args.torch_dtype,
+        callbacks=[custom_save_callback],  # Add our custom callback
+    )
+
+    # Set the trainer to save in PyTorch format instead of safetensors
+    training_args.save_safetensors = False
+
+    # Train and push the model to the Hub
+    # if list(pathlib.Path(training_args.output_dir).glob("checkpoint-*")):
+    #     trainer.train(resume_from_checkpoint=True)
+    # else:
+    #     trainer.train()
+
+    # Train and push the model to the Hub
+    trainer.train()
+
+
+if __name__ == "__main__":
+    # os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"
+    print(f"CUDA_VISIBLE_DEVICES: {os.environ.get('CUDA_VISIBLE_DEVICES')}")
+    parser = TrlParser((GRPOScriptArguments, DNALLMGRPOConfig, GRPOModelConfig))
+    script_args, training_args, model_args = parser.parse_args_and_config()
+    
+    # Ensure we use PyTorch's save mechanism instead of safetensors
+    training_args.save_safetensors = False
+
+    main(script_args, training_args, model_args)
+    
+    # parser.add_argument("--wandb_project", type=str, default="dna-text-finetune")
+    # parser.add_argument("--wandb_entity", type=str, default="adibvafa")
+
+    # args = parser.parse_args()

BioReason/reason_protein.py

@@ -0,0 +1,696 @@
+import os
+import re
+
+import pathlib
+from argparse import ArgumentParser
+from typing import List, Dict, Optional
+from dataclasses import dataclass, field
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+from torch.optim import AdamW
+from torch.utils.data import DataLoader, Dataset
+from transformers import get_cosine_schedule_with_warmup, AutoTokenizer
+
+from transformers import (
+    AutoTokenizer,
+    AutoModelForCausalLM,
+    AutoModelForMaskedLM,
+    AutoProcessor,
+)
+
+from datasets import load_dataset, DatasetDict
+
+from peft import get_peft_model, LoraConfig, prepare_model_for_kbit_training
+from transformers import BitsAndBytesConfig
+
+import pytorch_lightning as pl
+from pytorch_lightning.callbacks import ModelCheckpoint, LearningRateMonitor
+from pytorch_lightning.loggers import WandbLogger
+
+from trl import GRPOConfig, GRPOTrainer, ModelConfig, ScriptArguments, TrlParser, get_peft_config
+#from unsloth import FastLanguageModel, is_bfloat16_supported
+
+from bioreason.models.dna_llm import DNALLMModel
+from bioreason.models.protein_llm import ProteinLLMModel
+from bioreason.dna_modules import NucleotideDNAModule
+from bioreason.models.dl.processing_dl import DLProcessor
+from bioreason.trainer import DNALLMGRPOTrainer, DNALLMGRPOConfig
+from bioreason.models.evo2_tokenizer import Evo2Tokenizer, register_evo2_tokenizer
+register_evo2_tokenizer()
+
+# Custom TrainerCallback to override the saving mechanism
+from transformers import TrainerCallback, TrainerState, TrainerControl
+from transformers.trainer_utils import PREFIX_CHECKPOINT_DIR
+
+class SaveWithPyTorchCallback(TrainerCallback):
+    """Custom callback to save models with PyTorch's native save mechanism instead of safetensors"""
+    def on_save(self, args, state, control, **kwargs):
+        # Get the checkpoint folder
+        checkpoint_folder = os.path.join(
+            args.output_dir, f"{PREFIX_CHECKPOINT_DIR}-{state.global_step}"
+        )
+        os.makedirs(checkpoint_folder, exist_ok=True)
+        
+        # Save with PyTorch instead of safetensors
+        checkpoint_path = os.path.join(checkpoint_folder, "pytorch_model.bin")
+        model = kwargs.get("model")
+        
+        # Get model unwrapped from accelerator etc.
+        unwrapped_model = model.module if hasattr(model, "module") else model
+        
+        # Save using PyTorch directly
+        torch.save(unwrapped_model.state_dict(), checkpoint_path)
+        
+        # DNALLMModel doesn't have a direct config attribute, so we need to save
+        # the configs of its sub-models
+        if hasattr(unwrapped_model, "text_model"):
+            if hasattr(unwrapped_model.text_model, "config"):
+                unwrapped_model.text_model.config.save_pretrained(checkpoint_folder)
+            # Handle PEFT models which might have base_model
+            elif hasattr(unwrapped_model.text_model, "base_model") and hasattr(unwrapped_model.text_model.base_model, "config"):
+                unwrapped_model.text_model.base_model.config.save_pretrained(checkpoint_folder)
+        
+        # Print info about what's being saved
+        print(f"Saved model checkpoint to {checkpoint_folder}")
+        lora_params = [k for k in unwrapped_model.state_dict().keys() if "lora" in k]
+        print(f"Checkpoint contains {len(lora_params)} LoRA parameters")
+        
+        # Signal that we've saved
+        control.should_save = False
+        return control
+
+def _get_target_modules(model: ProteinLLMModel):
+    # Apply LoRA to all linear layers in the text model
+    target_modules = []
+
+    # Get all unique linear layer names
+    seen_names = set()
+    for name, module in model.text.named_modules():
+        if isinstance(module, torch.nn.Linear):
+            names = name.split(".")
+            target_name = names[-1]  # Use the last part of the name
+
+            # Skip output head but include all other linear layers
+            if target_name != "lm_head" and target_name not in seen_names:
+                target_modules.append(target_name)
+                seen_names.add(target_name)
+
+    
+    # Add attention-specific layers commonly found in transformers
+    attention_patterns = [
+        "q_proj",
+        "k_proj", 
+        "v_proj",
+        "o_proj",
+        "out_proj",
+        "query",
+        "key",
+        "value",
+        "gate_proj",
+        "up_proj", 
+        "down_proj",
+    ]
+    for pattern in attention_patterns:
+        if pattern not in seen_names:
+            target_modules.append(pattern)
+
+    # Return all unique layer names to apply LoRA to all layers
+    return list(target_modules)
+
+
+def extract_xml_answer(text: str) -> str:
+    # answer = text.split("<answer>")[-1]
+    # answer = answer.split("</answer>")[0]
+    answer = text.split("</think>")[-1]
+    return answer.strip()
+
+def extract_hash_answer(text: str) -> str | None:
+    if "####" not in text:
+        return None
+    return text.split("####")[1].strip()
+
+
+def get_kegg_questions() -> Dataset:
+    data = load_dataset('wanglab/kegg', 'default') # type: ignore
+    # 修改为蛋白质序列示例
+    example_protein_sequences = ["MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLAGG", 
+                                "MKTAYIAKQRQISFVKSHFSRQLEERLGLIEVQAPILSRVGDGTQDNLSGAEKAVQVKVKALPDAQFEVVHSLAKWKRISSKLLERGKTHYPPHTMVGTGVLVTKMRVAGQEPDVQGPHAGIVVQGAGDAPVVVKPVVEMLNRMVVVVSGSAAPVVVNNNNNGAAAAAAA",
+                                "MSQVQVQVQNQALNTLVKQLGRVLLQGKGRPPLQGFRIIEQNGGDSPPPPPPPPPPPPPPPPPPPPPPPPPPPPPPPPPPPPPP"]
+    num_protein_sequences = 2
+
+    data = data.map(lambda x: { # type: ignore
+        'prompt': [
+            {
+                'role': 'user',
+                'content': [
+                    *({'type': 'protein', 'text': None} for _ in range(num_protein_sequences)),
+                    {'type': 'text', 'text': x['question']},
+                ],
+            },
+        ],
+        'protein_sequences': [example_protein_sequences[0], example_protein_sequences[1]], # 使用蛋白质序列
+        'answer': x['answer'],
+    })  # type: ignore
+
+    return data
+
+
+# uncomment middle messages for 1-shot prompting
+
+# Reward functions
+
+
+
+def correctness_reward_func(prompts, completions, answer, **kwargs) -> list[float]:
+    responses = [completion[0]['content'] for completion in completions]
+    q = prompts[0][-1]['content']
+    extracted_responses = [extract_xml_answer(r) for r in responses]
+    # extracted_responses = [r.lower().replace("answer:", "").strip() for r in extracted_responses]
+    print('-'*20, f"Question:\n{q}", f"\nAnswer:\n{answer[0]}", f"\nResponse:\n{responses[0]}", f"\nExtracted:\n{extracted_responses[0]}")
+    return [2.0 if a.lower() in r.lower() else 0.0 for r, a in zip(extracted_responses, answer[0])]
+
+def less_than_4_reward_func(completions, **kwargs) -> list[float]:
+    responses = [completion[0]['content'] for completion in completions]
+    extracted_responses = [extract_xml_answer(r) for r in responses]
+    return [0.5 if len(r.split(' ')) <= 4 else 0.0 for r in extracted_responses]
+
+def strict_format_reward_func(completions, **kwargs) -> list[float]:
+    """Reward function that checks if the completion has a specific format."""
+    pattern = r"^<think>\n.*?\n</think>\n.*?\n$"
+    responses = [completion[0]["content"] for completion in completions]
+    matches = [re.match(pattern, r) for r in responses]
+    return [0.5 if match else 0.0 for match in matches]
+
+def soft_format_reward_func(completions, **kwargs) -> list[float]:
+    """Reward function that checks if the completion has a specific format."""
+    pattern = r"<think>.*?</think>\s*.*?"
+    responses = [completion[0]["content"] for completion in completions]
+    matches = [re.match(pattern, r) for r in responses]
+    return [0.5 if match else 0.0 for match in matches]
+
+def count_xml(text) -> float:
+    count = 0.0
+    if text.count("<think>\n") == 1:
+        count += 0.125
+    if text.count("\n</think>\n") == 1:
+        count += 0.125
+    return count
+
+def xmlcount_reward_func(completions, **kwargs) -> list[float]:
+    contents = [completion[0]["content"] for completion in completions]
+    return [count_xml(c) for c in contents]
+
+#补充奖励函数
+def repeatness_reward(s: str):
+    """计算文本重复度，返回值越高表示重复度越低"""
+    def ranks(l):
+        index = {v: i for i, v in enumerate(sorted(set(l)))}
+        return [index[v] for v in l]
+    
+    def suffixArray(s):
+        line = ranks(s)
+        n, k, ans, sa = len(s), 1, line, [0] * len(s)
+        while k < n - 1:
+            line = ranks(list(zip_longest(line, islice(line, k, None), fillvalue=-1)))
+            ans, k = line, k << 1
+        for i, k in enumerate(ans):
+            sa[k] = i
+        return ans, sa
+    
+    def lcp(arr, suffixArr, inv_suff):
+        n, ans, k = len(arr), [0] * len(arr), 0
+        for i in range(n):
+            if inv_suff[i] == n - 1:
+                k = 0
+                continue
+            j = suffixArr[inv_suff[i] + 1]
+            while i + k < n and j + k < n and arr[i + k] == arr[j + k]:
+                k += 1
+            ans[inv_suff[i]] = k
+            if k > 0:
+                k -= 1
+        return ans
+    
+    arr = [ord(i) for i in s]
+    n = len(arr)
+    if n <= 1:
+        return 0
+    c, sa = suffixArray(arr)
+    cnt = sum(lcp(arr, sa, c))
+    return 1 - cnt * 2 / (n * (n + 1))
+
+def format_reward(predict_str: str) -> float:
+    """
+    格式奖励函数，严格要求输出格式为：
+    <think>...</think><answer>...</answer>
+    中间不能有多余内容
+    """
+    pattern = r'^<think>.*?</think>\s*<answer>\s*.*?\s*</answer>$'
+    return 1.0 if re.fullmatch(pattern, predict_str.strip(), re.DOTALL) else 0.0
+
+def acc_reward(predict_str: str, ground_truth) -> float:
+    """
+    准确率奖励函数
+    要求<answer>中内容与ground_truth完全一致
+    """
+    match = re.search(r'<answer>\s*([^<]*?)\s*</answer>', predict_str)
+    if not match:
+        return 0.0
+    answer_content = match.group(1).strip()
+    
+    # 处理不同类型的ground_truth
+    if isinstance(ground_truth, str):
+        return 1.0 if answer_content == ground_truth else 0.0
+    elif isinstance(ground_truth, (int, float)):
+        try:
+            # 尝试将答案转换为数字进行比较
+            return 1.0 if float(answer_content) == float(ground_truth) else 0.0
+        except ValueError:
+            # 如果转换失败，尝试字符串比较
+            return 1.0 if answer_content == str(ground_truth) else 0.0
+    else:
+        # 其他类型，转换为字符串比较
+        return 1.0 if answer_content == str(ground_truth) else 0.0
+
+# 包装函数以适配现有的奖励函数接口
+def repeatness_reward_func(completions, **kwargs) -> list[float]:
+    """重复度奖励函数包装器"""
+    responses = [completion[0]['content'] for completion in completions]
+    return [repeatness_reward(r) for r in responses]
+
+def format_reward_func(completions, **kwargs) -> list[float]:
+    """格式奖励函数包装器"""
+    responses = [completion[0]['content'] for completion in completions]
+    return [format_reward(r) for r in responses]
+
+def acc_reward_func(prompts, completions, answer, **kwargs) -> list[float]:
+    """准确率奖励函数包装器"""
+    responses = [completion[0]['content'] for completion in completions]
+    
+    # 调试信息
+    print(f"DEBUG acc_reward_func - answer type: {type(answer)}, answer: {answer}")
+    
+    # 根据现有代码的模式，answer可能是一个嵌套结构
+    try:
+        if isinstance(answer, list) and len(answer) > 0:
+            # 如果answer[0]是一个列表，说明是批次数据
+            if isinstance(answer[0], list):
+                ground_truths = answer[0]
+            else:
+                # 如果answer[0]是单个值，为所有响应使用相同的真实答案
+                ground_truths = [answer[0]] * len(responses)
+        else:
+            # 如果answer不是期望的格式，返回全0
+            print(f"DEBUG: Unexpected answer format, returning zeros")
+            return [0.0] * len(responses)
+    except (IndexError, TypeError) as e:
+        print(f"DEBUG: Error processing answer: {e}, returning zeros")
+        return [0.0] * len(responses)
+    
+    print(f"DEBUG: ground_truths: {ground_truths}")
+    
+    # 确保responses和ground_truths长度一致
+    rewards = []
+    for i, response in enumerate(responses):
+        if i < len(ground_truths):
+            reward = acc_reward(response, ground_truths[i])
+            print(f"DEBUG: response {i}: '{response[:100]}...', ground_truth: '{ground_truths[i]}', reward: {reward}")
+        else:
+            # 如果ground_truths不够长，使用第一个值
+            reward = acc_reward(response, ground_truths[0] if ground_truths else "")
+            print(f"DEBUG: response {i} (fallback): reward: {reward}")
+        rewards.append(reward)
+    
+    return rewards
+
+
+#
+
+
+
+# Format into conversation
+def make_conversation(example):
+    return {
+        "prompt": [
+            {"role": "system", "content": SYSTEM_PROMPT},
+            {"role": "user", "content": example["problem"]},
+        ],
+    }
+
+def make_conversation_image(example):
+    return {
+        "prompt": [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "image"},
+                ],
+            },
+        ],
+    }
+
+@dataclass
+
+class GRPOModelConfig(ModelConfig):
+
+    model_name_or_path: str = field(default="Qwen/Qwen3-0.6B", metadata={"help": "Model checkpoint for LLM weights initialization."})
+    protein_model_name_or_path: str = field(default="esm2_t33_650M_UR50D", metadata={"help": "Model checkpoint for ESM-2 protein weights initialization."})
+    cache_dir: str = field(default=None, metadata={"help": "Path to model cache directory."})
+    max_length_text: int = field(default=800, metadata={"help": "Maximum length of text sequences."})
+    max_length_protein: int = field(default=800, metadata={"help": "Maximum length of protein sequences (number of amino acids)."})
+    sft_checkpoint: str = field(default=None, metadata={"help": "Path to the checkpoint for SFT."})
+    lora_r: int = field(default=32, metadata={"help": "LoRA R value."})
+    lora_alpha: int = field(default=64, metadata={"help": "LoRA alpha."})
+    lora_dropout: float = field(default=0.05, metadata={"help": "LoRA dropout."})
+    lora_modules_to_save: Optional[list[str]] = field(
+        default_factory=lambda: ["embed_tokens", "lm_head"],
+        metadata={"help": "Model layers to unfreeze & train with LoRA."},
+    )
+    
+    # Updated: Renamed `freeze_dna_modules` to `freeze_protein_model`
+    freeze_protein_model: bool = field(default=True, metadata={"help": "Whether to freeze the ESM-2 protein model during training."})
+    
+    num_query_tokens: int = field(default=32, metadata={"help": "Number of query tokens for QFormer."})
+    qformer_num_layers: int = field(default=6, metadata={"help": "Number of layers in QFormer."})
+    qformer_num_heads: int = field(default=8, metadata={"help": "Number of attention heads in QFormer."})
+    qformer_dropout: float = field(default=0.1, metadata={"help": "Dropout rate for QFormer."})
+
+@dataclass
+class GRPOScriptArguments(ScriptArguments):
+    """
+    Script arguments for the GRPO training script.
+    """
+    dataset_name: str = field(default="wanglab/kegg", metadata={"help": "Dataset name with default."})
+    data_file_paths: str = field(
+        default=None,
+        metadata={"help": "Paths to data files, separated by ':'"},
+    )
+    arrow_cache_dir: str = field(
+        default=None,
+        metadata={"help": "Path to arrow cache directory"},
+    )
+    val_split_ratio: float = field(
+        default=0.0,
+        metadata={"help": "Ratio of validation split, default 0.0"},
+    )
+    reward_funcs: list[str] = field(
+        # 更新默认奖励函数列表，包含新的三个函数
+        default_factory=lambda: ["repeatness", "format", "acc", "xmlcount", "soft_format"],
+        metadata={"help": "List of reward functions. Possible values: 'repeatness', 'format', 'acc', 'xmlcount', 'soft_format', 'strict_format', 'less_than_4', 'correctness'"},
+    )
+
+reward_funcs_registry = {
+    # "accuracy": accuracy_reward,
+    # "format": format_reward,
+    "repeatness": repeatness_reward_func,
+    "format": format_reward_func,
+    "acc": acc_reward_func,
+    "xmlcount": xmlcount_reward_func,
+    "soft_format": soft_format_reward_func,
+    "strict_format": strict_format_reward_func,
+    "less_than_4": less_than_4_reward_func,
+    "correctness": correctness_reward_func,
+}
+
+def get_vlm_module(model_name_or_path):
+    if any(mini_name in model_name_or_path.lower() for mini_name in ["qwen", "smol"]):
+        # 如果你有专门的蛋白质模块，使用它
+        try:
+            from bioreason.protein_modules import ProteinModule
+            return ProteinModule
+        except ImportError:
+            # 如果没有专门的蛋白质模块，检查DNAModule是否兼容
+            print("Warning: Using NucleotideDNAModule for protein processing. Consider creating a dedicated ProteinModule.")
+            return NucleotideDNAModule
+    else:
+        raise ValueError(f"Unsupported model: {model_name_or_path}")
+
+
+def _prep_for_training(model: ProteinLLMModel, model_args, protein_model_finetune: bool = False) -> LoraConfig:
+    """
+    准备ProteinLLMModel进行训练。
+    """
+    # Freeze protein encoder parameters if not finetuning
+    if not protein_model_finetune:
+        for param in model.protein_model.parameters():
+            param.requires_grad = False
+        print("Frozen protein model parameters")
+    else:
+        print("Protein model parameters will be finetuned")
+
+    # Get target modules for LoRA
+    target_modules = _get_target_modules(model)
+    print(f"LoRA target modules: {target_modules}")
+
+    lora_config = LoraConfig(
+        r=model_args.lora_r,
+        lora_alpha=model_args.lora_alpha,
+        lora_dropout=model_args.lora_dropout,
+        target_modules=target_modules,
+        init_lora_weights="gaussian",
+        bias="none",
+        task_type="CAUSAL_LM",
+    )
+
+    # Prepare text model for training
+    model.text_model = prepare_model_for_kbit_training(model.text_model)
+    model.text_model = get_peft_model(model.text_model, lora_config)
+
+    # Make QFormer projection layer trainable
+    for param in model.protein_projection.parameters():
+        param.requires_grad = True
+    print("QFormer projection layer set as trainable")
+
+    # Print trainable parameters info
+    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    total_params = sum(p.numel() for p in model.parameters())
+    print(f"Trainable parameters: {trainable_params:,} / {total_params:,} ({100 * trainable_params / total_params:.2f}%)")
+
+    return lora_config
+
+######################################################################
+######################################################################
+
+def main(script_args, training_args, model_args):
+
+    print(training_args.output_dir)
+    #pl.seed_everything(args.seed)
+    # os.environ["CUDA_LAUNCH_BLOCKING"] = "1"
+    torch.cuda.empty_cache()
+    torch.set_float32_matmul_precision("medium")
+
+    # Initialize model
+    # Load tokenizer for target text
+    # tokenizer = AutoTokenizer.from_pretrained(model_args.model_name_or_path)
+    # tokenizer.pad_token = tokenizer.eos_token
+
+    # Load model
+    # model = ProteinLLMModel(
+    #     text_model_name=model_args.model_name_or_path,
+    #     dna_model_name=model_args.dna_model_name_or_path,
+    #     cache_dir=model_args.cache_dir,
+    #     max_length_text=model_args.max_length_text,
+    #     max_length_dna=model_args.max_length_dna,
+    #     text_model_finetune=True,
+    #     dna_model_finetune=not model_args.freeze_dna_modules,
+    #     debug=False,
+    # )
+    print("Initializing ProteinLLMModel...")
+    model = ProteinLLMModel(
+        text_model_name=model_args.model_name_or_path,
+        protein_model_name=model_args.protein_model_name_or_path,
+        biomedbert_model_name=getattr(model_args, 'biomedbert_model_name', 
+                                    "microsoft/BiomedNLP-BiomedBERT-base-uncased-abstract-fulltext"),
+        cache_dir=model_args.cache_dir,
+        max_length_text=model_args.max_length_text,
+        max_length_protein=model_args.max_length_protein,
+        text_model_finetune=True,
+        protein_model_finetune=not model_args.freeze_protein_modules,
+        biomedbert_finetune=getattr(model_args, 'biomedbert_finetune', True),  # 新增：控制BiomedBERT微调
+        # Q-Former相关参数（简化了，因为直接使用BiomedBERT）
+        qformer_num_query_tokens=getattr(model_args, 'qformer_num_query_tokens', 8),  # 重命名为qformer_num_query_tokens
+    )
+    # load checkpoint
+    if model_args.sft_checkpoint is not None:
+        print(f"Loading SFT checkpoint from {model_args.sft_checkpoint}")
+        
+        # Determine if it's a directory (PEFT format) or file (PyTorch state dict)
+        is_directory = os.path.isdir(model_args.sft_checkpoint)
+        
+        if is_directory:
+            # It's a PEFT checkpoint directory - load properly with PEFT
+            from peft import PeftModel
+            
+            # First initialize the text model with PEFT
+            print("Loading as PEFT checkpoint directory")
+            model.text_model = PeftModel.from_pretrained(
+                model.text_model,
+                model_args.sft_checkpoint,
+                is_trainable=True
+            )
+            
+            # Verify loaded adapters
+            print("Loaded LoRA adapters:", model.text_model.active_adapter)
+            
+            # Optional: Merge weights into base model
+            print("Merging SFT LoRA weights into base model...")
+            model.text_model = model.text_model.merge_and_unload()
+            print("Successfully merged SFT knowledge into base model")
+            
+        else:
+            # It's a PyTorch state dict file
+            print("Loading as PyTorch state dict file")
+            checkpoint = torch.load(model_args.sft_checkpoint, map_location="cpu")
+            
+            # replace model.text_model with text_model for all in state dict
+            def new_key(k):
+                if k.startswith("=model."): return k[6:]
+                elif k.startswith("_forward_module."): return k[len("_forward_module."):]
+                else: return k
+            
+            if "state_dict" in checkpoint:
+                magic = {new_key(k): v for k, v in checkpoint["state_dict"].items()}
+            elif "module" in checkpoint:
+                magic = {new_key(k): v for k, v in checkpoint["module"].items()}
+            elif isinstance(checkpoint, dict) and all(isinstance(k, str) for k in checkpoint.keys()):
+                # Direct state dict - the checkpoint itself is the state dict
+                print("Detected direct state dict format")
+                magic = {new_key(k): v for k, v in checkpoint.items()}
+            else:
+                raise ValueError(f"Unsupported checkpoint format: {model_args.sft_checkpoint}")
+            
+            # Handle prefix mapping for different model architectures
+            lora_prefix = any("lora" in key for key in state_dict.keys())
+            
+            if lora_prefix:
+                print("Detected LoRA weights in state dict")
+                # First prepare model for LoRA training
+                _prep_for_training(model, model_args, protein_model_finetune=model_args.freeze_protein_modules)
+                
+                # Print diagnostic info
+                model_keys = set(model.state_dict().keys())
+                checkpoint_keys = set(state_dict.keys())
+                print(f"Model has {len(model_keys)} keys")
+                print(f"Checkpoint has {len(checkpoint_keys)} keys")
+                
+                # Intelligent key mapping for different prefixes
+                new_state_dict = {}
+                for k, v in state_dict.items():
+                    # Handle different common prefix patterns
+                    if "base_model.model" in k and k not in model_keys:
+                        new_k = k.replace("text_model.base_model.model", "text_model")
+                        if new_k in model_keys:
+                            new_state_dict[new_k] = v
+                            continue
+                    
+                    # Try removing/adding prefixes
+                    if k.startswith("text_model.") and k not in model_keys:
+                        new_k = "text_model.base_model.model." + k[len("text_model."):]
+                        if new_k in model_keys:
+                            new_state_dict[new_k] = v
+                            continue
+                    
+                    # Keep original key
+                    new_state_dict[k] = v
+                
+                state_dict = new_state_dict
+                print(f"After key mapping: {len(state_dict)} keys")
+                
+                # Load state dict with missing/unexpected keys allowed
+                result = model.load_state_dict(state_dict, strict=False)
+                
+                if len(result.unexpected_keys) > 0:
+                    print(f"Sample unexpected keys: {result.unexpected_keys[:5]}")
+                if len(result.missing_keys) > 0:
+                    print(f"Sample missing keys: {result.missing_keys[:5]}")
+                    
+                print(f"Loaded checkpoint with {len(result.missing_keys)} missing keys and {len(result.unexpected_keys)} unexpected keys")
+            else:
+                print("Standard weights detected - loading before LoRA setup")
+                
+                # Handle shared memory issue for embedding weights
+                for key in list(state_dict.keys()):
+                    if 'lm_head.weight' in key:
+                        state_dict[key] = state_dict[key].clone()
+                
+                # Load weights before setting up LoRA
+                result = model.load_state_dict(state_dict, strict=False)
+                print(f"Loaded checkpoint with {len(result.missing_keys)} missing keys and {len(result.unexpected_keys)} unexpected keys")
+                
+                # Now prepare for LoRA training
+                _prep_for_training(model, model_args, protein_model_finetune=model_args.freeze_protein_modules)
+    else:
+        # No checkpoint, just prepare for training
+        _prep_for_training(model, model_args, protein_model_finetune=not model_args.freeze_protein_model)
+
+    # Get reward functions
+    reward_funcs = [reward_funcs_registry[func] for func in script_args.reward_funcs]
+    # reward_funcs = [
+    #         xmlcount_reward_func,
+    #         soft_format_reward_func,
+    #         strict_format_reward_func,
+    #         int_reward_func,
+    #         correctness_reward_func,
+    #     ]
+    print("reward_funcs:", [func.__name__ for func in reward_funcs])
+
+    vlm_module_cls = get_vlm_module(model_args.model_name_or_path)
+    print("using vlm module:", vlm_module_cls.__name__)
+    question_prompt = vlm_module_cls.get_question_template()
+
+
+    dataset = get_kegg_questions()
+    
+    #dataset = get_gsm8k_questions(question_prompt)
+
+    print(dataset)
+
+    #print('ITEM ONE OF THE DATASET', dataset['train'][0])
+
+    # Custom callback to handle saving with PyTorch's native mechanism
+    custom_save_callback = SaveWithPyTorchCallback()
+
+    # Initialize the GRPO trainer with custom callback
+    trainer = DNALLMGRPOTrainer(
+        model=model,
+        reward_funcs=reward_funcs,
+        args=training_args,
+        dna_module=vlm_module_cls(),
+        train_dataset=dataset['train'],
+        eval_dataset=dataset['val'] if training_args.eval_strategy != "no" else None,
+        peft_config=get_peft_config(model_args),
+        attn_implementation=model_args.attn_implementation,
+        torch_dtype=model_args.torch_dtype,
+        callbacks=[custom_save_callback],  # Add our custom callback
+    )
+
+    # Set the trainer to save in PyTorch format instead of safetensors
+    training_args.save_safetensors = False
+
+    # Train and push the model to the Hub
+    # if list(pathlib.Path(training_args.output_dir).glob("checkpoint-*")):
+    #     trainer.train(resume_from_checkpoint=True)
+    # else:
+    #     trainer.train()
+
+    # Train and push the model to the Hub
+    trainer.train()
+
+
+if __name__ == "__main__":
+    # os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"
+    print(f"CUDA_VISIBLE_DEVICES: {os.environ.get('CUDA_VISIBLE_DEVICES')}")
+    parser = TrlParser((GRPOScriptArguments, DNALLMGRPOConfig, GRPOModelConfig))
+    script_args, training_args, model_args = parser.parse_args_and_config()
+    
+    # Ensure we use PyTorch's save mechanism instead of safetensors
+    training_args.save_safetensors = False
+
+    main(script_args, training_args, model_args)
+    
+    # parser.add_argument("--wandb_project", type=str, default="dna-text-finetune")
+    # parser.add_argument("--wandb_entity", type=str, default="adibvafa")
+
+    # args = parser.parse_args()

BioReason/requirements.txt

@@ -0,0 +1,13 @@
+torch
+torchvision
+transformers
+accelerate
+qwen-vl-utils
+jupyter
+datasets
+peft
+pytorch_lightning
+wandb
+trl[vllm]
+bitsandbytes
+deepspeed

BioReason/sh_reason.sh

@@ -0,0 +1,57 @@
+#!/bin/bash
+#SBATCH --job-name=Qwen3_1.7B_SFT_RL    # Name of the job
+#SBATCH --gres=gpu:4    # Number of GPUs
+#SBATCH -p a100    # Partition
+#SBATCH -c 12    # Number of cores
+#SBATCH --time=12:00:00    # Time limit
+#SBATCH --mem=128gb    # Memory limit
+#SBATCH --output=Qwen3_1.7B_SFT_RL_a100-%j.out    # Output file
+#SBATCH --error=Qwen3_1.7B_SFT_RL_a100-%j.err    # Error file
+
+## Environment Setup
+echo "CUDA_HOME: $CUDA_HOME"
+echo "PATH: $PATH"
+echo "LD_LIBRARY_PATH: $LD_LIBRARY_PATH"
+echo "which python: $(which python)"
+
+## Configuration Variables
+# Change these to match your setup
+SFT_CHECKPOINT=SFT_CHECKPOINT    # Change to the checkpoint of the SFT model
+CACHE_DIR=CACHE_DIR    # Change to the directory where the model weights are cached
+OUTPUT_DIR=OUTPUT_DIR    # Change to the directory where the model will be saved
+CONDA_ENV=CONDA_ENV    # Change to the conda environment
+
+## Setup Environment
+conda activate $CONDA_ENV    # Change to the conda environment
+cd .../BioReason/    # Change to the directory containing the script
+nvidia-smi    # Check GPU status
+
+## Dependencies
+# You might need to install this on a gpu session
+# pip install trl[vllm]
+
+## =============================================================================
+## Reinforcement Learning Training with DeepSpeed
+## =============================================================================
+
+# Run with DeepSpeed ZeRO Stage 2
+srun deepspeed --num_gpus=4 --num_nodes=1 \
+    reason.py \
+    --deepspeed grpo_trainer_lora_model/ds_config_stage2.json \
+    --num_generations 4 \
+    --per_device_train_batch_size 2 \
+    --bf16 true \
+    --ddp_find_unused_parameters false \
+    --sft_checkpoint $SFT_CHECKPOINT \
+    --model_name_or_path Qwen/Qwen3-1.7B \
+    --dna_model_name_or_path InstaDeepAI/nucleotide-transformer-v2-500m-multi-species \
+    --cache_dir $CACHE_DIR \
+    --output_dir $OUTPUT_DIR \
+    --save_strategy "steps" \
+    --save_steps 100 \
+    --save_total_limit 2 \
+    --use_vllm true \
+    --temperature 0.6 \
+    --top_p 0.95 \
+    --top_k 20 \
+    --num_train_epochs 1

BioReason/sh_train_dna_only.sh

@@ -0,0 +1,138 @@
+#!/bin/bash
+#SBATCH --job-name=train_dna    # Name of the job
+#SBATCH --time=8:00:00    # Time limit
+#SBATCH --partition=gpu_batch    # Partition
+#SBATCH --gpus=1    # Number of GPUs
+#SBATCH --ntasks=1    # Number of tasks
+#SBATCH --cpus-per-task=6    # Number of cores
+#SBATCH --mem=128gb    # Memory limit
+#SBATCH --output=train_dna_%j_%x.out    # Output file
+#SBATCH --error=train_dna_%j_%x.err    # Error file
+
+## Environment Setup
+echo "CUDA_HOME: $CUDA_HOME"
+echo "PATH: $PATH"
+echo "LD_LIBRARY_PATH: $LD_LIBRARY_PATH"
+echo "which python: $(which python)"
+
+## Configuration Variables
+# Change these to match your setup
+CONDA_ENV=CONDA_ENV                     # Change to your conda environment name
+CACHE_DIR=CACHE_DIR                     # Change to your HuggingFace cache directory
+WANDB_PROJECT=WANDB_PROJECT             # Change to your W&B project name
+
+## Setup Environment
+conda activate $CONDA_ENV               # Change to your conda environment
+cd .../BioReason/    # Change to the directory containing the script
+nvidia-smi                             # Check GPU status
+
+
+## =============================================================================
+## KEGG Dataset Training (DNA-only models)
+## =============================================================================
+
+# NT-500M on KEGG
+stdbuf -oL -eL srun python train_dna_only.py \
+    --cache_dir $CACHE_DIR \
+    --wandb_project $WANDB_PROJECT \
+    --dna_model_name InstaDeepAI/nucleotide-transformer-v2-500m-multi-species \
+    --strategy ddp \
+    --max_epochs 5 \
+    --num_gpus 1 \
+    --batch_size 1 \
+    --max_length_dna 2048 \
+    --truncate_dna_per_side 1024 \
+    --train_just_classifier True \
+    --learning_rate 3e-4 \
+    --dataset_type kegg \
+    --merge_val_test_set True
+
+# EVO2-1B on KEGG
+stdbuf -oL -eL srun python train_dna_only.py \
+    --cache_dir $CACHE_DIR \
+    --wandb_project $WANDB_PROJECT \
+    --dna_model_name evo2_1b_base \
+    --strategy ddp \
+    --max_epochs 5 \
+    --num_gpus 1 \
+    --batch_size 1 \
+    --max_length_dna 2048 \
+    --truncate_dna_per_side 1024 \
+    --train_just_classifier True \
+    --dna_is_evo2 True \
+    --dna_embedding_layer blocks.20.mlp.l3 \
+    --learning_rate 3e-4 \
+    --dataset_type kegg \
+    --merge_val_test_set True
+
+## =============================================================================
+## Variant Effect Prediction (VEP) Training
+## =============================================================================
+
+# NT-500M on VEP
+stdbuf -oL -eL srun python train_dna_only.py \
+    --cache_dir $CACHE_DIR \
+    --wandb_project $WANDB_PROJECT \
+    --dna_model_name InstaDeepAI/nucleotide-transformer-v2-500m-multi-species \
+    --strategy ddp \
+    --max_epochs 3 \
+    --num_gpus 1 \
+    --batch_size 2 \
+    --max_length_dna 2048 \
+    --truncate_dna_per_side 1024 \
+    --train_just_classifier True \
+    --learning_rate 3e-4 \
+    --dataset_type variant_effect_coding
+
+# EVO2-1B on VEP
+stdbuf -oL -eL srun python train_dna_only.py \
+    --cache_dir $CACHE_DIR \
+    --wandb_project $WANDB_PROJECT \
+    --dna_model_name evo2_1b_base \
+    --strategy ddp \
+    --max_epochs 3 \
+    --num_gpus 1 \
+    --batch_size 2 \
+    --max_length_dna 2048 \
+    --truncate_dna_per_side 1024 \
+    --train_just_classifier True \
+    --dna_is_evo2 True \
+    --dna_embedding_layer blocks.20.mlp.l3 \
+    --learning_rate 3e-4 \
+    --dataset_type variant_effect_coding
+
+## =============================================================================
+## Variant Effect Prediction Non-SNV Training
+## =============================================================================
+
+# NT-500M on VEP Non-SNV
+stdbuf -oL -eL srun python train_dna_only.py \
+    --cache_dir $CACHE_DIR \
+    --wandb_project $WANDB_PROJECT \
+    --dna_model_name InstaDeepAI/nucleotide-transformer-v2-500m-multi-species \
+    --strategy ddp \
+    --max_epochs 3 \
+    --num_gpus 1 \
+    --batch_size 2 \
+    --max_length_dna 2048 \
+    --truncate_dna_per_side 1024 \
+    --train_just_classifier True \
+    --learning_rate 3e-4 \
+    --dataset_type variant_effect_non_snv
+
+# EVO2-1B on VEP Non-SNV
+stdbuf -oL -eL srun python train_dna_only.py \
+    --cache_dir $CACHE_DIR \
+    --wandb_project $WANDB_PROJECT \
+    --dna_model_name evo2_1b_base \
+    --strategy ddp \
+    --max_epochs 3 \
+    --num_gpus 1 \
+    --batch_size 2 \
+    --max_length_dna 2048 \
+    --truncate_dna_per_side 1024 \
+    --train_just_classifier True \
+    --dna_is_evo2 True \
+    --dna_embedding_layer blocks.20.mlp.l3 \
+    --learning_rate 3e-4 \
+    --dataset_type variant_effect_non_snv

BioReason/sh_train_dna_qwen.sh

@@ -0,0 +1,191 @@
+#!/bin/bash
+#SBATCH --job-name=train_dna_qwen    # Name of the job
+#SBATCH --time=12:00:00    # Time limit
+#SBATCH --partition=gpu_batch    # Partition
+#SBATCH --gpus=1    # Number of GPUs
+#SBATCH --ntasks=1    # Number of tasks
+#SBATCH --cpus-per-task=8    # Number of cores
+#SBATCH --mem=128gb    # Memory limit
+#SBATCH --output=train_dna_qwen_%j_%x.out    # Output file
+#SBATCH --error=train_dna_qwen_%j_%x.err    # Error file
+
+## Environment Setup
+echo "CUDA_HOME: $CUDA_HOME"
+echo "PATH: $PATH"
+echo "LD_LIBRARY_PATH: $LD_LIBRARY_PATH"
+echo "which python: $(which python)"
+
+## Configuration Variables
+# Change these to match your setup
+CONDA_ENV=CONDA_ENV                     # Change to your conda environment name
+CACHE_DIR=CACHE_DIR                     # Change to your HuggingFace cache directory
+OUTPUT_DIR=OUTPUT_DIR                   # Change to your output/log directory
+WANDB_PROJECT=WANDB_PROJECT             # Change to your W&B project name
+
+## Setup Environment
+conda activate $CONDA_ENV               # Change to your conda environment
+cd .../BioReason/    # Change to the directory containing the script
+nvidia-smi                             # Check GPU status
+
+
+## =============================================================================
+## KEGG Dataset Training
+## =============================================================================
+
+# NT-500M + Qwen3-1.7B on KEGG
+stdbuf -oL -eL srun python train_dna_qwen.py \
+    --cache_dir $CACHE_DIR \
+    --wandb_project $WANDB_PROJECT \
+    --text_model_name Qwen/Qwen3-1.7B \
+    --dna_model_name InstaDeepAI/nucleotide-transformer-v2-500m-multi-species \
+    --strategy deepspeed_stage_2 \
+    --max_epochs 5 \
+    --num_gpus 1 \
+    --batch_size 1 \
+    --model_type dna-llm \
+    --dataset_type kegg \
+    --merge_val_test_set True \
+    --return_answer_in_batch True
+
+# EVO2-1B + Qwen3-1.7B on KEGG
+stdbuf -oL -eL srun python train_dna_qwen.py \
+    --cache_dir $CACHE_DIR \
+    --wandb_project $WANDB_PROJECT \
+    --text_model_name Qwen/Qwen3-1.7B \
+    --dna_model_name evo2_1b_base \
+    --strategy deepspeed_stage_2 \
+    --max_epochs 5 \
+    --num_gpus 1 \
+    --batch_size 1 \
+    --model_type dna-llm \
+    --dataset_type kegg \
+    --max_length_dna 2048 \
+    --truncate_dna_per_side 1024 \
+    --dna_is_evo2 True \
+    --dna_embedding_layer blocks.20.mlp.l3 \
+    --merge_val_test_set True \
+    --return_answer_in_batch True
+
+# Qwen3-4B on KEGG (LLM-only)
+stdbuf -oL -eL srun python train_dna_qwen.py \
+    --cache_dir $CACHE_DIR \
+    --wandb_project $WANDB_PROJECT \
+    --text_model_name Qwen/Qwen3-4B \
+    --dna_model_name InstaDeepAI/nucleotide-transformer-v2-500m-multi-species \
+    --strategy deepspeed_stage_2 \
+    --max_epochs 5 \
+    --num_gpus 1 \
+    --batch_size 1 \
+    --model_type llm \
+    --dataset_type kegg \
+    --max_length_dna 4 \
+    --max_length_text 8192 \
+    --truncate_dna_per_side 1024 \
+    --merge_val_test_set True \
+    --return_answer_in_batch True
+
+## =============================================================================
+## Variant Effect Prediction (VEP) Training
+## =============================================================================
+
+# NT-500M + Qwen3-4B on VEP
+stdbuf -oL -eL srun python train_dna_qwen.py \
+    --cache_dir $CACHE_DIR \
+    --wandb_project $WANDB_PROJECT \
+    --text_model_name Qwen/Qwen3-4B \
+    --dna_model_name InstaDeepAI/nucleotide-transformer-v2-500m-multi-species \
+    --strategy deepspeed_stage_2 \
+    --max_epochs 3 \
+    --num_gpus 1 \
+    --batch_size 2 \
+    --model_type dna-llm \
+    --dataset_type variant_effect_coding \
+    --return_answer_in_batch True
+
+# EVO2-1B + Qwen3-1.7B on VEP
+stdbuf -oL -eL srun python train_dna_qwen.py \
+    --cache_dir $CACHE_DIR \
+    --wandb_project $WANDB_PROJECT \
+    --text_model_name Qwen/Qwen3-1.7B \
+    --dna_model_name evo2_1b_base \
+    --strategy deepspeed_stage_2 \
+    --max_epochs 3 \
+    --num_gpus 1 \
+    --batch_size 2 \
+    --model_type dna-llm \
+    --dataset_type variant_effect_coding \
+    --max_length_dna 2048 \
+    --truncate_dna_per_side 1024 \
+    --dna_is_evo2 True \
+    --dna_embedding_layer blocks.20.mlp.l3 \
+    --return_answer_in_batch True
+
+# Qwen3-4B on VEP (LLM-only) - Testing max length text
+stdbuf -oL -eL srun python train_dna_qwen.py \
+    --cache_dir $CACHE_DIR \
+    --wandb_project $WANDB_PROJECT \
+    --text_model_name Qwen/Qwen3-4B \
+    --dna_model_name InstaDeepAI/nucleotide-transformer-v2-500m-multi-species \
+    --strategy deepspeed_stage_2 \
+    --max_epochs 3 \
+    --num_gpus 1 \
+    --batch_size 2 \
+    --model_type llm \
+    --dataset_type variant_effect_coding \
+    --max_length_dna 4 \
+    --max_length_text 4096 \
+    --truncate_dna_per_side 1024 \
+    --return_answer_in_batch True
+
+## =============================================================================
+## Variant Effect Prediction Non-SNV Training
+## =============================================================================
+
+# NT-500M + Qwen3-4B on VEP Non-SNV
+stdbuf -oL -eL srun python train_dna_qwen.py \
+    --cache_dir $CACHE_DIR \
+    --wandb_project $WANDB_PROJECT \
+    --text_model_name Qwen/Qwen3-4B \
+    --dna_model_name InstaDeepAI/nucleotide-transformer-v2-500m-multi-species \
+    --strategy deepspeed_stage_2 \
+    --max_epochs 1 \
+    --num_gpus 1 \
+    --batch_size 2 \
+    --model_type dna-llm \
+    --dataset_type variant_effect_non_snv \
+    --return_answer_in_batch True
+
+# EVO2-1B + Qwen3-4B on VEP Non-SNV
+stdbuf -oL -eL srun python train_dna_qwen.py \
+    --cache_dir $CACHE_DIR \
+    --wandb_project $WANDB_PROJECT \
+    --text_model_name Qwen/Qwen3-4B \
+    --dna_model_name evo2_1b_base \
+    --strategy deepspeed_stage_2 \
+    --max_epochs 3 \
+    --num_gpus 1 \
+    --batch_size 2 \
+    --model_type dna-llm \
+    --dataset_type variant_effect_non_snv \
+    --max_length_dna 2048 \
+    --truncate_dna_per_side 1024 \
+    --dna_is_evo2 True \
+    --dna_embedding_layer blocks.20.mlp.l3 \
+    --return_answer_in_batch True
+
+# Qwen3-4B on VEP Non-SNV (LLM-only) - Testing max length text
+stdbuf -oL -eL srun python train_dna_qwen.py \
+    --cache_dir $CACHE_DIR \
+    --wandb_project $WANDB_PROJECT \
+    --text_model_name Qwen/Qwen3-4B \
+    --dna_model_name InstaDeepAI/nucleotide-transformer-v2-500m-multi-species \
+    --strategy deepspeed_stage_2 \
+    --max_epochs 1 \
+    --num_gpus 1 \
+    --batch_size 2 \
+    --model_type llm \
+    --dataset_type variant_effect_non_snv \
+    --max_length_dna 4 \
+    --max_length_text 4096 \
+    --truncate_dna_per_side 1024 \
+    --return_answer_in_batch True

BioReason/train_dna_only.py

@@ -0,0 +1,502 @@
+import os
+import time
+import argparse
+import torch
+import wandb
+from torch.optim import AdamW
+from torch.utils.data import DataLoader
+from transformers import get_cosine_schedule_with_warmup, AutoTokenizer
+from datasets import load_dataset, concatenate_datasets
+import pytorch_lightning as pl
+from pytorch_lightning.callbacks import ModelCheckpoint, LearningRateMonitor
+from pytorch_lightning.loggers import WandbLogger
+from pytorch_lightning.strategies import DeepSpeedStrategy
+from bioreason.models.dna_only import DNAClassifierModel
+from bioreason.dataset.utils import truncate_dna
+from bioreason.dataset.kegg import dna_collate_fn
+from bioreason.dataset.variant_effect import clean_variant_effect_example
+from bioreason.models.evo2_tokenizer import Evo2Tokenizer, register_evo2_tokenizer
+register_evo2_tokenizer()
+
+
+class DNAClassifierModelTrainer(pl.LightningModule):
+    """
+    PyTorch Lightning module for training the DNA classifier.
+    """
+
+    def __init__(self, args):
+        """
+        Initialize the DNAClassifierModelTrainer.
+
+        Args:
+            args: Command line arguments
+        """
+        super().__init__()
+        self.save_hyperparameters(args)
+
+        # Load dataset and labels
+        self.dataset, self.labels = self.load_dataset()
+        self.label2id = {label: i for i, label in enumerate(self.labels)}
+
+        # Load model
+        self.dna_model = DNAClassifierModel(
+            dna_model_name=self.hparams.dna_model_name,
+            cache_dir=self.hparams.cache_dir,
+            max_length_dna=self.hparams.max_length_dna,
+            num_classes=len(self.labels),
+            dna_is_evo2=self.hparams.dna_is_evo2,
+            dna_embedding_layer=self.hparams.dna_embedding_layer,
+            train_just_classifier=self.hparams.train_just_classifier,
+        )
+        self.dna_tokenizer = self.dna_model.dna_tokenizer
+
+        # Set the training mode for the classifier and pooler
+        self.dna_model.pooler.train()
+        self.dna_model.classifier.train()
+
+        # Freeze the DNA model parameters
+        if self.hparams.dna_is_evo2:
+            self.dna_model_params = self.dna_model.dna_model.model.parameters()
+        else:
+            self.dna_model_params = self.dna_model.dna_model.parameters()
+
+        if self.hparams.train_just_classifier:
+            for param in self.dna_model_params:
+                param.requires_grad = False
+
+    def _step(self, prefix, batch_idx, batch):
+        """
+        Performs a single training/validation step.
+
+        Args:
+            batch: Dictionary containing the batch data
+            prefix: String indicating the step type ('train' or 'val')
+
+        Returns:
+            torch.Tensor: The computed loss for this batch
+        """
+        ref_ids = batch["ref_ids"].to(self.device)
+        alt_ids = batch["alt_ids"].to(self.device)
+        ref_attention_mask = batch["ref_attention_mask"].to(self.device)
+        alt_attention_mask = batch["alt_attention_mask"].to(self.device)
+        labels = batch["labels"].to(self.device)
+
+        # Forward pass
+        logits = self.dna_model(ref_ids=ref_ids, alt_ids=alt_ids, ref_attention_mask=ref_attention_mask, alt_attention_mask=alt_attention_mask)
+
+        # Calculate loss
+        loss_fn = torch.nn.CrossEntropyLoss()
+        loss = loss_fn(logits, labels)
+
+        # Calculate accuracy
+        preds = torch.argmax(logits, dim=1)
+        acc = (preds == labels).float().mean()
+
+        # Calculate F1 score, precision, and recall for binary classification
+        # Assuming label 1 is positive and label 0 is negative as mentioned
+        true_positives = ((preds == 1) & (labels == 1)).float().sum()
+        false_positives = ((preds == 1) & (labels == 0)).float().sum()
+        false_negatives = ((preds == 0) & (labels == 1)).float().sum()
+        
+        # Calculate precision, recall, and F1 score
+        precision = true_positives / (true_positives + false_positives + 1e-8)  # add small epsilon to avoid division by zero
+        recall = true_positives / (true_positives + false_negatives + 1e-8)
+        f1 = 2 * precision * recall / (precision + recall + 1e-8)
+
+        # Logging metrics
+        self.log(
+            f"{prefix}_loss",
+            loss,
+            on_step=True,
+            on_epoch=False,
+            prog_bar=True,
+            logger=True,
+        )
+        self.log(
+            f"{prefix}_acc",
+            acc,
+            on_step=True,
+            on_epoch=False,
+            prog_bar=True,
+            logger=True,
+        )
+        self.log(
+            f"{prefix}_loss_epoch",
+            loss,
+            on_step=False,
+            on_epoch=True,
+            prog_bar=True,
+            logger=True,
+            sync_dist=True,
+        )
+        self.log(
+            f"{prefix}_acc_epoch",
+            acc,
+            on_step=False,
+            on_epoch=True,
+            prog_bar=True,
+            logger=True,
+            sync_dist=True,
+        )
+        self.log(
+            f"{prefix}_precision",
+            precision,
+            on_step=True,
+            on_epoch=False,
+            prog_bar=True,
+            logger=True,
+        )
+        self.log(
+            f"{prefix}_precision_epoch",
+            precision,
+            on_step=False,
+            on_epoch=True,
+            prog_bar=True,
+            logger=True,
+            sync_dist=True,
+        )
+        self.log(
+            f"{prefix}_recall",
+            recall,
+            on_step=True,
+            on_epoch=False,
+            prog_bar=True,
+            logger=True,
+        )
+        self.log(
+            f"{prefix}_recall_epoch",
+            recall,
+            on_step=False,
+            on_epoch=True,
+            prog_bar=True,
+            logger=True,
+            sync_dist=True,
+        )
+        self.log(
+            f"{prefix}_f1",
+            f1,
+            on_step=True,
+            on_epoch=False,
+            prog_bar=True,
+            logger=True,
+        )
+        self.log(
+            f"{prefix}_f1_epoch",
+            f1,
+            on_step=False,
+            on_epoch=True,
+            prog_bar=True,
+            logger=True,
+            sync_dist=True,
+        )
+
+        if (prefix == "test") or (prefix == "train" and (self.global_step % 1000 == 0)) or (prefix == "val" and (batch_idx % 100 == 0)):
+            wandb_logger = self.logger.experiment
+                
+            pred_label = self.labels[preds[0]]
+            true_label = self.labels[labels[0]]
+            timestamp = time.time()
+            step_id = f"gen_{self.global_step}-{timestamp}"
+
+            wandb_logger.log(
+                {
+                    step_id: wandb.Table(
+                        columns=["timestamp", "prefix", "pred_label", "true_label"],
+                        data=[[timestamp, prefix, pred_label, true_label]],
+                    )
+                }
+            )
+            
+            print(f"Example {prefix} {batch_idx} {self.global_step}: Prediction: {pred_label}, Target: {true_label}")
+
+        return loss
+
+    def training_step(self, batch, batch_idx):
+        """Perform a training step."""
+        return self._step(prefix="train", batch_idx=batch_idx, batch=batch)
+
+    def validation_step(self, batch, batch_idx):
+        """Perform a validation step."""
+        return self._step(prefix="val", batch_idx=batch_idx, batch=batch)
+    
+    def test_step(self, batch, batch_idx):
+        """Perform a test step."""
+        return self._step(prefix="test", batch_idx=batch_idx, batch=batch)
+
+    def configure_optimizers(self):
+        """Configure optimizers and learning rate schedulers."""
+        # Only include parameters that require gradients
+        classifier_params = [
+            {
+                "params": self.dna_model.classifier.parameters(),
+                "lr": self.hparams.learning_rate,
+            },
+            {
+                "params": self.dna_model.pooler.parameters(),
+                "lr": self.hparams.learning_rate,
+            }
+        ]
+        dna_model_params = [
+            {
+                "params": self.dna_model_params,
+                "lr": self.hparams.learning_rate * 0.1,
+            },
+        ]
+
+        if self.hparams.train_just_classifier:
+            # Only train classifier parameters
+            optimizer = AdamW(
+                classifier_params,
+                weight_decay=self.hparams.weight_decay,
+            )
+        else:
+            # Train both DNA model and classifier with different learning rates
+            optimizer = AdamW(
+                classifier_params + dna_model_params,
+                weight_decay=self.hparams.weight_decay,
+            )
+
+        # Get total steps from trainer's estimated stepping batches
+        total_steps = self.trainer.estimated_stepping_batches
+        warmup_steps = int(0.1 * total_steps)
+
+        # Create scheduler
+        scheduler = get_cosine_schedule_with_warmup(
+            optimizer,
+            num_warmup_steps=warmup_steps,
+            num_training_steps=total_steps,
+        )
+
+        return [optimizer], [{"scheduler": scheduler, "interval": "step"}]
+    
+    def load_dataset(self):
+        """Load the dataset based on the dataset type."""
+        if self.hparams.dataset_type == "kegg":
+            dataset = load_dataset(self.hparams.kegg_data_dir_huggingface)
+            
+            if self.hparams.truncate_dna_per_side:
+                dataset = dataset.map(
+                    truncate_dna, fn_kwargs={"truncate_dna_per_side": self.hparams.truncate_dna_per_side}
+                )
+
+            labels = []
+            for split, data in dataset.items():
+                labels.extend(data["answer"])
+            labels = list(set(labels))
+        
+        elif self.hparams.dataset_type == "variant_effect_coding":
+            dataset = load_dataset("wanglab/bioR_tasks", "variant_effect_coding")
+            dataset = dataset.map(clean_variant_effect_example)
+
+            if self.hparams.truncate_dna_per_side:
+                dataset = dataset.map(
+                    truncate_dna, fn_kwargs={"truncate_dna_per_side": self.hparams.truncate_dna_per_side}
+                )
+
+            labels = []
+            for split, data in dataset.items():
+                labels.extend(data["answer"])
+            labels = sorted(list(set(labels)))
+        
+        elif self.hparams.dataset_type == "variant_effect_non_snv":
+            dataset = load_dataset("wanglab/bioR_tasks", "task5_variant_effect_non_snv")
+            dataset = dataset.rename_column("mutated_sequence", "variant_sequence")
+            dataset = dataset.map(clean_variant_effect_example)
+
+            if self.hparams.truncate_dna_per_side:
+                dataset = dataset.map(
+                    truncate_dna, fn_kwargs={"truncate_dna_per_side": self.hparams.truncate_dna_per_side}
+                )
+
+            labels = []
+            for split, data in dataset.items():
+                labels.extend(data["answer"])
+            labels = sorted(list(set(labels)))
+
+        else:
+            raise ValueError(f"Invalid dataset type: {self.hparams.dataset_type}")
+
+        print(f"Dataset:\n{dataset}\nLabels:\n{labels}\nNumber of labels:{len(labels)}")
+        return dataset, labels
+
+    def train_dataloader(self):
+        """Create and return the training DataLoader."""
+        if self.hparams.dataset_type == "kegg":
+            train_dataset = self.dataset["train"]
+            collate_fn = lambda b: dna_collate_fn(b, dna_tokenizer=self.dna_tokenizer, label2id=self.label2id, max_length=self.hparams.max_length_dna)
+        
+        elif self.hparams.dataset_type == "variant_effect_coding":
+            train_dataset = self.dataset["train"]
+            collate_fn = lambda b: dna_collate_fn(b, dna_tokenizer=self.dna_tokenizer, label2id=self.label2id, max_length=self.hparams.max_length_dna)
+        
+        elif self.hparams.dataset_type == "variant_effect_non_snv":
+            train_dataset = self.dataset["train"]
+            collate_fn = lambda b: dna_collate_fn(b, dna_tokenizer=self.dna_tokenizer, label2id=self.label2id, max_length=self.hparams.max_length_dna)
+        
+        else:
+            raise ValueError(f"Invalid dataset type: {self.hparams.dataset_type}")
+
+        return DataLoader(
+            train_dataset,
+            batch_size=self.hparams.batch_size,
+            shuffle=True,
+            collate_fn=collate_fn,
+            num_workers=self.hparams.num_workers,
+            persistent_workers=True,
+        )
+
+    def val_dataloader(self):
+        """Create and return the training DataLoader."""
+        if self.hparams.dataset_type == "kegg":
+
+            if self.hparams.merge_val_test_set:
+                val_dataset = concatenate_datasets([self.dataset['test'], self.dataset['val']])
+            else:
+                val_dataset = self.dataset["val"]
+
+            collate_fn = lambda b: dna_collate_fn(b, dna_tokenizer=self.dna_tokenizer, label2id=self.label2id, max_length=self.hparams.max_length_dna)
+        
+        elif self.hparams.dataset_type == "variant_effect_coding":
+            val_dataset = self.dataset["test"]
+            collate_fn = lambda b: dna_collate_fn(b, dna_tokenizer=self.dna_tokenizer, label2id=self.label2id, max_length=self.hparams.max_length_dna)
+        
+        elif self.hparams.dataset_type == "variant_effect_non_snv":
+            val_dataset = self.dataset["test"]
+            collate_fn = lambda b: dna_collate_fn(b, dna_tokenizer=self.dna_tokenizer, label2id=self.label2id, max_length=self.hparams.max_length_dna)
+        
+        else:
+            raise ValueError(f"Invalid dataset type: {self.hparams.dataset_type}")
+
+        return DataLoader(
+            val_dataset,
+            batch_size=self.hparams.batch_size,
+            shuffle=False,
+            collate_fn=collate_fn,
+            num_workers=self.hparams.num_workers,
+            persistent_workers=True,
+        )
+    
+    def test_dataloader(self):
+        """Create and return the test DataLoader."""
+        return self.val_dataloader()
+
+
+def main(args):
+    """Main function to run the training process."""
+    # Set random seed and environment variables
+    pl.seed_everything(args.seed)
+    torch.cuda.empty_cache()
+    torch.set_float32_matmul_precision("medium")
+
+    # Initialize model
+    model = DNAClassifierModelTrainer(args)
+
+    # Setup directories
+    run_name = f"{args.wandb_project}-{args.dataset_type}-{args.dna_model_name.split('/')[-1]}"
+    args.checkpoint_dir = f"{args.checkpoint_dir}/{run_name}-{time.strftime('%Y%m%d-%H%M%S')}"
+    args.output_dir = f"{args.output_dir}/{run_name}-{time.strftime('%Y%m%d-%H%M%S')}"
+    os.makedirs(args.output_dir, exist_ok=True)
+    os.makedirs(args.checkpoint_dir, exist_ok=True)
+
+    # Setup callbacks
+    callbacks = [
+        ModelCheckpoint(
+            dirpath=args.checkpoint_dir,
+            filename=f"{run_name}-" + "{epoch:02d}-{val_loss_epoch:.4f}",
+            save_top_k=2,
+            monitor="val_acc_epoch",
+            mode="max",
+            save_last=True,
+        ),
+        LearningRateMonitor(logging_interval="step"),
+    ]
+
+    # Setup logger
+    is_resuming = args.ckpt_path is not None
+    logger = WandbLogger(
+        project=args.wandb_project,
+        entity=args.wandb_entity,
+        save_dir=args.log_dir,
+        name=run_name,
+        resume="allow" if is_resuming else None,  # Allow resuming existing run
+    )
+
+    # Initialize trainer
+    trainer = pl.Trainer(
+        max_epochs=args.max_epochs,
+        accelerator="gpu",
+        devices=args.num_gpus,
+        strategy=(
+            "ddp"
+            if args.strategy == "ddp"
+            else DeepSpeedStrategy(stage=2, offload_optimizer=False, allgather_bucket_size=5e8, reduce_bucket_size=5e8)
+        ),
+        precision="bf16-mixed",
+        callbacks=callbacks,
+        logger=logger,
+        deterministic=False,
+        enable_checkpointing=True,
+        enable_progress_bar=True,
+        enable_model_summary=True,
+        log_every_n_steps=5,
+        accumulate_grad_batches=args.gradient_accumulation_steps,
+        gradient_clip_val=1.0,
+        val_check_interval=1 / 3,
+    )
+
+    # Train model
+    trainer.fit(model, ckpt_path=args.ckpt_path)
+    trainer.test(model, ckpt_path=args.ckpt_path if args.ckpt_path else "best")
+
+    # Save final model
+    final_model_path = os.path.join(args.output_dir, "final_model")
+    torch.save(model.dna_model.state_dict(), final_model_path)
+    print(f"Final model saved to {final_model_path}")
+
+
+if __name__ == "__main__":
+    os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+    parser = argparse.ArgumentParser(description="Train DNA Classifier")
+
+    # Model parameters
+    parser.add_argument(
+        "--dna_model_name",
+        type=str,
+        default="InstaDeepAI/nucleotide-transformer-v2-500m-multi-species",
+    )
+    parser.add_argument("--cache_dir", type=str, default="/model-weights")
+    parser.add_argument("--max_length_dna", type=int, default=1024)
+    parser.add_argument("--dna_is_evo2", type=bool, default=False)
+    parser.add_argument("--dna_embedding_layer", type=str, default=None)
+
+    # Training parameters
+    parser.add_argument("--strategy", type=str, default="ddp")
+    parser.add_argument("--batch_size", type=int, default=8)
+    parser.add_argument("--learning_rate", type=float, default=5e-5)
+    parser.add_argument("--weight_decay", type=float, default=0.01)
+    parser.add_argument("--max_epochs", type=int, default=5)
+    parser.add_argument("--max_steps", type=int, default=-1)
+    parser.add_argument("--gradient_accumulation_steps", type=int, default=8)
+    parser.add_argument("--num_workers", type=int, default=4)
+    parser.add_argument("--num_gpus", type=int, default=1)
+    parser.add_argument("--train_just_classifier", type=bool, default=True)
+    parser.add_argument("--dataset_type", type=str, choices=["kegg", "variant_effect_coding", "variant_effect_non_snv"], default="kegg")
+    parser.add_argument("--kegg_data_dir_huggingface", type=str, default="wanglab/kegg")
+    parser.add_argument("--truncate_dna_per_side", type=int, default=0)
+
+    # Output parameters
+    parser.add_argument("--output_dir", type=str, default="dna_classifier_output")
+    parser.add_argument(
+        "--checkpoint_dir", type=str, default="checkpoints"
+    )
+    parser.add_argument("--ckpt_path", type=str, default=None)
+    parser.add_argument("--log_dir", type=str, default="logs")
+    parser.add_argument("--wandb_project", type=str, default="dna-only-nt-500m")
+    parser.add_argument("--wandb_entity", type=str, default="adibvafa")
+    parser.add_argument("--merge_val_test_set", type=bool, default=True)
+
+    # Other parameters
+    parser.add_argument("--seed", type=int, default=23)
+
+    args = parser.parse_args()
+    main(args)

BioReason/train_dna_qwen.py

@@ -0,0 +1,1064 @@
+import csv
+import gc
+import io
+import multiprocessing
+import os
+import time
+import traceback
+from argparse import ArgumentParser
+from functools import partial
+from typing import *
+
+import pandas as pd
+import torch
+import wandb
+from datasets import DatasetDict, concatenate_datasets, load_dataset
+from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
+from torch.optim import AdamW
+from torch.utils.data import DataLoader
+from transformers import AutoTokenizer, get_cosine_schedule_with_warmup
+from transformers.tokenization_utils_base import BatchEncoding
+
+import pytorch_lightning as pl
+from pytorch_lightning.callbacks import LearningRateMonitor, ModelCheckpoint
+from pytorch_lightning.loggers import WandbLogger
+from pytorch_lightning.strategies import DeepSpeedStrategy
+
+from bioreason.dataset.kegg import get_format_kegg_function, qwen_dna_collate_fn
+from bioreason.dataset.utils import truncate_dna
+from bioreason.dataset.variant_effect import (
+    clean_variant_effect_example,
+    clean_variant_effect_non_snv_example,
+    get_format_variant_effect_function,
+)
+from bioreason.models.dl.processing_dl import DLProcessor
+from bioreason.models.dna_llm import DNALLMModel
+from bioreason.models.evo2_tokenizer import register_evo2_tokenizer
+
+register_evo2_tokenizer()
+
+# Set start method to 'spawn' for CUDA compatibility with multiprocessing
+torch.multiprocessing.set_sharing_strategy("file_system")
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+
+class DNALLMFineTuner(pl.LightningModule):
+    """
+    PyTorch Lightning module for fine-tuning DNA-LLM models.
+    """
+
+    def __init__(self, hparams):
+        """
+        Initialize the DNALLMFineTuner.
+
+        Args:
+            hparams: Hyperparameters for the model and training
+        """
+        super().__init__()
+        self.save_hyperparameters(hparams)
+
+        self.text_model_name = self.hparams.text_model_name
+        self.dna_model_name = self.hparams.dna_model_name
+        self.cache_dir = self.hparams.cache_dir
+        self.learning_rate = self.hparams.learning_rate
+        self.weight_decay = self.hparams.weight_decay
+        self.text_model_finetune = self.hparams.text_model_finetune
+        self.dna_model_finetune = self.hparams.dna_model_finetune
+        self.lora_rank = self.hparams.lora_rank
+        self.lora_alpha = self.hparams.lora_alpha
+        self.lora_dropout = self.hparams.lora_dropout
+        self.max_length_dna = self.hparams.max_length_dna
+        self.max_length_text = self.hparams.max_length_text
+        self.dna_is_evo2 = self.hparams.dna_is_evo2
+        self.dna_embedding_layer = self.hparams.dna_embedding_layer
+        self.return_answer_in_batch = self.hparams.return_answer_in_batch
+        self.merge_val_test_set = self.hparams.merge_val_test_set
+
+        # Store dataset configuration
+        self.dataset_type = self.hparams.dataset_type
+
+        # Load model
+        self.model = DNALLMModel(
+            text_model_name=self.text_model_name,
+            dna_model_name=self.dna_model_name,
+            cache_dir=self.cache_dir,
+            max_length_dna=self.max_length_dna,
+            max_length_text=self.max_length_text,
+            text_model_finetune=self.text_model_finetune,
+            dna_model_finetune=self.dna_model_finetune,
+            dna_is_evo2=self.dna_is_evo2,
+            dna_embedding_layer=self.dna_embedding_layer,
+        )
+
+        self.text_model = self.model.text_model
+        self.dna_model = self.model.dna_model
+        self.dna_projection = self.model.dna_projection
+
+        # Load tokenizer for target text
+        self.tokenizer = self.model.text_tokenizer
+
+        # Prepare model for training
+        self.lora_config = self._prep_for_training()
+
+    def _get_target_modules(self):
+        # Apply LoRA to all linear layers in the text model
+        target_modules = []
+
+        # Get all unique linear layer names
+        seen_names = set()
+        for name, module in self.text_model.named_modules():
+            if isinstance(module, torch.nn.Linear):
+                names = name.split(".")
+                target_name = names[-1]  # Use the last part of the name
+
+                # Skip output head but include all other linear layers
+                if target_name != "lm_head" and target_name not in seen_names:
+                    target_modules.append(target_name)
+                    seen_names.add(target_name)
+
+        # Add attention-specific layers
+        attention_patterns = [
+            "q_proj",
+            "k_proj",
+            "v_proj",
+            "out_proj",
+            "query",
+            "key",
+            "value",
+        ]
+        for pattern in attention_patterns:
+            if pattern not in seen_names:
+                target_modules.append(pattern)
+
+        # Return all unique layer names to apply LoRA to all layers
+        return list(target_modules)
+
+    def _prep_for_training(self) -> LoraConfig:
+        """
+        Load and configure the DNALLMModel.
+        """
+
+        # Freeze DNA encoder parameters
+        if self.dna_model_finetune:
+            pass
+        else:
+            if self.dna_is_evo2:
+                for param in self.dna_model.model.parameters():
+                    param.requires_grad = False
+            else:
+                for param in self.dna_model.parameters():
+                    param.requires_grad = False
+
+        if self.text_model_finetune:
+            target_modules = self._get_target_modules()
+
+            lora_config = LoraConfig(
+                r=self.lora_rank,
+                lora_alpha=self.lora_alpha,
+                lora_dropout=self.lora_dropout,
+                target_modules=target_modules,
+                init_lora_weights="gaussian",
+                bias="none",
+                task_type="CAUSAL_LM",
+            )
+
+            # Prepare text model for training
+            self.text_model = prepare_model_for_kbit_training(self.text_model)
+            self.text_model = get_peft_model(self.text_model, lora_config)
+        else:
+            # Freeze text model parameters
+            for param in self.text_model.parameters():
+                param.requires_grad = False
+
+        # Make projection layer trainable
+        for param in self.dna_projection.parameters():
+            param.requires_grad = True
+
+        return lora_config
+
+    def _step(self, batch: Dict, batch_idx: int, prefix: str) -> torch.Tensor:
+        """
+        Performs a single step for training, validation, or testing.
+
+        Args:
+            batch: Dictionary containing the batch data
+            batch_idx: Integer indicating the batch index
+            prefix: String indicating the step type ('train', 'val', or 'test')
+
+        Returns:
+            torch.Tensor: The computed loss for this batch
+        """
+        if prefix == "test":
+            return {"loss": torch.tensor(0.0, device=self.device)}
+
+        # Get batch data from the collate function
+        input_ids = batch["input_ids"].to(self.device)
+        attention_mask = batch["attention_mask"].to(self.device)
+        labels = batch["labels"].to(self.device) if "labels" in batch else None
+        dna_tokenized = batch.get("dna_tokenized")
+        if dna_tokenized is not None:
+            dna_tokenized = dna_tokenized.to(self.device)
+        batch_idx_map = batch.get("batch_idx_map")
+
+        # Forward pass through the model
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            dna_tokenized=dna_tokenized,
+            batch_idx_map=batch_idx_map,
+            labels=labels,
+        )
+
+        # Get the loss from model outputs
+        loss = outputs.loss
+
+        # Occasionally show generations for debugging purposes - ONLY during training/validation
+        # You can reduce the frequency of generations by increasing the step size to make the model train faster
+        if (prefix == "train" and (self.global_step % 3000 == 0)) or (prefix == "val" and (batch_idx % 300 == 0)):
+            try:
+                # Select first example from batch for demonstration
+                example_idx = 0
+
+                print(
+                    f"\n=== Sample Generation (step {self.global_step} / {self.trainer.estimated_stepping_batches}) ==="
+                )
+
+                # Get the tokens that define the assistant pattern
+                assistant_start_marker = "<|im_start|>assistant\n"
+                assistant_marker_tokens = self.tokenizer.encode(assistant_start_marker, add_special_tokens=False)
+                marker_tensor = torch.tensor(assistant_marker_tokens, device=input_ids.device)
+                marker_len = len(assistant_marker_tokens)
+
+                # Find non-padding tokens in input
+                non_pad = (input_ids[example_idx] != self.tokenizer.pad_token_id).nonzero(as_tuple=True)[0]
+                if len(non_pad) > 0:
+                    start_idx = non_pad[0].item()  # First non-padding token
+                else:
+                    start_idx = 0
+
+                # For each position, check if the next marker_len tokens match the pattern
+                matches = []
+                for pos in range(start_idx, input_ids.size(1) - marker_len + 1):
+                    if torch.all(input_ids[example_idx, pos : pos + marker_len] == marker_tensor):
+                        matches.append(pos)
+                        break  # Stop at first match
+
+                assistant_pos = matches[0] if matches else None
+
+                if assistant_pos is not None:
+                    # Get input up to and including the assistant marker
+                    gen_input_ids = input_ids[
+                        example_idx : example_idx + 1, start_idx : assistant_pos + marker_len
+                    ]
+                    gen_attention_mask = attention_mask[
+                        example_idx : example_idx + 1, start_idx : assistant_pos + marker_len
+                    ]
+
+                    # Extract DNA data for this example
+                    example_dna_data = None
+                    example_batch_map = None
+
+                    if dna_tokenized is not None and batch_idx_map is not None:
+                        # Find DNA sequences for this example
+                        example_indices = [i for i, idx in enumerate(batch_idx_map) if idx == example_idx]
+
+                        if len(example_indices) > 0:
+                            # Extract just this example's DNA data
+                            example_dna_data = BatchEncoding(
+                                {
+                                    "input_ids": dna_tokenized.input_ids[example_indices].to(self.device),
+                                    "attention_mask": dna_tokenized.attention_mask[example_indices].to(self.device),
+                                }
+                            )
+
+                            # For generation we need all sequences mapped to index 0
+                            example_batch_map = [0] * len(example_indices)
+
+                    # Generate text
+                    with torch.no_grad():
+                        generated = self.model.generate(
+                            input_ids=gen_input_ids,
+                            attention_mask=gen_attention_mask,
+                            dna_tokenized=example_dna_data,
+                            batch_idx_map=example_batch_map,
+                            max_new_tokens=800,
+                            temperature=0.6,
+                            top_p=0.95,
+                            top_k=20,
+                            do_sample=True,
+                        )
+
+                    # Decode and display
+                    user_input = self.tokenizer.decode(gen_input_ids[0], skip_special_tokens=False).strip()
+                    generation = self.tokenizer.decode(generated[0], skip_special_tokens=False).strip()
+
+                    # Free memory early
+                    del generated, gen_input_ids, gen_attention_mask, example_dna_data, example_batch_map
+                    gc.collect()
+
+                    print(f"=====[Sample {prefix} {batch_idx}]=====")
+                    print(f"=====[User input]=====\n{user_input}")
+                    print(f"=====[Complete generation]=====\n{generation}")
+
+                    # Get ground truth if available
+                    ground_truth = ""
+                    if labels is not None:
+                        # Find all positions where we have valid labels (not -100)
+                        valid_label_pos = (labels[example_idx] != -100).nonzero(as_tuple=True)[0]
+
+                        if len(valid_label_pos) > 0:
+                            # Check if valid labels start after assistant marker
+                            if valid_label_pos[0] >= assistant_pos + marker_len:
+                                ground_truth = self.tokenizer.decode(
+                                    input_ids[example_idx, valid_label_pos], skip_special_tokens=False
+                                ).strip()
+                                print(f"=====[Ground truth]=====\n{ground_truth}")
+
+                    # Log to wandb
+                    timestamp = time.time()
+                    step_id = f"gen_{self.global_step}-{timestamp}"
+                    wandb_logger = self.logger.experiment
+                    wandb_logger.log(
+                        {
+                            step_id: wandb.Table(
+                                columns=["timestamp", "prefix", "batch_idx", "user_input", "generation", "ground_truth"],
+                                data=[[timestamp, prefix, batch_idx, user_input, generation, ground_truth]],
+                            )
+                        }
+                    )
+
+                    # Clean up memory
+                    del user_input, generation, ground_truth
+                    torch.cuda.empty_cache()
+                    gc.collect()
+
+                else:
+                    print("No assistant marker found in the input sequence")
+
+            except Exception as e:
+                print(f"Error during sample generation: {str(e)}")
+                traceback.print_exc()
+
+        # Get current learning rate (skip during test as scheduler might not be available)
+        if prefix != "test":
+            current_lr = self.lr_schedulers().get_last_lr()[0]
+        else:
+            current_lr = 0
+
+        # Logging metrics
+        self.log(
+            f"{prefix}_loss",
+            loss,
+            on_step=True,
+            on_epoch=False,
+            prog_bar=True,
+            logger=True,
+        )
+        self.log(
+            f"{prefix}_loss_epoch",
+            loss,
+            on_step=False,
+            on_epoch=True,
+            prog_bar=True,
+            logger=True,
+            sync_dist=True,
+        )
+        
+        # Only log learning rate during training/validation
+        if prefix != "test":
+            self.log(
+                "lr",
+                current_lr,
+                on_step=True,
+                on_epoch=True,
+                prog_bar=True,
+                logger=True,
+                sync_dist=True,
+            )
+
+        return loss
+
+    def training_step(self, batch: Dict, batch_idx: int) -> torch.Tensor:
+        """Perform a single training step."""
+        return self._step(batch, batch_idx, prefix="train")
+
+    def validation_step(self, batch: Dict, batch_idx: int) -> torch.Tensor:
+        """Perform a single validation step."""
+        return self._step(batch, batch_idx, prefix="val")
+    
+    def test_step(self, batch: Dict, batch_idx: int) -> torch.Tensor:
+        """Perform a single test step."""
+        return self._step(batch, batch_idx, prefix="test")
+
+    def configure_optimizers(self):
+        """
+        Configure optimizers and learning rate schedulers.
+
+        Returns:
+            Tuple[List, List]: A tuple containing a list of optimizers and schedulers
+        """
+        optimizer = AdamW(self.parameters(), lr=self.learning_rate, weight_decay=self.weight_decay)
+
+        total_steps = self.trainer.estimated_stepping_batches
+        warmup_steps = int(0.1 * total_steps)
+
+        scheduler = get_cosine_schedule_with_warmup(
+            optimizer,
+            num_warmup_steps=warmup_steps,
+            num_training_steps=total_steps,
+        )
+
+        return [optimizer], [{"scheduler": scheduler, "interval": "step"}]
+
+    def train_dataloader(self) -> DataLoader:
+        """Create and return the training DataLoader."""
+        # Load dataset based on type specified in hyperparameters
+
+        if  self.hparams.dataset_type == "kegg":
+            # Use Hugging Face dataset if provided
+            dataset = load_dataset(self.hparams.kegg_data_dir_huggingface)
+            dataset = dataset.map(get_format_kegg_function(self.hparams.model_type))
+
+            labels = []
+            for split, data in dataset.items():
+                labels.extend(data["answer"])
+            self.labels = sorted(list(set(labels)))
+
+            train_dataset = dataset["train"]
+
+            if self.hparams.truncate_dna_per_side:
+                train_dataset = train_dataset.map(
+                    truncate_dna, fn_kwargs={"truncate_dna_per_side": self.hparams.truncate_dna_per_side}
+                )
+            
+            processor = DLProcessor(
+                tokenizer=self.model.text_tokenizer,
+                dna_tokenizer=self.model.dna_tokenizer,
+            )
+
+            # Create partial function with all required arguments except the batch
+            collate_fn = partial(
+                qwen_dna_collate_fn,
+                processor=processor,
+                max_length_text=self.max_length_text,
+                max_length_dna=self.max_length_dna,
+                return_answer_in_batch=self.return_answer_in_batch,
+            )
+                
+
+        elif self.hparams.dataset_type == "variant_effect_coding":
+            dataset = load_dataset(self.hparams.variant_effect_coding_data_dir_huggingface)
+            cleaned_dataset = dataset.map(clean_variant_effect_example)
+            dataset = dataset.map(get_format_variant_effect_function(self.hparams.model_type))
+
+            labels = []
+            for split, data in cleaned_dataset.items():
+                labels.extend(data["answer"])
+            self.labels = sorted(list(set(labels)))
+
+            train_dataset = dataset["train"]
+
+            if self.hparams.truncate_dna_per_side:
+                train_dataset = train_dataset.map(
+                    truncate_dna, fn_kwargs={"truncate_dna_per_side": self.hparams.truncate_dna_per_side}
+                )
+            
+            processor = DLProcessor(
+                    tokenizer=self.model.text_tokenizer,
+                    dna_tokenizer=self.model.dna_tokenizer,
+                )
+            
+            # Create partial function with all required arguments except the batch
+            collate_fn = partial(
+                qwen_dna_collate_fn,
+                processor=processor,
+                max_length_text=self.max_length_text,
+                max_length_dna=self.max_length_dna,
+                return_answer_in_batch=self.return_answer_in_batch,
+            )
+        
+        elif self.hparams.dataset_type == "variant_effect_non_snv":
+            dataset = load_dataset(self.hparams.variant_effect_non_snv_data_dir_huggingface)
+            dataset = dataset.map(clean_variant_effect_non_snv_example)
+            cleaned_dataset = dataset.map(clean_variant_effect_example)
+            dataset = dataset.rename_column("mutated_sequence", "variant_sequence")
+
+            labels = []
+            for split, data in cleaned_dataset.items():
+                labels.extend(data["answer"])
+            self.labels = sorted(list(set(labels)))
+
+            train_dataset = dataset["train"]
+
+            if self.hparams.truncate_dna_per_side:
+                train_dataset = train_dataset.map(
+                    truncate_dna, fn_kwargs={"truncate_dna_per_side": self.hparams.truncate_dna_per_side}
+                )
+            train_dataset = train_dataset.map(get_format_variant_effect_function(self.hparams.model_type))
+
+            processor = DLProcessor(
+                tokenizer=self.model.text_tokenizer,
+                dna_tokenizer=self.model.dna_tokenizer,
+            )
+
+            # Create partial function with all required arguments except the batch
+            collate_fn = partial(
+                qwen_dna_collate_fn,
+                processor=processor,
+                max_length_text=self.max_length_text,
+                max_length_dna=self.max_length_dna,
+                return_answer_in_batch=self.return_answer_in_batch,
+            )
+
+        else:
+            raise ValueError(f"Unknown dataset type: {self.hparams.dataset_type}")
+
+        return DataLoader(
+            train_dataset,
+            batch_size=self.hparams.batch_size,
+            shuffle=True,
+            collate_fn=collate_fn,
+            num_workers=self.hparams.num_workers,
+            persistent_workers=False,
+            pin_memory=False,
+        )
+
+    def val_dataloader(self) -> DataLoader:
+        """Create and return the validation DataLoader."""
+
+        if self.hparams.dataset_type == "kegg":
+            # Use Hugging Face dataset
+            dataset = load_dataset(self.hparams.kegg_data_dir_huggingface)
+            dataset = dataset.map(get_format_kegg_function(self.hparams.model_type))
+
+            if self.hparams.merge_val_test_set:
+                val_dataset = concatenate_datasets([dataset['test'], dataset['val']])
+            else:
+                val_dataset = dataset["val"]
+
+            labels = []
+            for split, data in dataset.items():
+                labels.extend(data["answer"])
+            self.labels = sorted(list(set(labels)))
+
+            if self.hparams.truncate_dna_per_side:
+                val_dataset = val_dataset.map(
+                    truncate_dna, fn_kwargs={"truncate_dna_per_side": self.hparams.truncate_dna_per_side}
+                )
+
+            processor = DLProcessor(
+                tokenizer=self.model.text_tokenizer,
+                dna_tokenizer=self.model.dna_tokenizer,
+            )
+
+            # Create partial function with all required arguments except the batch
+            collate_fn = partial(
+                qwen_dna_collate_fn,
+                processor=processor,
+                max_length_text=self.max_length_text,
+                max_length_dna=self.max_length_dna,
+                return_answer_in_batch=self.return_answer_in_batch,
+            )
+
+        elif self.hparams.dataset_type == "variant_effect_coding":
+            dataset = load_dataset(self.hparams.variant_effect_coding_data_dir_huggingface)
+            cleaned_dataset = dataset.map(clean_variant_effect_example)
+            dataset = dataset.map(get_format_variant_effect_function(self.hparams.model_type))
+
+            labels = []
+            for split, data in cleaned_dataset.items():
+                labels.extend(data["answer"])
+            self.labels = sorted(list(set(labels)))
+
+            val_dataset = dataset["test"]
+
+            if self.hparams.truncate_dna_per_side:
+                val_dataset = val_dataset.map(
+                    truncate_dna, fn_kwargs={"truncate_dna_per_side": self.hparams.truncate_dna_per_side}
+                )
+            
+            processor = DLProcessor(
+                tokenizer=self.model.text_tokenizer,
+                dna_tokenizer=self.model.dna_tokenizer,
+            )
+
+            # Create partial function with all required arguments except the batch
+            collate_fn = partial(
+                qwen_dna_collate_fn,
+                processor=processor,
+                max_length_text=self.max_length_text,
+                max_length_dna=self.max_length_dna,
+                return_answer_in_batch=self.return_answer_in_batch,
+            )
+        
+        elif self.hparams.dataset_type == "variant_effect_non_snv":
+            dataset = load_dataset(self.hparams.variant_effect_non_snv_data_dir_huggingface)
+            cleaned_dataset = dataset.map(clean_variant_effect_example)
+            dataset = dataset.map(clean_variant_effect_non_snv_example)
+
+            labels = []
+            for split, data in cleaned_dataset.items():
+                labels.extend(data["answer"])
+            self.labels = sorted(list(set(labels)))
+
+            dataset = dataset.rename_column("mutated_sequence", "variant_sequence")
+            val_dataset = dataset["test"]
+
+            if self.hparams.truncate_dna_per_side:
+                val_dataset = val_dataset.map(
+                    truncate_dna, fn_kwargs={"truncate_dna_per_side": self.hparams.truncate_dna_per_side}
+                )
+            val_dataset = val_dataset.map(get_format_variant_effect_function(self.hparams.model_type))
+
+            processor = DLProcessor(
+                tokenizer=self.model.text_tokenizer,
+                dna_tokenizer=self.model.dna_tokenizer,
+            )
+
+            # Create partial function with all required arguments except the batch
+            collate_fn = partial(
+                qwen_dna_collate_fn,
+                processor=processor,
+                max_length_text=self.max_length_text,
+                max_length_dna=self.max_length_dna,
+                return_answer_in_batch=self.return_answer_in_batch,
+            )
+
+        else:
+            raise ValueError(f"Unknown dataset type: {self.hparams.dataset_type}")
+
+        return DataLoader(
+            val_dataset,
+            batch_size=self.hparams.batch_size,
+            shuffle=False,
+            collate_fn=collate_fn,
+            num_workers=self.hparams.num_workers,
+            persistent_workers=False,
+            pin_memory=False,
+        )
+    
+    def test_dataloader(self) -> DataLoader:
+        """Create and return the test DataLoader."""
+        return self.val_dataloader()
+    
+    # Only for VEP datasets, for KEGG use the resulting generations in W&B
+    def on_test_epoch_end(self):
+        """
+        Called at the end of test epoch to generate text for all test examples
+        and calculate accuracy, precision, recall, and F1 score based on whether 
+        the label appears in the generated response.
+        """
+        # Get wandb logger
+        wandb_logger = self.logger.experiment
+        wandb_logger.log({"test_progress": 0.0, "status": "starting test generation"})
+        
+        # Set model to eval mode
+        self.model.eval()
+        
+        # Get test dataloader
+        test_dataloader = self.test_dataloader()
+        total_batches = len(test_dataloader)
+        
+        # Get negative and positive labels
+        neg_label = self.labels[0]  # Negative label (first item)
+        pos_label = self.labels[1]  # Positive label (second item)
+        
+        # Log label information
+        wandb_logger.log({
+            "positive_label": pos_label,
+            "negative_label": neg_label
+        })
+        print(f"Using labels - Positive: '{pos_label}', Negative: '{neg_label}'")
+        
+        # Initialize counters and storage for generations
+        total_examples = 0
+        true_positives = 0
+        true_negatives = 0
+        false_positives = 0
+        false_negatives = 0
+        processed_batches = 0
+        generations = []
+        
+        # Process each batch in the test dataloader
+        for batch_idx, batch in enumerate(test_dataloader):
+            # Log batch start to wandb
+            wandb_logger.log({
+                "test_progress": batch_idx / total_batches,
+                "status": f"processing batch {batch_idx}/{total_batches}"
+            })
+            
+            # Get batch data
+            input_ids = batch["input_ids"].to(self.device)
+            attention_mask = batch["attention_mask"].to(self.device)
+            answer = batch["answer"]
+            dna_tokenized = batch.get("dna_tokenized")
+            if dna_tokenized is not None:
+                dna_tokenized = dna_tokenized.to(self.device)
+            batch_idx_map = batch.get("batch_idx_map")
+            
+            # Get assistant marker position
+            assistant_start_marker = "<|im_start|>assistant\n"
+            assistant_marker_tokens = self.tokenizer.encode(assistant_start_marker, add_special_tokens=False)
+            marker_tensor = torch.tensor(assistant_marker_tokens, device=input_ids.device)
+            marker_len = len(assistant_marker_tokens)
+            
+            # Log batch metadata to wandb
+            wandb_logger.log({
+                "batch_size": input_ids.shape[0],
+                "input_sequence_length": input_ids.shape[1]
+            })
+            
+            # Process examples in the batch
+            examples_in_batch = 0
+            for example_idx in range(input_ids.size(0)):
+                # Log example progress to wandb
+                if total_examples % 10 == 0:
+                    current_accuracy = (true_positives + true_negatives) / max(1, total_examples)
+                    wandb_logger.log({
+                        "examples_processed": total_examples,
+                        "current_accuracy": current_accuracy
+                    })
+                
+                # Find non-padding tokens
+                non_pad = (input_ids[example_idx] != self.tokenizer.pad_token_id).nonzero(as_tuple=True)[0]
+                start_idx = non_pad[0].item() if len(non_pad) > 0 else 0
+                
+                # Find assistant marker position
+                assistant_pos = None
+                for pos in range(start_idx, input_ids.size(1) - marker_len + 1):
+                    if torch.all(input_ids[example_idx, pos:pos + marker_len] == marker_tensor):
+                        assistant_pos = pos
+                        break
+                
+                # Log to wandb if assistant marker was found
+                wandb_logger.log({"assistant_marker_found": assistant_pos is not None})
+                
+                if assistant_pos is not None:
+                    # Prepare input for generation
+                    gen_input_ids = input_ids[example_idx:example_idx + 1, start_idx:assistant_pos + marker_len]
+                    gen_attention_mask = attention_mask[example_idx:example_idx + 1, start_idx:assistant_pos + marker_len]
+                    
+                    # Extract DNA data for this example
+                    example_dna_data = None
+                    example_batch_map = None
+                    
+                    if dna_tokenized is not None and batch_idx_map is not None:
+                        example_indices = [i for i, idx in enumerate(batch_idx_map) if idx == example_idx]
+                        
+                        if example_indices:
+                            example_dna_data = BatchEncoding({
+                                "input_ids": dna_tokenized.input_ids[example_indices].to(self.device),
+                                "attention_mask": dna_tokenized.attention_mask[example_indices].to(self.device),
+                            })
+                            example_batch_map = [0] * len(example_indices)
+                    
+                    # Log generation start to wandb
+                    wandb_logger.log({"status": f"generating for example {example_idx} in batch {batch_idx}"})
+                    
+                    # Generate text
+                    with torch.no_grad():
+                        generated = self.model.generate(
+                            input_ids=gen_input_ids,
+                            attention_mask=gen_attention_mask,
+                            dna_tokenized=example_dna_data,
+                            batch_idx_map=example_batch_map,
+                            max_new_tokens=800,
+                            temperature=0.6,
+                            top_p=0.95,
+                            top_k=20,
+                            do_sample=True,
+                        )
+                    
+                    # Decode user input and generated text
+                    user_input = self.tokenizer.decode(gen_input_ids[0], skip_special_tokens=False).strip()
+                    generation = self.tokenizer.decode(generated[0], skip_special_tokens=False).strip()
+                    
+                    # Get ground truth and clean it if needed
+                    ground_truth = answer[example_idx]
+                    if ";" in ground_truth:
+                        ground_truth = ground_truth.split(";")[0]
+                    
+                    # Determine if this is a positive or negative example
+                    is_positive_example = ground_truth.lower() == pos_label.lower()
+                    is_negative_example = ground_truth.lower() == neg_label.lower()
+                    
+                    # Check if the generated text contains the ground truth
+                    generation_contains_ground_truth = ground_truth.lower() in generation.lower()
+                    
+                    # Update metrics based on the classification
+                    total_examples += 1
+                    examples_in_batch += 1
+                    
+                    if is_positive_example and generation_contains_ground_truth:
+                        true_positives += 1
+                    elif is_positive_example and not generation_contains_ground_truth:
+                        false_negatives += 1
+                    elif is_negative_example and generation_contains_ground_truth:
+                        true_negatives += 1
+                    elif is_negative_example and not generation_contains_ground_truth:
+                        false_positives += 1
+                    
+                    # Add metadata about the prediction
+                    prediction_category = (
+                        "TP" if is_positive_example and generation_contains_ground_truth else
+                        "FN" if is_positive_example and not generation_contains_ground_truth else
+                        "TN" if is_negative_example and generation_contains_ground_truth else
+                        "FP"
+                    )
+                    
+                    # Store generation data
+                    generations.append({
+                        "batch_idx": batch_idx,
+                        "example_idx": example_idx,
+                        "user_input": user_input,
+                        "generation": generation,
+                        "ground_truth": ground_truth,
+                        "contains_ground_truth": generation_contains_ground_truth,
+                        "is_positive_example": is_positive_example,
+                        "prediction_category": prediction_category
+                    })
+                    
+                    # Clean up memory
+                    torch.cuda.empty_cache()
+                    gc.collect()
+            
+            # Log batch completion to wandb
+            processed_batches += 1
+            
+            # Calculate current metrics
+            current_accuracy = (true_positives + true_negatives) / max(total_examples, 1)
+            current_precision = true_positives / max(true_positives + false_positives, 1)
+            current_recall = true_positives / max(true_positives + false_negatives, 1)
+            current_f1 = 2 * current_precision * current_recall / max(current_precision + current_recall, 1e-8)
+            
+            wandb_logger.log({
+                "batches_processed": processed_batches,
+                "examples_processed": total_examples,
+                "examples_in_last_batch": examples_in_batch,
+                "current_accuracy": current_accuracy,
+                "current_precision": current_precision,
+                "current_recall": current_recall,
+                "current_f1": current_f1,
+                "progress_percentage": (batch_idx + 1) / total_batches * 100
+            })
+        
+        # Calculate final metrics
+        accuracy = (true_positives + true_negatives) / max(total_examples, 1)
+        precision = true_positives / max(true_positives + false_positives, 1)
+        recall = true_positives / max(true_positives + false_negatives, 1)
+        f1 = 2 * precision * recall / max(precision + recall, 1e-8)
+        
+        # Log final metrics to wandb
+        wandb_logger.log({
+            "test_accuracy": accuracy,
+            "test_precision": precision,
+            "test_recall": recall,
+            "test_f1": f1,
+            "true_positives": true_positives,
+            "false_positives": false_positives,
+            "true_negatives": true_negatives,
+            "false_negatives": false_negatives,
+            "total_examples_processed": total_examples,
+            "positive_examples": true_positives + false_negatives,
+            "negative_examples": true_negatives + false_positives,
+            "test_status": "completed"
+        })
+        
+        # Create a confusion matrix
+        confusion_matrix = {
+            "True Positives": true_positives,
+            "False Positives": false_positives,
+            "True Negatives": true_negatives,
+            "False Negatives": false_negatives
+        }
+        wandb_logger.log({"confusion_matrix": confusion_matrix})
+        
+        # Create a table with all the generations
+        if generations:
+            columns = [
+                "batch_idx",
+                "example_idx",
+                "user_input",
+                "generation", 
+                "ground_truth", 
+                "contains_ground_truth", 
+                "is_positive_example", 
+                "prediction_category"
+            ]
+            data = []
+            for g in generations:
+                # Handle any missing keys
+                row = [g.get(c, "") for c in columns]
+                data.append(row)
+            
+            wandb_logger.log({
+                f"test_generations_{time.strftime('%Y%m%d-%H%M%S')}:": wandb.Table(columns=columns, data=data)
+            })
+        
+        # Save generations to a CSV file
+        model_name = self.hparams.text_model_name.split('/')[-1]
+        if self.hparams.ckpt_path:
+            csv_path = os.path.join(self.hparams.ckpt_path, f"{time.strftime('%Y%m%d-%H%M%S')}-test_generations_{model_name}.csv")
+        else:
+            csv_path = os.path.join(self.hparams.checkpoint_dir, f"{time.strftime('%Y%m%d-%H%M%S')}-test_generations_{model_name}.csv")
+        
+        try:
+            with open(csv_path, 'w', newline='', encoding='utf-8') as f:
+                if generations:
+                    writer = csv.DictWriter(f, fieldnames=generations[0].keys())
+                    writer.writeheader()
+                    for g in generations:
+                        writer.writerow(g)
+                    
+                    wandb_logger.log({"csv_saved": True, "csv_path": csv_path})
+        except Exception as e:
+            wandb_logger.log({"csv_saved": False, "csv_path": csv_path, "error": str(e)})
+        
+        # Log a summary of the metrics
+        summary = (
+            f"Test Results Summary:\n"
+            f"Total examples: {total_examples}\n"
+            f"Accuracy: {accuracy:.4f}\n"
+            f"Precision: {precision:.4f}\n"
+            f"Recall: {recall:.4f}\n"
+            f"F1 Score: {f1:.4f}\n"
+            f"TP: {true_positives}, FP: {false_positives}, TN: {true_negatives}, FN: {false_negatives}"
+        )
+        print(summary)
+        wandb_logger.log({"test_summary": summary})
+        
+        # Force garbage collection
+        torch.cuda.empty_cache()
+        gc.collect()
+        
+        return {
+            "test_accuracy": accuracy,
+            "test_precision": precision, 
+            "test_recall": recall,
+            "test_f1": f1
+        }
+
+
+def main(args: ArgumentParser):
+    """
+    Main function to run the DNA-Text fine-tuning process.
+
+    Args:
+        args (ArgumentParser): Parsed command-line arguments
+    """
+    # Set random seed and environment variables
+    pl.seed_everything(args.seed)
+    torch.cuda.empty_cache()
+    torch.set_float32_matmul_precision("medium")
+
+    # Setup directories
+    run_name = f"{args.wandb_project}-{args.dataset_type}-{args.text_model_name.split('/')[-1]}"
+    args.checkpoint_dir = f"{args.checkpoint_dir}/{run_name}-{time.strftime('%Y%m%d-%H%M%S')}"
+
+    # Initialize model
+    model = DNALLMFineTuner(args)
+
+    # Setup callbacks
+    callbacks = [
+        ModelCheckpoint(
+            dirpath=args.checkpoint_dir,
+            filename=f"{run_name}-" + "{epoch:02d}-{val_loss_epoch:.4f}",
+            save_top_k=2,
+            monitor="val_loss_epoch",
+            mode="min",
+            save_last=True,
+        ),
+        LearningRateMonitor(logging_interval="step"),
+    ]
+
+    # Setup logger
+    is_resuming = args.ckpt_path is not None
+    logger = WandbLogger(
+        project=args.wandb_project,
+        entity=args.wandb_entity,
+        save_dir=args.log_dir,
+        name=run_name,
+        resume="allow" if is_resuming else None,  # Allow resuming existing run
+    )
+
+    # Initialize the PyTorch Lightning Trainer
+    trainer = pl.Trainer(
+        max_epochs=args.max_epochs,
+        accelerator="gpu",
+        devices=args.num_gpus,
+        strategy=(
+            "ddp"
+            if args.strategy == "ddp"
+            else DeepSpeedStrategy(stage=2, offload_optimizer=False, allgather_bucket_size=5e8, reduce_bucket_size=5e8)
+        ),
+        precision="bf16-mixed",
+        callbacks=callbacks,
+        logger=logger,
+        deterministic=False,
+        enable_checkpointing=True,
+        enable_progress_bar=True,
+        enable_model_summary=True,
+        log_every_n_steps=5,
+        accumulate_grad_batches=args.gradient_accumulation_steps,
+        gradient_clip_val=1.0,
+        val_check_interval=1 / 3,
+    )
+
+    # Start the training process
+    trainer.fit(model, ckpt_path=args.ckpt_path)
+    trainer.test(model, ckpt_path=args.ckpt_path if args.ckpt_path else "best")
+
+if __name__ == "__main__":
+    parser = ArgumentParser()
+
+    # Model configuration
+    parser.add_argument("--model_type", type=str, choices=["llm", "dna-llm"], default="dna-llm")
+    parser.add_argument("--text_model_name", type=str, default="Qwen/Qwen3-1.7B")
+    parser.add_argument("--dna_model_name", type=str, default="InstaDeepAI/nucleotide-transformer-v2-500m-multi-species")
+    parser.add_argument("--text_model_finetune", type=bool, default=True)
+    parser.add_argument("--dna_model_finetune", type=bool, default=False)
+    parser.add_argument("--dna_is_evo2", type=bool, default=False)
+    parser.add_argument("--dna_embedding_layer", type=str, default=None)
+    
+    # Training parameters
+    parser.add_argument("--seed", type=int, default=23)
+    parser.add_argument("--batch_size", type=int, default=1)
+    parser.add_argument("--max_epochs", type=int, default=5)
+    parser.add_argument("--learning_rate", type=float, default=5e-5)
+    parser.add_argument("--weight_decay", type=float, default=0.01)
+    parser.add_argument("--gradient_accumulation_steps", type=int, default=8)
+    parser.add_argument("--max_length_dna", type=int, default=1024)
+    parser.add_argument("--max_length_text", type=int, default=1024)
+    parser.add_argument("--truncate_dna_per_side", type=int, default=1024)
+    parser.add_argument("--return_answer_in_batch", type=bool, default=False)
+    
+    # LoRA parameters
+    parser.add_argument("--lora_rank", type=int, default=32)
+    parser.add_argument("--lora_alpha", type=int, default=64)
+    parser.add_argument("--lora_dropout", type=float, default=0.05)
+    
+    # Infrastructure and paths
+    parser.add_argument("--checkpoint_dir", type=str, default="checkpoints")
+    parser.add_argument("--log_dir", type=str, default="logs")
+    parser.add_argument("--cache_dir", type=str, default="/model-weights")
+    parser.add_argument("--ckpt_path", type=str, default=None)
+    parser.add_argument("--num_workers", type=int, default=4)
+    parser.add_argument("--num_gpus", type=int, default=1)
+    parser.add_argument("--strategy", type=str, default="ddp")
+    
+    # Dataset configuration
+    parser.add_argument("--dataset_type", type=str, choices=["kegg", "variant_effect_coding", "variant_effect_non_snv"], default="kegg")
+    parser.add_argument("--use_qwen_dna_collate_fn", type=bool, default=True)
+    parser.add_argument("--kegg_data_dir_local", type=str, default="data/kegg")
+    parser.add_argument("--kegg_data_dir_huggingface", type=str, default="wanglab/kegg")
+    parser.add_argument("--variant_effect_coding_data_dir_huggingface", type=str, default="wanglab/variant_effect_coding")
+    parser.add_argument("--variant_effect_non_snv_data_dir_huggingface", type=str, default="wanglab/variant_effect_non_snv")
+    parser.add_argument("--merge_val_test_set", type=bool, default=False)
+    
+    # Logging and monitoring
+    parser.add_argument("--wandb_project", type=str, default="nt-500m-qwen3-1.7b-finetune")
+    parser.add_argument("--wandb_entity", type=str)
+    
+    args = parser.parse_args()
+
+    main(args)